import time
import asyncio
import aiohttp
from PIL import Image
from pydantic import BaseModel
import logging

logger = logging.getLogger(__name__)

api_key = '8e32d44e3007447cb4be6ee52c5d3110'


class UploadInfo(BaseModel):
    fileName: str
    fileType: str


class CreateInfo(BaseModel):
    taskId: str  # 创建的任务 ID，可用于查询状态或获取结果
    taskStatus: str  # 初始状态，可能为：QUEUED、RUNNING、FAILED
    clientId: str  # 平台内部标识，用于排错，无需关注
    netWssUrl: str  # WebSocket 地址（当前不稳定，不推荐使用）
    promptTips: str  # ComfyUI 校验信息（字符串格式的 JSON），可用于识别配置异常节点


class RunHubResponse(BaseModel):
    code: int  # 状态码，0 表示成功
    msg: str  # 提示信息
    data: UploadInfo | CreateInfo | str | None = None  # 数据对象

    class Config:
        extra = 'allow'  # 允许添加额外字段


async def upload(img_path: str) -> RunHubResponse:
    with open(img_path, 'rb') as f:
        img_data = f.read()

    form = aiohttp.FormData()
    form.add_field('apiKey', api_key)
    form.add_field('file', img_data, filename='image.jpg', content_type='image/jpeg')
    form.add_field('fileType', 'image')

    url = 'https://www.runninghub.cn/task/openapi/upload'

    async with aiohttp.ClientSession() as session:
        async with session.post(url, data=form) as resp:
            response = await resp.json()

    return RunHubResponse.model_validate(response)


async def create(workflow_id: str, node_info_list: list[dict[str, str]]) -> RunHubResponse:
    url = 'https://www.runninghub.cn/task/openapi/create'
    json_data = {'apiKey': api_key, 'workflowId': workflow_id, 'nodeInfoList': node_info_list}

    async with aiohttp.ClientSession() as session:
        async with session.post(url, json=json_data) as resp:
            response = await resp.json()

    return RunHubResponse.model_validate(response)


async def status(task_id: str) -> RunHubResponse:
    # 查询状态
    url = 'https://www.runninghub.cn/task/openapi/status'
    payload = {'apiKey': api_key, 'taskId': task_id}

    async with aiohttp.ClientSession() as session:
        async with session.post(url, json=payload) as resp:
            response = await resp.json()

    # ["QUEUED","RUNNING","FAILED","SUCCESS"]
    return RunHubResponse.model_validate(response)


async def outputs(task_id: str) -> dict:
    # 获取结果
    url = 'https://www.runninghub.cn/task/openapi/outputs'
    payload = {'apiKey': api_key, 'taskId': task_id}

    async with aiohttp.ClientSession() as session:
        async with session.post(url, json=payload) as resp:
            response = await resp.json()

    return response


async def 花纹提取_api(img_path: str, save_path: str, prompt: str = '') -> bool:
    """
    异步花纹提取API

    Args:
        img_path: 输入图片路径
        save_path: 输出图片路径
        prompt: 自定义提示词，为空则使用默认提示词

    Returns:
        bool: 处理是否成功
    """
    try:
        upload_res = await upload(img_path=img_path)
        if upload_res.code != 0 or not upload_res.data:
            logger.error(f"Qwen上传失败: code={upload_res.code}, msg={upload_res.msg}")
            return False
            
        # 确保 data 是 UploadInfo 类型
        if not hasattr(upload_res.data, 'fileName'):
            logger.error(f"Qwen上传返回数据格式错误: {upload_res.data}")
            return False
            
        logger.info(f"Qwen上传成功: {upload_res.data.fileName}")

        workflow_id = '1980864078929379330'
        if len(prompt) == 0:
            prompt = '提取桌布上的花纹，自动补全空白，使得所有位置饱满并且完美衔接，去除所有的皱纹和扭曲和凸凹不平，图案自动摆正对齐并且铺平，使直线变得笔直，平行的花纹更有规律，没有残缺的花纹和折痕和断痕，铺满画布，完整的图案。简单的纯色背景'

        node_info_list = [
            {
                'nodeId': '78',
                'fieldName': 'image',
                'fieldValue': upload_res.data.fileName,
            },
            {
                'nodeId': '103',
                'fieldName': 'text',
                'fieldValue': prompt,
            },
        ]
        create_res = await create(workflow_id=workflow_id, node_info_list=node_info_list)
        
        if create_res.code != 0 or not create_res.data:
            logger.error(f"Qwen任务创建失败: code={create_res.code}, msg={create_res.msg}")
            return False
            
        # 确保 data 是 CreateInfo 类型
        if not hasattr(create_res.data, 'taskId'):
             logger.error(f"Qwen任务创建返回数据格式错误: {create_res.data}")
             return False
             
        task_id = create_res.data.taskId
        logger.info(f"Qwen任务创建成功: {task_id}")

        # 轮询检查状态
        max_retries = 120  # 最多等待10分钟（120次 * 5秒）
        retry_count = 0

        while retry_count < max_retries:
            status_res = await status(task_id=task_id)
            if status_res.code == 0:
                if status_res.data == 'QUEUED':
                    logger.info('Qwen队列排队中...')
                elif status_res.data == 'RUNNING':
                    logger.info('Qwen正在处理中...')
                elif status_res.data == 'FAILED':
                    logger.error(f'Qwen处理失败: {status_res}')
                    return False
                elif status_res.data == 'SUCCESS':
                    logger.info('Qwen处理完成，开始下载结果')
                    outputs_res = await outputs(task_id=task_id)
                    img_url = outputs_res['data'][0]['fileUrl']

                    # 下载结果图片
                    async with aiohttp.ClientSession() as session:
                        async with session.get(img_url) as resp:
                            img_data = await resp.read()

                    with open(save_path, 'wb') as f:
                        f.write(img_data)

                    logger.info(f"Qwen结果保存成功: {save_path}")
                    try:
                        from utils.api_cost_tracker import record
                        record("qwen_enhance", count=1)
                    except Exception:
                        pass
                    return True

                await asyncio.sleep(5)  # 每5秒检查一次
                retry_count += 1
            else:
                logger.error(f'Qwen处理失败: {status_res}')
                return False

        logger.error(f"Qwen处理超时，超过{max_retries * 5}秒")
        return False

    except Exception as e:
        logger.error(f"Qwen花纹提取异常: {e}")
        import traceback
        logger.error(f"异常堆栈: {traceback.format_exc()}")
        return False


async def 清晰化_api(img_path: str, save_path: str) -> bool:
    """
    高清增强：对透视矫正后的图案进行清晰化处理。
    使用与花纹提取相同的 ComfyUI 工作流，但提示词聚焦于清晰度增强。

    Args:
        img_path:  输入图片路径（透视矫正后的结果）
        save_path: 输出图片路径

    Returns:
        bool: 处理是否成功
    """
    prompt = (
        "对这张已展平的图案进行高清增强处理："
        "提升整体清晰度和锐利度，修复模糊边缘，补全细节纹理，"
        "使图案线条清晰笔直，颜色鲜艳均匀，"
        "去除噪点和压缩痕迹，输出印刷级高质量平面图，"
        "背景保持纯白色，不要改变图案内容和构图。"
    )
    return await 花纹提取_api(img_path=img_path, save_path=save_path, prompt=prompt)


# 测试代码（注释掉）
# if __name__ == "__main__":
#     asyncio.run(花纹提取_api(img_path=r'1.jpg', save_path='save1.png', prompt=''))