import os
import pandas as pd
import gradio as gr
from openai import OpenAI

# 初始化 OpenAI 客户端
client = OpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url="https://api.deepseek.com"
)

def read_csv_prefix(file, n_lines=128):
    """读取 CSV 文件的前 n 行并提取统一前缀"""
    df = pd.read_csv(file.name)
    if 'prompt' in df.columns:
        return df['prompt'].head(n_lines).tolist()
    else:
        raise ValueError("CSV 文件中没有 'prompt' 列")

def generate_prompts(user_input, csv_prompts, prefix, num_examples=3, language='English'):
    """基于用户输入和 CSV 提示生成多个新提示词"""
    # 构建系统提示
    system_prompt = (
        f"你是一位自然诗人，请将以下内容转化为{'英文' if language == 'English' else '中文'}的风景描写。"
        "请确保风格与提供的 CSV 中的描述一致，并以统一前缀 '{prefix}' 开头。"
        "请生成 {num_examples} 条不同风格的结果，每条结果不要加编号，且保持简洁自然。注意：输出语言必须与要求一致。"
    )

    # 拼接用户输入和 CSV 内容
    user_content = f"转换以下内容：\n{user_input}\n\n参考示例：\n" + "\n".join(csv_prompts[:5])

    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": system_prompt.format(prefix=prefix, num_examples=num_examples)},
            {"role": "user", "content": user_content}
        ],
        temperature=0.7,
        max_tokens=300 * num_examples
    )

    # 将模型输出按换行分割为多个提示词
    raw_output = response.choices[0].message.content.strip()
    prompts = [line.strip() for line in raw_output.split('\n') if line.strip()]

    # 限制数量
    prompts = prompts[:num_examples]

    # 处理后文本：每行以指定前缀开头（先判断是否已包含）
    processed_prompts = []
    for p in prompts:
        if prefix not in p:
            processed_prompts.append(f"{prefix} {p}")
        else:
            processed_prompts.append(p)

    # 输出格式
    output_text = "\n\n".join(prompts)
    processed_text = "\n".join(processed_prompts)

    return output_text, processed_text

# Gradio 界面定义
with gr.Blocks(title="提示词生成器", theme=gr.themes.Soft()) as app:
    gr.Markdown("## 🎨 基于 CSV 的提示词生成器")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="输入文本",
                placeholder="例如：森林里的阳光透过树叶洒落...",
                lines=5
            )
            csv_file = gr.File(label=".csv 文件", file_types=[".csv"])
            line_count = gr.Number(label="读取行数", value=128, minimum=1, maximum=10000)
            prefix_input = gr.Textbox(label="提示词前缀", value="In the style of anime landscape,", lines=1)
            num_examples = gr.Number(label="生成示例数量", value=3, minimum=1, maximum=10)
            language_choice = gr.Radio(choices=["English", "Chinese"], label="输出语言", value="English")
            btn = gr.Button("生成提示词", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(
                label="生成的提示词",
                lines=8,
                interactive=False
            )
            processed_text = gr.Textbox(
                label="处理后文本（每行以指定前缀开头）",
                lines=8,
                interactive=False
            )
    
    # 示例数据
    examples = gr.Examples(
        examples=[
            ["用英文描述，清晨的阳光洒在湖面上"],
            ["保留角色性别和外貌，用英文，描述角色的动作场面"]
        ],
        inputs=[input_text],
        label="点击试试示例"
    )

    # 事件绑定
    btn.click(
        fn=lambda text, file, lines, prefix, num, lang: generate_prompts(text, read_csv_prefix(file, lines), prefix, int(num), lang),
        inputs=[input_text, csv_file, line_count, prefix_input, num_examples, language_choice],
        outputs=[output_text, processed_text]
    )

app.launch(share=True)