from pydantic import BaseModel
from openaiapi.client import generate as client_openaiapi
from ollama.client import generate as client_ollama
import unittest
import json
from utils import clean_md_whitespace

class Document(BaseModel):
    recommend: str
    options: list[str]

def generate(content, type="openaiapi", auto_replace=True):
    categories ={
        "options": [
            "新时代码农",
            "HomeLab:中年男人的快乐源泉",
            "AI:人工智能",
            "生活:生下来活下去",
            "转载内容",
            "经验分享",
            "我的项目",
            "闲聊杂谈",
            "软件推荐",
            "好物推荐",
            "翻译内容"
        ]
    }
    if not auto_replace:
        return json.dumps(categories)

    # 构造 prompt
    prompt = f"""
    请根据我提供给你的博客内容的核心主题为其选择一个最合适的文章分类。
    附加说明：
    1. 生成的分类必须是我给定的以下几个选项之一:
        {categories['options']}
    2. 如果涉及到编程与开发相关的关键字, 比如 Java, Spring, Redis, Rabbitmq, SQL, MySQL, JVM等, 全部划分到 '新时代码农' 这个分类。
    3. 你必须以JSON格式响应，键为'recommend'，值是字符串格式的分类名称, 键为'options'，值是我给定的分类列表。
    """

    if type == "openaiapi":
        return client_openaiapi(prompt, content, response_format=Document)
    elif type == "ollama":
        prompt = prompt + f"""
        请分析'CONTENT START HERE'和'CONTENT END HERE'之间的文本

        CONTENT START HERE

        {content}

        CONTENT END HERE
        """
        return client_ollama(prompt)

class Document(BaseModel):
    recommend: str
    options: list[str]

import requests

def generate(prompt, model="qwen2"):
    # 设置 Ollama API 的 URL
    url = f"http://localhost:11434/api/generate"
    data = {
        "model": model,
        "prompt": prompt,
        "format": "json",
        "stream": False
    }

    # 发送 POST 请求到 Ollama API
    try:
        response = requests.post(url, json=data, stream=False)
        response.raise_for_status()  # 检查请求是否成功
        data = response.json()

        # 提取生成的摘要
        if "response" in data:
            return data["response"]
        else:
            print("没有从响应中获取到摘要内容！")
            return None
    except requests.exceptions.RequestException as e:
        print(f"请求失败: {e}")
        return None

from openai import OpenAI

def generate(prompt, content, usemodel="glm-4-9b-chat-1m", response_format=None):
    client = OpenAI(base_url="http://localhost:1234/v1", api_key="lm-studio")

    completion = client.beta.chat.completions.parse(
        model=usemodel,
        messages=[
            {"role": "system", "content": prompt},
            {"role": "user", "content": content}
        ],
        temperature=0.7,  # 根据需要调整这个参数
        response_format=response_format,  # 动态传入的格式
    )
    return completion.choices[0].message.content

import os
import sys
import time
import json
import threading
import queue
from utils import log, get_process_md_files, split_md, dump_md_yaml, clean_md_whitespace, get_md_category, save_processed_file, load_processed_files
from generate_category import generate as generate_category_from_ai

# 配置路径
PROCESSED_FILE = "./processed_category_files.txt"  # 已处理的文件记录

# 全局队列
task_queue = queue.Queue()
result_queue = queue.Queue()

def fetch_category_from_ai(md_file, auto_replace111):
    """
    调用 ai 接口生成分类。
    这是占位函数，具体逻辑请实现后替换。
    """
    content = clean_md_whitespace(md_file)
    # print(blog_content)
    category = generate_category_from_ai(content, auto_replace=auto_replace111)
    if category:
        return json.loads(category)
    return None

def replace_md_category(md_file, original_categorie, new_category):
    # 调用函数并获取 body 和 data
    result = split_md(md_file)
    if not result:
        return

    data = result['data']
    body = result['body']

    # 替换 categories
    data['categories'] = [new_category]
    dump_md_yaml(md_file, data, body)  # 保存更新后的 YAML 和 body

    log(f"修改分类: {md_file}")
    log(f"{original_categorie} -> {new_category}")


def ai_worker(auto_replace):
    """
    负责调用 ai 的线程。
    """
    while True:
        md_file = task_queue.get()
        if md_file is None:  # 检查是否结束
            break
        data = fetch_category_from_ai(md_file, auto_replace)
        result_queue.put((md_file, data))
        task_queue.task_done()

def interactive_worker(auto_replace):
    """
    负责交互式处理的线程。
    """
    processed_files = load_processed_files(PROCESSED_FILE)

    while True:
        md_file, data = result_queue.get()

        if data is None:  # 标识当前任务还未完成，阻塞等待
            log(f"正在等待为 {md_file} 生成分类...")
            result_queue.put((md_file, None))  # 重新将任务放回队列以等待生成完成
            time.sleep(1)  # 等待片刻再检查
            continue

        if not data:  # 如果明确返回空列表，说明分类生成失败
            log(f"未生成分类，跳过 {md_file}")
            continue

        original_category = get_md_category(md_file)

        if auto_replace:
            # 提取分类列表
            new_category = data.get("recommend", '')
            replace_md_category(md_file, original_category, new_category)
        else:
            # 提取分类列表
            options = data.get("options", [])
            # 处理分类逻辑
            print(f"\n为文件 {md_file} 生成的分类如下：")
            for i, option in enumerate(options):
                print(f"{i + 1}: {option}")

            choice = input("请选择分类编号: ").strip()
            if choice.isdigit():
                choice = int(choice)
                if choice > 0 and choice <= len(options):
                    new_category = options[choice - 1]
                    replace_md_category(md_file, original_category, new_category)
                else:
                    log(f"未替换 {md_file} 的分类。")

        save_processed_file(md_file, PROCESSED_FILE)
        processed_files.add(md_file)
        result_queue.task_done()

def main():
    # True 使用 AI 自动替换, False 使用交互式替换
    auto_replace = True

    dicts = get_process_md_files(sys.argv[1:])
    # 确保发布目录存在
    os.makedirs(dicts.get('publish_dir'), exist_ok=True)

    md_files_to_process = dicts.get('files')

    # 加载已处理文件
    processed_files = load_processed_files(PROCESSED_FILE)

    # 获取所有 Markdown 文件
    md_files_to_process = [f for f in md_files_to_process if f not in processed_files]

    # 检查是否有未处理的文件
    if not md_files_to_process:
        log("没有需要处理的 Markdown 文件。")
        return

    # 启动 Ollama 调用线程
    threading.Thread(target=ai_worker, args=(auto_replace,), daemon=True).start()

    # 将任务添加到队列
    for md_file in md_files_to_process:
        task_queue.put(md_file)

    # 启动交互线程
    threading.Thread(target=interactive_worker, args=(auto_replace,), daemon=True).start()

    # 等待任务完成
    task_queue.join()
    result_queue.join()
    log("==================category 处理完成==================")

if __name__ == "__main__":
    main()

# True 使用 AI 自动替换, False 使用交互式替换
auto_replace = True