import os
import sys
import time
import base64
import subprocess
import asyncio
import requests

from telegram import Update
from telegram.ext import Application, ContextTypes, MessageHandler, filters

# ======================
# 基础配置
# ======================

BOT_TOKEN = os.getenv("TG_BOT_TOKEN")

# Ollama
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://127.0.0.1:11434/api/chat")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3.2:latest")

# A1111（Stable Diffusion WebUI）
SD_URL = os.getenv("SD_URL", "http://127.0.0.1:7860")
SD_TXT2IMG = f"{SD_URL}/sdapi/v1/txt2img"

# 小红书发布
DEFAULT_XHS_IMAGE = os.getenv("XHS_IMAGE", "test.jpg")  # 出图失败就用这个兜底

# 出图参数（默认走轻量一点，减少超时）
SD_STEPS = int(os.getenv("SD_STEPS", "18"))
SD_CFG = float(os.getenv("SD_CFG", "6"))
SD_W = int(os.getenv("SD_W", "768"))
SD_H = int(os.getenv("SD_H", "768"))
SD_SAMPLER = os.getenv("SD_SAMPLER", "DPM++ 2M Karras")

# A1111 请求超时（第一次加载模型会很慢）
SD_TIMEOUT = int(os.getenv("SD_TIMEOUT", "600"))  # 秒


# ======================
# 判断是否是发小红书意图
# ======================
def is_xhs_intent(text: str) -> bool:
    keywords = [
        "发小红书",
        "帮我发小红书",
        "发布小红书",
        "发一条小红书",
        "帮我发个小红书",
        "发个小红书",
        "帮我发个小红书笔记",
        "发小红书笔记",
    ]
    return any(k in (text or "") for k in keywords)


# ======================
# Ollama 普通聊天
# ======================
def ollama_chat(user_text: str) -> str:
    payload = {
        "model": OLLAMA_MODEL,
        "messages": [
            {"role": "system", "content": "你是一个简洁、直接的中文助手。"},
            {"role": "user", "content": user_text},
        ],
        "stream": False,
    }
    r = requests.post(OLLAMA_URL, json=payload, timeout=120)
    r.raise_for_status()
    data = r.json()
    return (data.get("message") or {}).get("content", "").strip() or "（Ollama 无回复）"


# ======================
# 生成小红书文案 + 图片提示词（无JSON）
# 输出格式：
# 第一行：标题
# 第二行开始：正文
# 最后一行以 IMAGE: 开头：图片提示词（英文为主）
# ======================
def ollama_make_xhs_copy_and_image_prompt(user_need: str):
    prompt = f"""为小红书生成一条图文笔记，并给出配图提示词。

输出规则（非常重要）：
- 第一行只写标题（不超过20字）
- 第二行开始写正文（80-200字，分段，3-6个话题标签）
- 最后一行必须以 IMAGE: 开头，后面是图片生成提示词（英文为主，简洁，不要特殊符号）
- 不要JSON
- 不要代码块
- 不要多余解释
- 不要花括号引号等特殊符号

用户需求：
{user_need}
"""

    payload = {
        "model": OLLAMA_MODEL,
        "messages": [
            {"role": "system", "content": "你是专业小红书运营文案与配图策划助手。"},
            {"role": "user", "content": prompt},
        ],
        "stream": False,
    }

    r = requests.post(OLLAMA_URL, json=payload, timeout=120)
    r.raise_for_status()
    text = (r.json().get("message") or {}).get("content", "").strip()

    # 清理常见全角符号
    text = (
        text.replace("｛", "")
            .replace("｝", "")
            .replace("：", ":")
            .replace("“", "")
            .replace("”", "")
    )

    lines = [ln.strip() for ln in text.split("\n") if ln.strip()]
    if not lines:
        return "小红书笔记", user_need, "cozy coffee shop, warm light, product photo, minimal style"

    title = lines[0].strip()

    image_prompt = None
    body_lines = []
    for ln in lines[1:]:
        if ln.upper().startswith("IMAGE:"):
            image_prompt = ln.split(":", 1)[1].strip()
        else:
            body_lines.append(ln)

    body = "\n".join(body_lines).strip() or user_need
    if not image_prompt:
        image_prompt = "cozy coffee shop, warm light, product photo, minimal style, high quality"

    return title, body, image_prompt


# ======================
# A1111 生成图片（稳健版：超时加长 + 重试降分辨率/步数）
# 返回本地图片路径
# ======================
def sd_generate_image(image_prompt: str) -> str:
    os.makedirs("images", exist_ok=True)

    positive = f"{image_prompt}, clean composition, commercial photography, high detail"
    negative = "low quality, blurry, watermark, text, logo, deformed, extra fingers, bad anatomy"

    attempts = [
        {"w": SD_W, "h": SD_H, "steps": SD_STEPS},
        {"w": 640, "h": 640, "steps": 16},
        {"w": 512, "h": 512, "steps": 14},
    ]

    last_err = None

    for a in attempts:
        payload = {
            "prompt": positive,
            "negative_prompt": negative,
            "steps": a["steps"],
            "cfg_scale": SD_CFG,
            "width": a["w"],
            "height": a["h"],
            "sampler_name": SD_SAMPLER,
            "batch_size": 1,
            "n_iter": 1,
        }

        try:
            r = requests.post(SD_TXT2IMG, json=payload, timeout=SD_TIMEOUT)
            r.raise_for_status()
            data = r.json()

            imgs = data.get("images") or []
            if not imgs:
                raise RuntimeError("A1111 returned no images")

            b64 = imgs[0]
            if "," in b64:
                b64 = b64.split(",", 1)[1]

            raw = base64.b64decode(b64)

            ts = time.strftime("%Y%m%d_%H%M%S")
            path = os.path.abspath(os.path.join("images", f"xhs_{ts}.png"))
            with open(path, "wb") as f:
                f.write(raw)

            return path

        except Exception as e:
            last_err = e
            continue

    raise RuntimeError(f"txt2img failed after retries: {last_err}")


# ======================
# 调用小红书发布 CLI
# ======================
def run_xhs_publish(title: str, body: str, image_path: str) -> str:
    cmd = [
        sys.executable,
        "xhs_publish_cli.py",
        "--image", image_path,
        "--title", title,
        "--body", body,
    ]

    r = subprocess.run(cmd, capture_output=True, text=True)

    if r.returncode == 0:
        return "小红书已发布✅"

    return (
        "发布失败❌\n"
        f"code={r.returncode}\n"
        f"STDOUT:\n{(r.stdout or '')[-2000:]}\n\n"
        f"STDERR:\n{(r.stderr or '')[-2000:]}"
    )


def safe_delete(path: str):
    try:
        if path and os.path.exists(path):
            os.remove(path)
    except:
        pass


# ======================
# Telegram 消息处理
# ======================
async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user_text = (update.message.text or "").strip()
    if not user_text:
        return

    # 自然语言触发：发小红书
    if is_xhs_intent(user_text):
        await update.message.reply_text("收到，我先生成文案和配图提示词…")

        title, body, image_prompt = ollama_make_xhs_copy_and_image_prompt(user_text)

        await update.message.reply_text(
            f"标题：{title}\n\n配图提示词：{image_prompt}\n\n开始生成图片（A1111可能需要较久）…"
        )

        generated_path = None
        try:
            # 出图放线程
            generated_path = await asyncio.to_thread(sd_generate_image, image_prompt)
            await update.message.reply_text("图片生成完成，开始发布…（会弹出浏览器）")
            image_path = generated_path
        except Exception as e:
            await update.message.reply_text(f"出图失败，改用默认图片发布。\n原因：{e}")
            image_path = DEFAULT_XHS_IMAGE

        # 发布放线程
        result = await asyncio.to_thread(run_xhs_publish, title, body, image_path)
        await update.message.reply_text(result)

        # ✅ 发布后删除生成的图片（只删生成的，不删默认图片）
        if generated_path and generated_path != DEFAULT_XHS_IMAGE:
            safe_delete(generated_path)

        return

    # 普通聊天
    reply = ollama_chat(user_text)
    await update.message.reply_text(reply)


# ======================
# 启动
# ======================
def main():
    if not BOT_TOKEN:
        raise SystemExit('Missing TG_BOT_TOKEN. 先 export TG_BOT_TOKEN="你的token"')

    app = Application.builder().token(BOT_TOKEN).build()
    app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))

    print("TG bot running…")

    app.run_polling(
        drop_pending_updates=True,
        allowed_updates=Update.ALL_TYPES,
        close_loop=False,
    )


if __name__ == "__main__":
    main()