import gradio as gr
import requests
import json
import base64
import io
import os
import time
import re
import shutil
import platform
import subprocess
from PIL import Image

# ================= 0. 基础配置与工具 =================

# 强制不走代理连接本地 SD
os.environ['NO_PROXY'] = '127.0.0.1,localhost'

CONFIG_FILE = "config.json"
OUTPUT_DIR = "xhs_workspace"
os.makedirs(OUTPUT_DIR, exist_ok=True)

class ConfigManager:
    @staticmethod
    def load():
        if os.path.exists(CONFIG_FILE):
            try:
                with open(CONFIG_FILE, 'r', encoding='utf-8') as f:
                    return json.load(f)
            except:
                pass
        return {
            "api_key": "", 
            "base_url": "https://api.openai.com/v1", 
            "sd_url": "http://127.0.0.1:7860",
            "model": "gpt-3.5-turbo"
        }

    @staticmethod
    def save(config_data):
        with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
            json.dump(config_data, f, indent=4, ensure_ascii=False)

# ================= 1. 核心逻辑功能 =================

def get_llm_models(api_key, base_url):
    if not api_key or not base_url:
        return gr.update(choices=[]), "⚠️ 请先填写配置"
    try:
        url = f"{base_url.rstrip('/')}/models"
        headers = {"Authorization": f"Bearer {api_key}"}
        response = requests.get(url, headers=headers, timeout=10)
        if response.status_code == 200:
            data = response.json()
            models = [item['id'] for item in data.get('data', [])]
            
            # 保存配置
            cfg = ConfigManager.load()
            cfg['api_key'] = api_key
            cfg['base_url'] = base_url
            ConfigManager.save(cfg)
            
            # 修复警告：允许自定义值
            return gr.update(choices=models, value=models[0] if models else None), f"✅ 已连接，加载 {len(models)} 个模型"
        return gr.update(), f"❌ 连接失败: {response.status_code}"
    except Exception as e:
        return gr.update(), f"❌ 错误: {e}"

def generate_copy(api_key, base_url, model, topic, style):
    if not api_key: return "", "", "", "❌ 缺 API Key"
    
    # --- 核心修改：优化了 Prompt，增加字数和违禁词限制 ---
    system_prompt = """
    你是一个小红书爆款内容专家。请根据用户主题生成内容。
    
    【标题规则】(严格执行)：
    1. 长度限制：必须控制在 18 字以内（含Emoji），绝对不能超过 20 字！
    2. 格式要求：Emoji + 爆点关键词 + 核心痛点。
    3. 禁忌：禁止使用“第一”、“最”、“顶级”等绝对化广告法违禁词。
    4. 风格：二极管标题（震惊/后悔/必看/避雷/哭了），具有强烈的点击欲望。

    【正文规则】：
    1. 口语化，多用Emoji，分段清晰，不堆砌长句。
    2. 结尾必须有 5 个以上相关话题标签(#)。

    【绘图 Prompt】：
    生成对应的 Stable Diffusion 英文提示词，强调：masterpiece, best quality, 8k, soft lighting, ins style。
    
    返回 JSON 格式：
    {"title": "...", "content": "...", "sd_prompt": "..."}
    """
    
    try:
        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
        payload = {
            "model": model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"主题：{topic}\n风格：{style}"}
            ],
            "response_format": {"type": "json_object"}
        }
        resp = requests.post(f"{base_url.rstrip('/')}/chat/completions", headers=headers, json=payload, timeout=60)
        
        content = resp.json()['choices'][0]['message']['content']
        content = re.sub(r'```json\s*|```', '', content).strip()
        data = json.loads(content)
        
        # --- 双重保险：Python 强制截断 ---
        title = data.get('title', '')
        # 如果 LLM 不听话超过了20字，强制截断并保留前19个字+省略号，或者直接保留前20个
        if len(title) > 20:
            title = title[:20] 
            
        return title, data.get('content', ''), data.get('sd_prompt', ''), "✅ 文案生成完毕"
    except Exception as e:
        return "", "", "", f"❌ 生成失败: {e}"

def get_sd_models(sd_url):
    try:
        resp = requests.get(f"{sd_url}/sdapi/v1/sd-models", timeout=3)
        if resp.status_code == 200:
            models = [m['title'] for m in resp.json()]
            return gr.update(choices=models, value=models[0] if models else None), "✅ SD 已连接"
        return gr.update(choices=[]), "❌ SD 连接失败"
    except:
        return gr.update(choices=[]), "❌ SD 未启动或端口错误"

def generate_images(sd_url, prompt, neg_prompt, model, steps, cfg):
    if not model: return None, "❌ 未选择模型"
    
    # 切换模型
    try:
        requests.post(f"{sd_url}/sdapi/v1/options", json={"sd_model_checkpoint": model})
    except:
        pass # 忽略切换错误，继续尝试生成
    
    payload = {
        "prompt": prompt,
        "negative_prompt": neg_prompt,
        "steps": steps,
        "cfg_scale": cfg,
        "width": 768,
        "height": 1024,
        "batch_size": 2
    }
    
    try:
        resp = requests.post(f"{sd_url}/sdapi/v1/txt2img", json=payload, timeout=120)
        images = []
        for i in resp.json()['images']:
            img = Image.open(io.BytesIO(base64.b64decode(i)))
            images.append(img)
        return images, "✅ 图片生成完毕"
    except Exception as e:
        return None, f"❌ 绘图失败: {e}"

def one_click_export(title, content, images):
    if not title: return "❌ 无法导出：没有标题"
    
    safe_title = re.sub(r'[\\/*?:"<>|]', "", title)[:20]
    folder_name = f"{int(time.time())}_{safe_title}"
    folder_path = os.path.join(OUTPUT_DIR, folder_name)
    os.makedirs(folder_path, exist_ok=True)
    
    with open(os.path.join(folder_path, "文案.txt"), "w", encoding="utf-8") as f:
        f.write(f"{title}\n\n{content}")
        
    if images:
        for idx, img in enumerate(images):
            img.save(os.path.join(folder_path, f"图{idx+1}.png"))
            
    try:
        if platform.system() == "Windows":
            os.startfile(folder_path)
        elif platform.system() == "Darwin":
            subprocess.call(["open", folder_path])
        else:
            subprocess.call(["xdg-open", folder_path])
        return f"✅ 已导出至: {folder_path}"
    except:
        return f"✅ 已导出: {folder_path}"

# ================= 2. UI 界面构建 =================

cfg = ConfigManager.load()

with gr.Blocks(title="小红书全自动工作台", theme=gr.themes.Soft()) as app:
    gr.Markdown("## 🍒 小红书 AI 爆文生产工坊")
    
    state_images = gr.State([]) 
    
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Accordion("⚙️ 系统设置 (自动保存)", open=True):
                api_key = gr.Textbox(label="LLM API Key", value=cfg['api_key'], type="password")
                base_url = gr.Textbox(label="Base URL", value=cfg['base_url'])
                sd_url = gr.Textbox(label="SD URL", value=cfg['sd_url'])
                
                with gr.Row():
                    btn_connect = gr.Button("🔗 连接并获取模型", size="sm")
                    btn_refresh_sd = gr.Button("🔄 刷新 SD", size="sm")
                
                # 修复点 1：允许自定义值，防止报错
                llm_model = gr.Dropdown(label="选择 LLM 模型", value=cfg['model'], allow_custom_value=True, interactive=True)
                sd_model = gr.Dropdown(label="选择 SD 模型", allow_custom_value=True, interactive=True)
                status_bar = gr.Markdown("等待就绪...")

            gr.Markdown("### 💡 内容构思")
            topic = gr.Textbox(label="笔记主题", placeholder="例如：优衣库早春穿搭")
            style = gr.Dropdown(["好物种草", "干货教程", "情绪共鸣", "生活Vlog"], label="风格", value="好物种草")
            btn_step1 = gr.Button("✨ 第一步：生成文案方案", variant="primary")

        with gr.Column(scale=1):
            gr.Markdown("### 📝 文案确认")
            # 修复点 2：去掉了 show_copy_button 参数，兼容旧版 Gradio
            res_title = gr.Textbox(label="标题 (AI生成)", interactive=True)
            res_content = gr.TextArea(label="正文 (AI生成)", lines=10, interactive=True)
            res_prompt = gr.TextArea(label="绘图提示词", lines=4, interactive=True)
            
            with gr.Accordion("🎨 绘图参数", open=False):
                neg_prompt = gr.Textbox(label="反向词", value="nsfw, lowres, bad anatomy, text, error")
                steps = gr.Slider(15, 50, value=25, label="步数")
                cfg_scale = gr.Slider(1, 15, value=7, label="相关性 (CFG)")
            
            btn_step2 = gr.Button("🎨 第二步：开始绘图", variant="primary")

        with gr.Column(scale=1):
            gr.Markdown("### 🖼️ 视觉结果")
            gallery = gr.Gallery(label="生成预览", columns=1, height="auto")
            btn_export = gr.Button("📂 一键导出 (文案+图片)", variant="stop")
            export_msg = gr.Markdown("")

    # ================= 3. 事件绑定 =================
    
    btn_connect.click(fn=get_llm_models, inputs=[api_key, base_url], outputs=[llm_model, status_bar])
    btn_refresh_sd.click(fn=get_sd_models, inputs=[sd_url], outputs=[sd_model, status_bar])
    
    btn_step1.click(
        fn=generate_copy,
        inputs=[api_key, base_url, llm_model, topic, style],
        outputs=[res_title, res_content, res_prompt, status_bar]
    )
    
    def on_img_gen(sd_url, p, np, m, s, c):
        imgs, msg = generate_images(sd_url, p, np, m, s, c)
        return imgs, imgs, msg 

    btn_step2.click(
        fn=on_img_gen,
        inputs=[sd_url, res_prompt, neg_prompt, sd_model, steps, cfg_scale],
        outputs=[gallery, state_images, status_bar]
    )
    
    btn_export.click(
        fn=one_click_export,
        inputs=[res_title, res_content, state_images],
        outputs=[export_msg]
    )
    
    app.load(fn=get_sd_models, inputs=[sd_url], outputs=[sd_model, status_bar])

if __name__ == "__main__":
    app.launch(inbrowser=True)