import json
import requests
from pathlib import Path
from urllib.parse import urlparse
from PIL import Image
import io

import jingrow
from jingrow.ai.utils.utils import update_status_progress, update_progress, send_notification, get_jingrow_api_headers, get_jingrow_api_url

# 默认的系统消息和用户内容
default_system_message = """
请用中英文分别描述该图片，使用结构化描述，描述的内容用于ai绘画，因此请优化内容，不要用这是开头，使之适合用作ai绘画prompts。
输出格式为：
{
    "中文描述": "中文内容",
    "英文描述": "英文内容"
}
"""
default_user_content = "请用中英文分别生成该图片的内容描述。"

# 从配置中获取系统消息和用户内容，如果没有设置则使用默认值
system_message = jingrow.db.get_single_value("Ai Settings", "imgtxt_system_message") or default_system_message
user_content = jingrow.db.get_single_value("Ai Settings", "imgtxt_user_content") or default_user_content

# 根据配置选择模型
image_to_text_model = jingrow.db.get_single_value("Ai Settings", "image_to_text_model")

# 从 Ai Settings 读取配置
chatgpt_api_url = jingrow.db.get_single_value("Ai Settings", "chatgpt_api_url")
chatgpt_api_key = jingrow.db.get_single_value("Ai Settings", "chatgpt_api_key")
chatgpt_image_to_text_model = jingrow.db.get_single_value("Ai Settings", "chatgpt_image_to_text_model")

deepseek_api_url = jingrow.db.get_single_value("Ai Settings", "deepseek_api_url")
ai_settings = jingrow.get_single("Ai Settings")
deepseek_api_key = ai_settings.get_password("deepseek_api_key")
deepseek_image_to_text_model = jingrow.db.get_single_value("Ai Settings", "deepseek_image_to_text_model")

# 读取 Doubao 配置
doubao_api_url = jingrow.db.get_single_value("Ai Settings", "doubao_api_url")
doubao_api_key = ai_settings.get_password("doubao_api_key")
doubao_image_to_text_model = jingrow.db.get_single_value("Ai Settings", "doubao_image_to_text_model")

upload_url = jingrow.db.get_single_value("Ai Settings", "upload_url")

# 定义请求ChatGPT的函数
def send_to_chatgpt(system_message, user_content, image_url):
    # 使用更简单的 API 请求格式
    payload = {
        "model": chatgpt_image_to_text_model,
        "messages": [
            {
                "role": "system", 
                "content": system_message
            },
            {
                "role": "user", 
                "content": f"{user_content}\n\n图片链接: {image_url}"
            }
        ],
        "temperature": 0.9,
        "top_p": 0.9
    }
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {chatgpt_api_key}"
    }
    
    response = requests.post(chatgpt_api_url, headers=headers, json=payload)
    if response.status_code != 200:
        jingrow.log_error(f"Error: {response.status_code}, {response.text}", "GPT API 请求失败")
        return None
    return response.json()

# 定义请求DeepSeek的函数
def send_to_deepseek(system_message, user_content, image_url):
    # 使用更简单的 API 请求格式
    payload = {
        "model": deepseek_image_to_text_model,
        "messages": [
            {
                "role": "system", 
                "content": system_message
            },
            {
                "role": "user", 
                "content": f"{user_content}\n\n图片链接: {image_url}"
            }
        ],
        "temperature": 0.9,
        "top_p": 0.9
    }
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {deepseek_api_key}"
    }
    
    response = requests.post(deepseek_api_url, headers=headers, json=payload)
    if response.status_code != 200:
        jingrow.log_error(f"Error: {response.status_code}, {response.text}", "DeepSeek API 请求失败")
        return None
    return response.json()

# 定义请求Doubao的函数
def send_to_doubao(system_message, user_content, image_url):
    payload = {
        "model": doubao_image_to_text_model,
        "messages": [
            {
                "role": "system", 
                "content": system_message
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": user_content
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": image_url
                        }
                    }
                ]
            }
        ],
        "temperature": 0.9,
        "top_p": 0.9
    }
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {doubao_api_key}"
    }
    
    response = requests.post(doubao_api_url, headers=headers, json=payload)
    if response.status_code != 200:
        jingrow.log_error(f"Error: {response.status_code}, {response.text}", "Doubao API 请求失败")
        return None
    return response.json()

# 调用Jingrow API获取图片描述
def call_jdescribe_api(image_url, system_message=None, user_content=None):
    try:
        # 获取API URL
        api_url = f"{get_jingrow_api_url()}/jdescribe/get"
        
        # 获取认证头信息
        headers = get_jingrow_api_headers()
        if not headers:
            send_notification("Jingrow API 未设置，请在 AI设置 中设置 Jingrow Api Key 和 Jingrow Api Secret")
            return None
            
        request_data = {
            "image_url": image_url
        }
        
        # 如果提供了自定义提示词，添加到请求数据中
        if system_message:
            request_data["system_message"] = system_message
        if user_content:
            request_data["user_content"] = user_content
            
        resp = requests.post(api_url, json=request_data, headers=headers)
        if resp.status_code != 200:
            return None
        try:
            resp_json = resp.json()
        except Exception:
            return None
        # 优先找data字段，其次找message.data
        data = resp_json.get("data") or (resp_json.get("message") or {}).get("data")
        return data or None
    except Exception:
        return None

# 定义处理描述任务的函数
def process_describe_task(pg_name, pagetype):
    try:
        # 获取文档并检查是否存在
        pg = jingrow.get_pg(pagetype, pg_name)
        if not pg:
            jingrow.log_error("找不到文档", "描述任务失败")
            return False

        # 收集所有参考图片
        reference_images = []
        for i in range(1, 5):
            img = getattr(pg, f"reference_image_{i}", None)
            if img:
                reference_images.append(img)
        if not reference_images:
            jingrow.log_error("缺少参考图片", "描述任务失败")
            return False

        # 取第一个有效图片
        image_path = reference_images[0]
        if image_path.startswith(('http://', 'https://')):
            full_image_url = image_path
        else:
            full_image_url = f"{jingrow.utils.get_url()}{image_path}"

        # 更新任务状态
        update_status_progress(pagetype, pg_name, "任务进行中", 20)

        # 获取图片描述
        cn_description, en_description = "", ""
        try:
            if image_to_text_model == "Jingrow":
                update_progress(pg_name, 40, pagetype)
                data = call_jdescribe_api(full_image_url, system_message, user_content)
                if isinstance(data, dict):
                    cn_description = data.get("cn_description", "")
                    en_description = data.get("en_description", "")
                update_progress(pg_name, 60, pagetype)
            else:
                update_progress(pg_name, 40, pagetype)
                if image_to_text_model == "DeepSeek":
                    gpt_response = send_to_deepseek(system_message, user_content, full_image_url)
                elif image_to_text_model == "Doubao":
                    gpt_response = send_to_doubao(system_message, user_content, full_image_url)
                else:
                    gpt_response = send_to_chatgpt(system_message, user_content, full_image_url)
                update_progress(pg_name, 60, pagetype)
                if gpt_response and gpt_response.get("choices"):
                    try:
                        message = gpt_response["choices"][0].get("message", {}).get("content", "")
                        response_data = json.loads(message)
                        cn_description = response_data.get("中文描述", "")
                        en_description = response_data.get("英文描述", "")
                    except Exception as e:
                        jingrow.log_error(f"解析响应失败: {str(e)}", "描述失败")
                        return False
                else:
                    jingrow.log_error("AI响应无效", "描述失败")
                    return False

            if not (cn_description and en_description):
                jingrow.log_error("未获取描述内容", "描述失败")
                return False

            # 更新文档
            jingrow.db.set_value(pagetype, pg_name, "prompts_zh", cn_description)
            jingrow.db.set_value(pagetype, pg_name, "prompts", en_description)
            update_status_progress(pagetype, pg_name, "已描述", 100)
            return True

        except Exception as e:
            jingrow.log_error(f"处理图片描述异常: {str(e)}", "描述失败")
            # 更新任务状态为失败
            update_status_progress(pagetype, pg_name, "任务未完成", 0)
            return False

    except Exception as e:
        jingrow.log_error(f"处理任务异常: {str(e)}", "描述失败")
        update_status_progress(pagetype, pg_name, "任务未完成", 0)
        return False

