"""
生成图片 API - 供 pywebview 前端调用
流程：上传图片 -> workflow_run -> 轮询 query_result -> 解析返回图片 URL
"""

import base64
import json
import os
import tempfile
import time
import re
import traceback

import requests

from coze import upload_file, workflow_run, query_result
from config import workflow_id
from ali_oss import upload_data_urls as oss_upload_data_urls


def _data_url_to_temp_file(data_url: str, allow_video: bool = False) -> str:
    """将 base64 data URL 转为临时文件路径"""
    # data:image/png;base64,xxxx 或 data:video/mp4;base64,xxxx
    if allow_video:
        match = re.match(r'data:(?:image|video)/(\w+);base64,(.+)', data_url)
    else:
        match = re.match(r'data:image/(\w+);base64,(.+)', data_url)
    if not match:
        raise ValueError('无效的 data URL 格式')
    mime = match.group(1).lower()
    ext_map = {'png': 'png', 'webp': 'png', 'jpeg': 'jpg', 'jpg': 'jpg', 'mp4': 'mp4', 'webm': 'webm'}
    ext = ext_map.get(mime, 'jpg')
    data = base64.b64decode(match.group(2))
    fd, path = tempfile.mkstemp(suffix=f'.{ext}')
    try:
        os.write(fd, data)
    finally:
        os.close(fd)
    return path


def _upload_video(video_data_url: str) -> str:
    """上传视频，返回 file_id"""
    if not video_data_url or not isinstance(video_data_url, str):
        raise ValueError('无效的视频数据')
    path = _data_url_to_temp_file(video_data_url, allow_video=True)
    try:
        resp = upload_file(path)
        if resp.get('code') == 0 and resp.get('data', {}).get('id'):
            return resp['data']['id']
        raise RuntimeError(f'视频上传失败: {resp.get("msg", resp)}')
    finally:
        try:
            os.unlink(path)
        except Exception:
            pass


def _upload_images(image_data_urls: list) -> list:
    """上传多张图片，返回 file_id 列表"""
    file_ids = []
    temp_paths = []
    try:
        for data_url in (image_data_urls or []):
            if not data_url or not isinstance(data_url, str):
                continue
            if data_url.startswith("http"):
                fd, path = tempfile.mkstemp(suffix=f'.png')
                data_content = requests.get(data_url).content
                os.write(fd, data_content)
            else:
                path = _data_url_to_temp_file(data_url)
            temp_paths.append(path)
            resp = upload_file(path)
            if resp.get('code') == 0 and resp.get('data', {}).get('id'):
                file_ids.append(resp['data']['id'])
            else:
                raise RuntimeError(f'上传失败: {resp.get("msg", resp)}')
        return file_ids
    finally:
        for p in temp_paths:
            try:
                os.unlink(p)
            except Exception:
                pass


def _parse_output(output_str: str) -> list:
    """解析 query_result 中的 output，提取图片 URL 列表"""
    try:
        outer = json.loads(output_str)
        inner_str = outer.get('Output', '{}')
        inner = json.loads(inner_str)
        data_str = inner.get('data', '[]')
        data = json.loads(data_str)
        urls = data.get("images")
        long_image_url = data.get("merged_image_url")
        return urls if isinstance(urls, list) else [],long_image_url
    except Exception:
        return [],None


def _parse_output_text(output_str: str) -> list:
    """解析 query_result 中的 output，提取图片 URL 列表"""
    try:
        outer = json.loads(output_str)
        inner_str = outer.get('Output', '{}')
        inner = json.loads(inner_str)
        data_str = inner.get('data', '')
        if isinstance(data_str,str):
            data = json.loads(data_str).get("reverse_prompt")
            return [data]
        pattern = r'【?图像 \d+】.*?(?=【?图像 \d+】|$)'
        segments = re.findall(pattern, data_str, re.DOTALL)
        return segments
    except Exception:
        return ""


def _poll_until_done(wf_id: str, execute_id: str, interval: float =10, timeout: int = 600*2, is_text: bool = False) -> list:
    """轮询直到成功或超时，返回图片 URL 列表"""
    start = time.time()
    while time.time() - start < timeout:
        resp = query_result(wf_id, execute_id)
        if resp.get('code') != 0:
            raise RuntimeError(f'查询失败: {resp.get("msg", resp)}')
        items = resp.get('data') or []
        if not items:
            time.sleep(interval)
            continue
        item = items[0]
        status = item.get('execute_status', '')
        if status == 'Success':
            output = item.get('output', '{}')
            if is_text:
                return _parse_output_text(output)
            return _parse_output(output)
        if status and status not in ('Running', 'Pending', ''):
            raise RuntimeError(f'执行失败: {status}')
        time.sleep(interval)
    raise RuntimeError('生成超时')


def _res_to_2k(res: str) -> str:
    """统一分辨率为 2K/4K"""
    r = (res or '2k').strip().upper()
    return '4K' if r == '4K' else '2K'


def generate(params: dict) -> dict:
    """
    生成图片
    params: {
        menu: int,  # 1=图片反推 2=图片编辑 3=随机海报 4=克隆海报 5=服饰穿搭
        prompt: str,  # 用户自定义指令 (menu=1 必填)
        ref_images: list,  # base64 data URLs - 参考图
        video: str,  # base64 data URL - 视频 (menu=1 可选，仅支持1个)
        model_images: list,  # base64 data URLs - 多模特图 (menu!=1，最多5张)
        name: str, desc: str, ratio: str, resolution: str, count: int,
        language: str, style: str, batch_prompt: list, brand_name: str,
        Ingredients: str, activity: str, mode: str, texts: list,
        proc_images: list, layout_image: str,
    }
    返回: { success: bool, urls: list, prompts: list, error: str }
    """
    original_urls = []
    api_key = (params.get('api_key') or '').strip()
    try:
        menu = int(params.get('menu', 2))
        ref_data = params.get('ref_images') or []
        proc_data = params.get('proc_images') or []
        layout_data = params.get('layout_image')
        video_data = params.get('video')
        model_data = params.get('model_images') or []

        # menu 1: 反推词，简化参数
        if menu == 1:
            base_params = {'menu': 1, 'prompt': str(params.get('prompt', '')).strip() or ''}

            ref_ids = _upload_images(ref_data) if ref_data else []
            if ref_ids:
                base_params['ref_images'] = [{'file_id': fid} for fid in ref_ids]
            if video_data:
                try:
                    vid = _upload_video(video_data)
                    base_params['video'] = {'file_id': vid}
                except Exception as ve:
                    return {'success': False, 'urls': [], 'prompts': [], 'error': f'视频上传失败: {ve}'}
            base_params = {k: v for k, v in base_params.items() if v is not None and v != ''}
            base_params["api_key"] = api_key
            resp = workflow_run(workflow_id, base_params)
            if resp.get('code') != 0:
                return {'success': False, 'urls': [], 'prompts': [], 'error': resp.get('msg', str(resp))}
            execute_id = resp.get('execute_id')
            if not execute_id:
                return {'success': False, 'urls': [], 'prompts': [], 'error': '未返回 execute_id'}
            prompts = _poll_until_done(workflow_id, str(execute_id), is_text=True)
            # menu 1 返回提示词列表，统一转为字符串
            # prompts = [str(x) for x in result_list] if result_list else []
            # prompts = [prompts]
            return {'success': True, 'urls': [], 'prompts': prompts, 'error': ''}

        # menu != 1: 原有逻辑
        all_ref = list(ref_data)
        if layout_data:
            all_ref = [layout_data] + list(ref_data)
        all_originals = list(all_ref) + list(proc_data) + list(model_data)

        if all_originals:
            try:
                original_urls = oss_upload_data_urls(all_originals, prefix="originals")
            except Exception:
                traceback.print_exc()
                pass

        ref_ids = _upload_images(all_ref) if all_ref else []
        proc_ids = _upload_images(proc_data) if proc_data else []
        model_ids = _upload_images(model_data) if model_data else []

        res = _res_to_2k(params.get('resolution', '2K'))

        base_params = {
            'name': str(params.get('name', '')).strip(),
            'ratio': str(params.get('ratio', '')).strip(),
            'menu': menu,
            'resolution': res,
            'count': int(params.get('count', 1)),
            'desc': str(params.get('desc', '')).strip(),
            'language': str(params.get('language', '中文')).strip() or '中文',
            'style': str(params.get('style', '')).strip(),
            'prompt': str(params.get('prompt', '')).strip(),
            'batch_prompt': params.get('batch_prompt') or [],
            'brand_name': str(params.get('brand_name', '')).strip(),
            'Ingredients': str(params.get('Ingredients', '')).strip(),
            'activity': str(params.get('activity', '')).strip(),
            'mode': str(params.get('mode', '1')).strip() or '1',
            'texts': params.get('text') or [],
            "api_key" : api_key
        }

        if ref_ids:
            base_params['ref_images'] = [{'file_id': fid} for fid in ref_ids]
        if proc_ids:
            base_params['proc_images'] = [{'file_id': fid} for fid in proc_ids]
        if model_ids:
            base_params['model_images'] = [{'file_id': fid} for fid in model_ids]
        base_params = {k: v for k, v in base_params.items() if v is not None and v != ''}

        resp = workflow_run(workflow_id, base_params)
        if resp.get('code') != 0:
            return {'success': False, 'urls': [], 'error': resp.get('msg', str(resp))}
        execute_id = resp.get('execute_id')
        if not execute_id:
            return {'success': False, 'urls': [], 'error': '未返回 execute_id'}

        urls,long_image_url = _poll_until_done(workflow_id, str(execute_id))
        return {'success': True, 'urls': urls, 'original_urls': original_urls, 'error': '',"long_image_url":long_image_url}
    except Exception as e:
        return {'success': False, 'urls': [], 'prompts': [], 'original_urls': [], 'error': str(e),"long_image_url":""}