import os
from openai import OpenAI
import gradio as gr
from dotenv import load_dotenv
import tempfile
import atexit
from logtail import LogtailHandler
import logging
import datetime

# 加载环境变量
load_dotenv()

# 初始化OpenAI客户端
client = OpenAI(
    api_key=os.getenv("SHU"),
    base_url="https://api.shubiaobiao.cn/v1/"
)


# === 初始化 logtail logger ===
logtail_token = os.getenv("LOGTAIL_TOKEN")
logtail_handler = LogtailHandler(
    source_token=logtail_token,
    host='https://s1195313.eu-nbg-2.betterstackdata.com',
)

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.handlers.clear()  # 清空其他 handlers（可选）
logger.addHandler(logtail_handler)
# =======================================

def increment_user_counter():
    """使用Logtail记录点击事件（替换原文件写入方法）"""
    try:
        logger.info("用户点击下载按钮", extra={
            'event_type': 'download_click',
            'timestamp': datetime.datetime.now().isoformat(),
            'service_version': '1.0'
        })
    except Exception as e:
        print(f"日志记录失败: {str(e)}")


def call_openai_api(prompt, temperature=0.7, model="deepseek-v3"):
    """调用OpenAI API生成内容"""
    try:
        chat_completion = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": "你是一个专业的专利代理人，正在完成一份专利文档的撰写。",
                },
                {
                    "role": "user",
                    "content": prompt,
                },
            ],
            model=model,
            temperature=temperature,
            max_tokens=8000,
            n=1
        )

        return chat_completion.choices[0].message.content.strip()
    except Exception as e:
        print(f"API调用出错：{str(e)}")
        return f"生成失败：{str(e)}"


# 定义每个部分的Agent Prompt模板
BACKGROUND_PROMPT = """
你是一名资深的专利代理人，现在需要根据以下输入信息来撰写一份专利申请技术交底书的第一部分内容。
请根据用户提供的背景材料和参考文献信息，详说明本发明所处的技术领域背景和已有技术方案。

要求：
1. 内容必须清晰、详实，包括宏观技术背景和具体现有技术方案。
2. 若有参考文献（专利、期刊、书籍）请在文中标明出处。
3. 内容可读性强，让读者无需查阅外部文献即可理解。

用户提供的信息：
{user_input}

请按照如下格式输出：
一、详细介绍技术背景，并描述已有的与本发明最相近似的实现方案。
1.背景技术
2.现有技术方案详细介绍
"""

SHORTCOMING_PROMPT = """
请根据以下背景与已有技术方案描述，以及用户给出的需求信息，撰写第二部分内容：
二、现有技术的缺点是什么？
要求：
- 针对上部分已有技术方案，指出目前存在的不足、问题点或局限性。
- 可根据用户输入信息进行适当扩展。

现有信息（第一部分输出）：
{previous_content}

用户提供的需求：
{user_input}

请按照如下格式输出：
二、现有技术的缺点是什么？
……
"""

PROBLEM_PROMPT = """
请根据以下信息编写第三部分内容：
三、本发明解决的技术问题是：
要求：
- 对第二部分所列出的缺点和不足，描述本发明试图解决的具体技术问题。
- 内容应与本发明拟解决的问题相呼应。

已有输出：
第一部分：
{background_content}

第二部分：
{shortcoming_content}

用户提供的需求：
{user_input}

三、本发明解决的技术问题是：
……
"""

SOLUTION_PROMPT = """
请根据以下信息撰写第四部分内容：
四、本发明技术方案的详细阐述：
要求：
- 详细描述本发明的技术方案，包括各功能模块/步骤的技术实现方式、结构图（可文字描述）和原理说明。
- 确保是完整且清晰的技术方案描述。

已有输出：
第一部分：
{background_content}
第二部分：
{shortcoming_content}
第三部分：
{problem_content}

用户提供的需求：
{user_input}

请按照如下格式输出：
四、本发明技术方案的详细阐述：
……
"""

KEYPOINT_PROMPT = """
请根据以下信息撰写第五部分内容：
五、本发明的关键点和欲保护点是：
要求：
- 从第四部分的技术方案中提炼出关键创新点，以简洁的列点形式呈现。
- 这些点应是有利于申请权利要求保护的技术核心。

已有输出：
第四部分：
{solution_content}

用户提供的需求：
{user_input}

请按照如下格式输出：
五、本发明的关键点和欲保护点是：
……
"""

ADVANTAGE_PROMPT = """
请根据以下信息撰写第六部分内容：
六、与第二条所述最好的现有技术相比，本发明的优：
要求：
- 简要介绍本发明相对现有技术的有益效果和优势。
- 结合技术方案来描述，做到有理有据。

已有输出：
第二部分（现有技术的缺点）：
{shortcoming_content}
第三部分（本发明解决的技术问题）：
{problem_content}
第四部分（本发明技术方案的详细阐述）：
{solution_content}

用户提供的需求：
{user_input}

请按照如下格式输出：
六、与第二条所述的最好的现有技术相比，本发明的优点：
……
"""

ALTERNATIVE_PROMPT = """
请根据以下信息撰写第七部分内容：
七、针对第四部分中的技术方案，是否还有别的替代方案？
要求：
- 阐述可能的替代技术方案或实现方式，以拓宽专利保护范围。
- 可以是部分元件/步骤的替换，也可以是整体方案的替代。

已有输出：
第四部分：
{solution_content}

用户提供的需求：
{user_input}

请按照如下格式输出：
七、针对第四部分中的技术方案，是否还有别的替代方案？
……
"""

def generate_patent_document(
    bg_input, 
    shortcoming_input, 
    problem_input, 
    solution_input, 
    keypoint_input,
    advantage_input, 
    alternative_input,
    model_choice,
):
    try:
        final_document = ""
        
        # 生成背景技术部分
        # progress(0, desc="正在生成背景技术部分...")
        yield "正在生成背景技术部分...", None
        bg_prompt = BACKGROUND_PROMPT.format(user_input=bg_input)
        background_content = call_openai_api(bg_prompt, model=model_choice)
        final_document += f"{background_content}\n\n"
        yield "正在生成背景技术部分...", final_document

        # 生成现有技术缺点部分
        # progress(0.15, desc="正在生成现有技术缺点部分...")
        yield "正在生成现有技术缺点部分...", final_document
        short_prompt = SHORTCOMING_PROMPT.format(
            previous_content=background_content,
            user_input=shortcoming_input
        )
        shortcoming_content = call_openai_api(short_prompt, model=model_choice)
        final_document += f"{shortcoming_content}\n\n"
        yield "正在生成现有技术缺点部分...", final_document

        # 生成技术问题部分
        # progress(0.3, desc="正在生成技术问题部分...")
        yield "正在生成技术问题部分...", final_document
        problem_prompt_full = PROBLEM_PROMPT.format(
            background_content=background_content,
            shortcoming_content=shortcoming_content,
            user_input=problem_input
        )
        problem_content = call_openai_api(problem_prompt_full, model=model_choice)
        final_document += f"{problem_content}\n\n"
        yield "正在生成技术问题部分...", final_document

        # 生成技术方案部分
        # progress(0.45, desc="正在生成技术方案部分...")
        yield "正在生成技术方案部分...", final_document
        solution_prompt_full = SOLUTION_PROMPT.format(
            background_content=background_content,
            shortcoming_content=shortcoming_content,
            problem_content=problem_content,
            user_input=solution_input
        )
        solution_content = call_openai_api(solution_prompt_full, model=model_choice)
        final_document += f"{solution_content}\n\n"
        yield "正在生成技术方案部分...", final_document

        # 生成关键点部分
        # progress(0.6, desc="正在生成关键点部分...")
        yield "正在生成关键点部分...", final_document
        keypoint_prompt_full = KEYPOINT_PROMPT.format(
            solution_content=solution_content,
            user_input=keypoint_input
        )
        keypoint_content = call_openai_api(keypoint_prompt_full, model=model_choice)
        final_document += f"{keypoint_content}\n\n"
        yield "正在生成关键点部分...", final_document

        # 生成优点部分
        # progress(0.75, desc="正在生成优点部分...")
        yield "正在生成优点部分...", final_document
        advantage_prompt_full = ADVANTAGE_PROMPT.format(
            shortcoming_content=shortcoming_content,
            problem_content=problem_content,
            solution_content=solution_content,
            user_input=advantage_input
        )
        advantage_content = call_openai_api(advantage_prompt_full, model=model_choice)
        final_document += f"{advantage_content}\n\n"
        yield "正在生成优点部分...", final_document

        # 生成替代方案部分
        # progress(0.9, desc="正在生成替代方案部分...")
        yield "正在生成替代方案部分...", final_document
        alternative_prompt_full = ALTERNATIVE_PROMPT.format(
            solution_content=solution_content,
            user_input=alternative_input
        )
        alternative_content = call_openai_api(alternative_prompt_full, model=model_choice)
        final_document += f"{alternative_content}"
        
        # progress(1.0, desc="生成完成！")
        yield "生成完成！", final_document

    except Exception as e:
        error_message = f"生成过程中发生错误：{str(e)}"
        yield error_message, error_message

def generate_filename(text):
    """根据专利交底书内容生成合适的文件名"""
    try:
        prompt = f"""
        请根据以下专利交底书内容，生成一个简短的文件名（不超过50个字符）。
        文件名应该能反映出发明的核心内容和技术领域。
        只需要返回文件名，不需要其他解释。

        专利交底书内容：
        {text[:1000]}
        """
        filename = call_openai_api(prompt).strip()
        # 确保文件名合法
        filename = "".join(c for c in filename if c not in r'\/:*?"<>|')
        return filename + ".txt"
    except Exception as e:
        print(f"生成文件名失败：{str(e)}")
        return "专利交底书.txt"

def download_text(text):
    """将文本保存为临时文件并返回文件路径"""
    filename = generate_filename(text)
    temp_dir = tempfile.gettempdir()
    file_path = os.path.join(temp_dir, filename)
    
    with open(file_path, mode="w", encoding="utf-8") as f:
        f.write(text)
    return file_path

def clear_all():
    """清空所有输入和输出"""
    try:
        return (
            "", "", "", "", "", "", "",  # 7 个 Textbox
            "",                             # 1 个 Markdown (final_output)
            None                            # 1 个 File (download_file)
        )
    except Exception as e:
        print(f"清空过程中出错：{str(e)}")
        return (
            "", "", "", "", "", "", "",
            "",
            None
        )
def cleanup_temp_files():
    """清理临时文件"""
    temp_dir = tempfile.gettempdir()
    for file in os.listdir(temp_dir):
        if file.endswith(".txt"):
            try:
                file_path = os.path.join(temp_dir, file)
                if os.path.exists(file_path):
                    os.remove(file_path)
            except Exception as e:
                print(f"删除临时文件失败：{str(e)}")
                continue

atexit.register(cleanup_temp_files)

custom_css = """
.gradio-container {
    background-color: #f0f2f6;
    padding: 20px;
}
.gradio-container h2 {
    color: #333333;
}
.gr-button-primary {
    background-color: #4B6CB7;
    color: white;
}
.gr-button-secondary {
    background-color: #A5A5A5;
    color: white;
}
.gr-textbox, .gr-markdown {
    border-radius: 5px;
    padding: 10px;
    border: 1px solid #ccc;
}
"""

with gr.Blocks(theme=gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="indigo",
    neutral_hue="slate"
)) as demo:
    gr.Markdown("## 📝 专利交底书生成系统\n### **基于语言模型技术，为您自动生成专利文档。**")
    gr.Markdown("### ❗❗❗ 大模型引用的任何专利号或者论文资料等都需要二次审查❗❗❗")
    gr.Markdown('- 请按顺序填写下列所有部分的信息，然后点击\"🚀 开始生成\"按钮。')
    gr.Markdown("- 请提供分点的自然语言信息，得到更高质量的输出结果")

    # 添加模型选择下拉菜单
    model_dropdown = gr.Dropdown(
        choices=["deepseek-v3", "gpt-4o"],
        value="deepseek-v3",
        label="选择模型",
        info="请选择要使用的AI模型"
    )

    with gr.Accordion("📥 输入信息", open=True):
        bg = gr.Textbox(
            label="一、背景技术与现有技术方案输入",
            placeholder="请输入背景技术和现有技术方案...",
            lines=4
        )
        # 添加背景技术示例
        gr.Examples(
            examples=[["在移动设备领域，智能手机的电池续航一直是一个重要问题。目前主流的充电技术包括有线快充和无线充电。现有技术中，无线充电通常采用电磁感应原理，需要将设备放置在特定的充电板上，且充电效率相对较低。"]],
            inputs=[bg],
            label="背景技术示例"
        )

        sh = gr.Textbox(
            label="二、现有技术缺点输入",
            placeholder="请输入现有技术的缺点...",
            lines=4
        )
        # 添加现有技术缺点示例
        gr.Examples(
            examples=[["现有的无线充电技术存在充电距离短、效率低、发热严重等问题。同时，用户必须将设备精确放置在充电区域内，使用不够便捷。"]],
            inputs=[sh],
            label="现有技术缺点示例"
        )

        pr = gr.Textbox(
            label="三、本发明解决的技术问题输入",
            placeholder="请输入本发明解决的技术问题...",
            lines=4
        )
        # 添加技术问题示例
        gr.Examples(
            examples=[["本发明旨在提供一种新型远距离无线充电技术，解决现有无线充电距离短、效率低的问题，实现更便捷的充电体验。"]],
            inputs=[pr],
            label="技术问题示例"
        )

        so = gr.Textbox(
            label="四、本发明技术方案输入",
            placeholder="请输入本发明的技术方案...",
            lines=4
        )
        # 添加技术方案示例
        gr.Examples(
            examples=[["本发明提出一种基于定向电磁波的远距离无线充电系统，包含智能定向天线阵列和高效能量转换模块。系统通过毫米波技术实现精确对准和能量传输，可在3米范围内实现稳定充电。"]],
            inputs=[so],
            label="技术方案示例"
        )

        kp = gr.Textbox(
            label="五、关键点输入",
            placeholder="请输入本发明的关键点...",
            lines=3
        )
        # 添加关键点示例
        gr.Examples(
            examples=[["智能定向天线阵列设计、高效能量转换模块、远距离无线充电控制算法"]],
            inputs=[kp],
            label="关键点示例"
        )

        adv = gr.Textbox(
            label="六、本发明优点输入",
            placeholder="请输入本发明的优点...",
            lines=3
        )
        # 添加优点示例
        gr.Examples(
            examples=[["充电距离显著提升至3米，充电效率提高30%，支持多设备同时充电，智能识别待充电设备位置"]],
            inputs=[adv],
            label="优点示例"
        )

        alt = gr.Textbox(
            label="七、替代方案输入",
            placeholder="请输入替代方案...",
            lines=3
        )
        # 添加替代方案示例
        gr.Examples(
            examples=[["天线阵列可替换为相控阵天线系统，能量转换模块可采用不同的整流电路方案，控制算法可使用其他机器学习方法"]],
            inputs=[alt],
            label="替代方案示例"
        )

    with gr.Row():
        generate_button = gr.Button("🚀 开始生成", variant="primary")
        clear_button = gr.Button("🧹 清空所有", variant="secondary")

    progress_bar = gr.Progress()
    status_box = gr.Textbox(label="生成状态", interactive=False)  # 添加标签和设置为只读

    final_output = gr.Markdown(label="生成的专利交底书文本")
    with gr.Row():
        download_button = gr.Button("⬇️ 创建下载文件", variant="secondary")
        download_file = gr.File(label="点击下载文件", visible=False)

    download_button.click(
        fn=increment_user_counter,            # 第一步：记录日志
        inputs=[],                        # 不需要输入
        outputs=[]             # 不返回修改提示
    ).then(
        fn=download_text,                 # 第二步：执行下载逻辑
        inputs=[final_output],
        outputs=[download_file]
    )

    generate_button.click(
        fn=generate_patent_document,  # 第二步：再去真正生成文档
        inputs=[bg, sh, pr, so, kp, adv, alt, model_dropdown],
        outputs=[status_box, final_output]
    ).then(
        fn=lambda x: gr.update(visible=True),
        inputs=[final_output],
        outputs=[download_file]
    )


    clear_button.click(
        fn=clear_all,
        inputs=[],
        outputs=[bg, sh, pr, so, kp, adv, alt, final_output, download_file]
    )

demo.launch()