from __future__ import annotations

import base64
import json
import math
import time
import uuid
from copy import deepcopy
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
from urllib.parse import urlencode, urljoin

from .config import AnimaToolConfig
from .history import HistoryManager


def _round_up(x: int, base: int) -> int:
    if base <= 1:
        return x
    return int(math.ceil(x / base) * base)


def _parse_aspect_ratio(ratio: str) -> float:
    """
    解析 "16:9" -> 16/9。
    """
    s = (ratio or "").strip()
    if ":" not in s:
        raise ValueError(f"aspect_ratio 必须形如 '16:9'，收到：{ratio!r}")
    a_str, b_str = s.split(":", 1)
    a = float(a_str.strip())
    b = float(b_str.strip())
    if a <= 0 or b <= 0:
        raise ValueError(f"aspect_ratio 两边必须 > 0，收到：{ratio!r}")
    return a / b


def estimate_size_from_ratio(
    *,
    aspect_ratio: str,
    target_megapixels: float = 1.0,
    round_to: int = 16,
) -> Tuple[int, int]:
    """
    只给定长宽比时，估算 width/height，使像素数接近 target_megapixels。
    宽高会向上取整到 round_to 的倍数。
    
    注意：Anima 基于 Cosmos 架构，VAE 缩放 8 倍后还需被 spatial_patch_size=2 整除，
    所以 round_to 必须至少是 16（8×2），否则会报错 "should be divisible by spatial_patch_size"。
    """
    r = _parse_aspect_ratio(aspect_ratio)
    target_px = max(1.0, float(target_megapixels)) * 1_000_000.0
    w = int(math.sqrt(target_px * r))
    h = int(math.sqrt(target_px / r))
    w = _round_up(max(64, w), round_to)
    h = _round_up(max(64, h), round_to)
    return w, h


def align_dimension(value: int, round_to: int = 16) -> int:
    """
    将尺寸对齐到 round_to 的倍数（向上取整）。
    用于处理用户直接传入的 width/height。
    """
    return _round_up(max(64, int(value)), round_to)


def _join_csv(*parts: str) -> str:
    cleaned: List[str] = []
    for p in parts:
        if p is None:
            continue
        s = str(p).strip()
        if not s:
            continue
        cleaned.append(s)
    return ", ".join(cleaned)


def build_anima_positive_text(prompt_json: Dict[str, Any]) -> str:
    """
    按 Anima 推荐顺序拼接正面提示词（全部逗号连接，不分行）。
    顺序：[质量/安全/年份] [人数] [角色] [作品] [画师] [风格] [外观] [标签] [环境] [自然语言]
    
    说明：nltags（自然语言补充）放在最后，因为它是"实在没法用 tag 才写"的兜底描述。
    """
    return _join_csv(
        prompt_json.get("quality_meta_year_safe", ""),
        prompt_json.get("count", ""),
        prompt_json.get("character", ""),
        prompt_json.get("series", ""),
        prompt_json.get("artist", ""),
        prompt_json.get("style", ""),
        prompt_json.get("appearance", ""),
        prompt_json.get("tags", ""),
        prompt_json.get("environment", ""),
        prompt_json.get("nltags", ""),
    )


@dataclass(frozen=True)
class GeneratedImage:
    filename: str
    subfolder: str
    folder_type: str
    view_url: str
    saved_path: Optional[str] = None
    content: Optional[bytes] = None  # 原始图片数据


class AnimaExecutor:
    _SUPPORTED_MODEL_TYPES = ("loras", "diffusion_models", "vae", "text_encoders")
    """
    将结构化 JSON 注入 ComfyUI prompt 并执行，获取输出图片。
    """

    def __init__(self, config: Optional[AnimaToolConfig] = None):
        self.config = config or AnimaToolConfig()
        self._client_id = str(uuid.uuid4())

        # 远端 ComfyUI 返回的模型名称分隔符（Windows 常为 "\\"，Linux 常为 "/"）
        self._remote_model_path_sep_cache: Dict[str, str] = {}

        template_path = Path(__file__).resolve().parent / "workflow_template.json"
        with template_path.open("r", encoding="utf-8") as f:
            self._workflow_template: Dict[str, Any] = json.load(f)

        # 生成历史管理器
        self.history = HistoryManager()

    # -------------------------
    # Model listing / metadata
    # -------------------------
    def _read_lora_metadata(self, lora_name: str) -> Optional[Dict[str, Any]]:
        """读取 LoRA 的 sidecar 元数据文件（同名 .json）。"""
        import os
        if not self.config.comfyui_models_dir:
            return None

        models_dir = Path(self.config.comfyui_models_dir)
        # 兼容处理：lora_name 可能带路径，针对 Windows 统一斜杠并去除首尾空格
        clean_name = lora_name.strip().replace("/", os.sep).replace("\\", os.sep)
        # 核心修复：如果 clean_name 以斜杠开头，Path / 拼接会变成绝对路径导致失败
        while clean_name.startswith(os.sep):
            clean_name = clean_name[1:]
        
        # 尝试几种可能的路径拼接方式
        search_paths = [
            # 1. 标准路径: models/loras/subfolder/file.safetensors.json
            models_dir / "loras" / f"{clean_name}.json",
            # 2. 移除扩展名后的路径: models/loras/subfolder/file.json
            models_dir / "loras" / f"{os.path.splitext(clean_name)[0]}.json",
            # 3. 如果 clean_name 本身已经包含 loras 目录（容错）
            models_dir / f"{clean_name}.json" if "loras" in clean_name.lower() else None
        ]

        for meta_path in search_paths:
            if meta_path and meta_path.exists():
                try:
                    return json.loads(meta_path.read_text(encoding="utf-8"))
                except Exception:
                    continue
        return None

    def list_models(self, model_type: str) -> List[Dict[str, Any]]:
        """列出 ComfyUI 模型文件。

        - model_type=loras：强制只返回存在 sidecar 元数据（.json）的 LoRA。
        - 其他类型：返回 ComfyUI API 的原始列表。
        - 返回的 name 统一使用正斜杠（/）作为分隔符，避免 Windows 反斜杠转义问题。
        """
        model_type = (model_type or "").strip()
        if model_type not in self._SUPPORTED_MODEL_TYPES:
            raise ValueError(f"不支持的 model_type={model_type!r}，仅支持：{self._SUPPORTED_MODEL_TYPES}")

        url = urljoin(self.config.comfyui_url.rstrip("/") + "/", f"models/{model_type}")
        files = self._http_get_json(url)

        if not isinstance(files, list):
            raise RuntimeError(f"ComfyUI /models/{model_type} 返回异常：{files!r}")

        results: List[Dict[str, Any]] = []
        for raw_name in files:
            if not isinstance(raw_name, str) or not raw_name.strip():
                continue
            
            # 统一使用正斜杠格式返回，避免 Windows 反斜杠在 JSON 中的转义问题
            normalized_name = raw_name.replace("\\", "/")
            item: Dict[str, Any] = {"name": normalized_name}

            if model_type == "loras":
                # 读取 sidecar 时仍使用原始路径（因为文件系统可能需要系统分隔符）
                meta = self._read_lora_metadata(raw_name)
                if not meta:
                    # 强制要求：不提供 json sidecar 的 LoRA 不允许被 list 出来
                    continue
                item["metadata"] = meta

            results.append(item)

        return results

    def _detect_remote_model_path_sep(self, model_type: str) -> str:
        """探测远端 ComfyUI 返回的模型名路径分隔符。

        ComfyUI 在 Windows 下通常使用 "\\" 返回子目录模型名，在 Linux/macOS 下通常使用 "/"。
        这里通过调用 /models/{model_type} 做一次探测并缓存。
        """
        model_type = (model_type or "").strip()
        if model_type in self._remote_model_path_sep_cache:
            return self._remote_model_path_sep_cache[model_type]

        url = urljoin(self.config.comfyui_url.rstrip("/") + "/", f"models/{model_type}")
        files = self._http_get_json(url)

        sep = "/"
        if isinstance(files, list):
            for it in files:
                if not isinstance(it, str):
                    continue
                if "\\" in it:
                    sep = "\\"
                    break
                if "/" in it:
                    sep = "/"
                    break

        self._remote_model_path_sep_cache[model_type] = sep
        return sep

    def _normalize_remote_model_name(self, name: str, model_type: str) -> str:
        """将用户输入的 name 规范化为远端 ComfyUI 可接受的模型名格式。"""
        import os

        s = (name or "").strip()
        if not s:
            return s

        # 先统一把两种分隔符都视为路径分隔
        remote_sep = self._detect_remote_model_path_sep(model_type)
        s = s.replace("/", remote_sep).replace("\\", remote_sep)

        # 去掉可能的前导分隔符（否则在某些 Path 拼接/校验场景会变成绝对路径语义）
        while s.startswith(remote_sep):
            s = s[len(remote_sep):]
        # 也顺手把当前系统分隔符的前导去掉
        while s.startswith(os.sep):
            s = s[len(os.sep):]

        return s

    # -------------------------
    # Workflow helpers
    # -------------------------
    def _inject_loras(self, wf: Dict[str, Any], loras: Any) -> None:
        """在 UNET 与 KSampler(model) 之间注入多 LoRA（仅 UNET）。

        loras: [{"name": "xxx.safetensors", "weight": 0.8}, ...]
        """
        if not loras:
            return
        if not isinstance(loras, list):
            raise ValueError("loras 必须是数组：[{name, weight}, ...]")

        # 兼容模板变化：以 KSampler 的 model 输入为起点
        if "19" not in wf or "inputs" not in wf["19"] or "model" not in wf["19"]["inputs"]:
            raise RuntimeError("workflow_template.json 缺少 KSampler(19).inputs.model，无法注入 LoRA")

        prev_model = wf["19"]["inputs"]["model"]

        # 生成不会冲突的数字 node id
        numeric_ids = [int(k) for k in wf.keys() if str(k).isdigit()]
        next_id = (max(numeric_ids) + 1) if numeric_ids else 1

        for i, lora in enumerate(loras):
            if not isinstance(lora, dict):
                continue
            name = str(lora.get("name") or "").strip()
            if not name:
                continue
            # ComfyUI 会对 lora_name 做枚举校验，必须与 /models/loras 返回的字符串完全一致
            name = self._normalize_remote_model_name(name, "loras")
            weight = float(lora.get("weight", 1.0))

            node_id = str(next_id + i)
            wf[node_id] = {
                "class_type": "LoraLoaderModelOnly",
                "inputs": {
                    "model": prev_model,
                    "lora_name": name,
                    "strength_model": weight,
                },
            }
            prev_model = [node_id, 0]

        wf["19"]["inputs"]["model"] = prev_model

    # -------------------------
    # HTTP helpers (requests 优先, 无则 urllib)
    # -------------------------
    def _http_post_json(self, url: str, payload: Dict[str, Any]) -> Dict[str, Any]:
        try:
            import requests  # type: ignore
        except Exception:
            requests = None  # type: ignore

        data = json.dumps(payload).encode("utf-8")
        if requests is not None:
            r = requests.post(url, json=payload, timeout=self.config.timeout_s)
            r.raise_for_status()
            return r.json()

        import urllib.request

        req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
        with urllib.request.urlopen(req, timeout=self.config.timeout_s) as resp:
            raw = resp.read().decode("utf-8", errors="replace")
            return json.loads(raw)

    def _http_get_json(self, url: str) -> Any:
        try:
            import requests  # type: ignore
        except Exception:
            requests = None  # type: ignore

        if requests is not None:
            r = requests.get(url, timeout=self.config.timeout_s)
            r.raise_for_status()
            return r.json()

        import urllib.request

        with urllib.request.urlopen(url, timeout=self.config.timeout_s) as resp:
            raw = resp.read().decode("utf-8", errors="replace")
            return json.loads(raw)

    def _http_get_bytes(self, url: str) -> bytes:
        try:
            import requests  # type: ignore
        except Exception:
            requests = None  # type: ignore

        if requests is not None:
            r = requests.get(url, timeout=self.config.timeout_s)
            r.raise_for_status()
            return r.content

        import urllib.request

        with urllib.request.urlopen(url, timeout=self.config.timeout_s) as resp:
            return resp.read()

    # -------------------------
    # Core workflow injection
    # -------------------------
    def _inject(self, prompt_json: Dict[str, Any]) -> Dict[str, Any]:
        wf = deepcopy(self._workflow_template)

        # 模型文件：优先使用参数指定，其次使用配置，最后使用模板默认值
        clip_name = prompt_json.get("clip_name") or self.config.clip_name
        unet_name = prompt_json.get("unet_name") or self.config.unet_name
        vae_name = prompt_json.get("vae_name") or self.config.vae_name
        
        wf["45"]["inputs"]["clip_name"] = str(clip_name)
        wf["44"]["inputs"]["unet_name"] = str(unet_name)
        wf["15"]["inputs"]["vae_name"] = str(vae_name)

        # 可选：LoRA 注入（仅 UNET）
        self._inject_loras(wf, prompt_json.get("loras"))

        # 文本
        positive = (prompt_json.get("positive") or "").strip()
        if not positive:
            positive = build_anima_positive_text(prompt_json)
        negative = (prompt_json.get("neg") or prompt_json.get("negative") or "").strip()

        wf["11"]["inputs"]["text"] = positive
        wf["12"]["inputs"]["text"] = negative

        # 分辨率
        width = prompt_json.get("width")
        height = prompt_json.get("height")
        aspect_ratio = (prompt_json.get("aspect_ratio") or "").strip()
        round_to = int(prompt_json.get("round_to") or self.config.round_to)

        if (width is None or height is None) and aspect_ratio:
            # 仅提供 aspect_ratio 时自动计算
            w, h = estimate_size_from_ratio(
                aspect_ratio=aspect_ratio,
                target_megapixels=float(prompt_json.get("target_megapixels") or self.config.target_megapixels),
                round_to=round_to,
            )
            width, height = w, h
        elif width is not None and height is not None:
            # 用户直接指定了 width/height，也需要对齐到 round_to 的倍数
            # 避免 "should be divisible by spatial_patch_size" 错误
            width = align_dimension(width, round_to)
            height = align_dimension(height, round_to)

        if width is None or height is None:
            # 默认方形 1MP（1024 是 16 的倍数）
            width, height = 1024, 1024

        wf["28"]["inputs"]["width"] = int(width)
        wf["28"]["inputs"]["height"] = int(height)
        wf["28"]["inputs"]["batch_size"] = int(prompt_json.get("batch_size") or 1)

        # 采样参数
        seed = prompt_json.get("seed")
        if seed is None:
            seed = int.from_bytes(uuid.uuid4().bytes[:4], "big", signed=False)
        wf["19"]["inputs"]["seed"] = int(seed)

        wf["19"]["inputs"]["steps"] = int(prompt_json.get("steps") or wf["19"]["inputs"]["steps"])
        wf["19"]["inputs"]["cfg"] = float(prompt_json.get("cfg") or wf["19"]["inputs"]["cfg"])
        wf["19"]["inputs"]["sampler_name"] = str(prompt_json.get("sampler_name") or wf["19"]["inputs"]["sampler_name"])
        wf["19"]["inputs"]["scheduler"] = str(prompt_json.get("scheduler") or wf["19"]["inputs"]["scheduler"])
        wf["19"]["inputs"]["denoise"] = float(prompt_json.get("denoise") or wf["19"]["inputs"]["denoise"])

        # 文件名前缀
        wf["52"]["inputs"]["filename_prefix"] = str(prompt_json.get("filename_prefix") or wf["52"]["inputs"]["filename_prefix"])

        return wf

    # -------------------------
    # Health check
    # -------------------------
    def check_comfyui_health(self) -> Tuple[bool, str]:
        """
        检查 ComfyUI 是否可访问。
        返回 (is_healthy, message)
        """
        try:
            url = urljoin(self.config.comfyui_url.rstrip("/") + "/", "system_stats")
            self._http_get_json(url)
            return True, f"ComfyUI 运行正常 ({self.config.comfyui_url})"
        except Exception as e:
            error_msg = str(e)
            # 提供友好的错误提示
            if "Connection refused" in error_msg or "连接" in error_msg:
                return False, (
                    f"无法连接到 ComfyUI ({self.config.comfyui_url})\n"
                    f"请确认：\n"
                    f"  1. ComfyUI 已启动\n"
                    f"  2. 地址和端口正确（可通过 COMFYUI_URL 环境变量修改）\n"
                    f"  3. 防火墙未阻止连接"
                )
            elif "timeout" in error_msg.lower() or "超时" in error_msg:
                return False, (
                    f"连接 ComfyUI 超时 ({self.config.comfyui_url})\n"
                    f"可能原因：网络延迟、ComfyUI 负载过高"
                )
            else:
                return False, f"ComfyUI 连接错误: {error_msg}"

    # -------------------------
    # ComfyUI execution
    # -------------------------
    def queue_prompt(self, prompt: Dict[str, Any]) -> str:
        url = urljoin(self.config.comfyui_url.rstrip("/") + "/", "prompt")
        payload = {"prompt": prompt, "client_id": self._client_id}
        
        try:
            resp = self._http_post_json(url, payload)
        except Exception as e:
            # 先检查 ComfyUI 是否可访问
            is_healthy, health_msg = self.check_comfyui_health()
            if not is_healthy:
                raise RuntimeError(health_msg) from e
            raise
        
        prompt_id = str(resp.get("prompt_id") or "")
        if not prompt_id:
            # 检查是否有错误信息
            error = resp.get("error") or resp.get("node_errors")
            if error:
                raise RuntimeError(f"ComfyUI 执行错误：{error}")
            raise RuntimeError(f"ComfyUI /prompt 返回异常：{resp}")
        return prompt_id

    def wait_history(self, prompt_id: str) -> Dict[str, Any]:
        url = urljoin(self.config.comfyui_url.rstrip("/") + "/", f"history/{prompt_id}")
        deadline = time.time() + float(self.config.timeout_s)
        last = None
        while time.time() < deadline:
            data = self._http_get_json(url)
            last = data
            if isinstance(data, dict) and prompt_id in data:
                return data[prompt_id]
            time.sleep(float(self.config.poll_interval_s))
        raise TimeoutError(f"等待 ComfyUI 生成超时：prompt_id={prompt_id}, last={last}")

    def _extract_images(self, prompt_id: str, history_item: Dict[str, Any]) -> List[GeneratedImage]:
        outputs = history_item.get("outputs") or {}
        images: List[GeneratedImage] = []
        for node_id, node_out in outputs.items():
            if not isinstance(node_out, dict):
                continue
            for im in node_out.get("images") or []:
                filename = str(im.get("filename") or "")
                subfolder = str(im.get("subfolder") or "")
                folder_type = str(im.get("type") or "output")
                if not filename:
                    continue
                qs = urlencode({"filename": filename, "subfolder": subfolder, "type": folder_type})
                view_url = urljoin(self.config.comfyui_url.rstrip("/") + "/", f"view?{qs}")
                images.append(
                    GeneratedImage(
                        filename=filename,
                        subfolder=subfolder,
                        folder_type=folder_type,
                        view_url=view_url,
                        saved_path=None,
                    )
                )
        return images

    def _download_images(self, images: List[GeneratedImage]) -> List[GeneratedImage]:
        """下载图片并保存到本地，同时保留原始 bytes 用于 base64 编码"""
        out_dir = Path(self.config.output_dir)
        out_dir.mkdir(parents=True, exist_ok=True)

        downloaded: List[GeneratedImage] = []
        for im in images:
            # 总是下载内容（用于 base64）
            content = self._http_get_bytes(im.view_url)
            
            saved_path = None
            if self.config.download_images:
                # 复刻 ComfyUI 的 subfolder 结构（可选）
                sub_dir = out_dir / (im.subfolder or "")
                sub_dir.mkdir(parents=True, exist_ok=True)
                dst = sub_dir / im.filename
                dst.write_bytes(content)
                saved_path = str(dst)
            
            downloaded.append(
                GeneratedImage(
                    filename=im.filename,
                    subfolder=im.subfolder,
                    folder_type=im.folder_type,
                    view_url=im.view_url,
                    saved_path=saved_path,
                    content=content,
                )
            )
        return downloaded

    def _get_mime_type(self, filename: str) -> str:
        """根据文件名推断 MIME 类型"""
        ext = Path(filename).suffix.lower()
        return {
            ".png": "image/png",
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg",
            ".webp": "image/webp",
            ".gif": "image/gif",
        }.get(ext, "image/png")

    def check_models(self) -> Tuple[bool, str]:
        """
        检查模型文件是否存在（如果配置了 COMFYUI_MODELS_DIR）。
        返回 (is_ok, message)
        """
        if not self.config.check_models or not self.config.comfyui_models_dir:
            return True, "模型检查已跳过（未配置 COMFYUI_MODELS_DIR）"
        
        all_exist, missing = self.config.check_models_exist()
        if all_exist:
            return True, "所有模型文件已就绪"
        
        missing_str = "\n".join(f"  - {m}" for m in missing)
        return False, (
            f"缺少以下模型文件：\n{missing_str}\n\n"
            f"请从 HuggingFace 下载：https://huggingface.co/circlestone-labs/Anima\n"
            f"并放置到 ComfyUI/models/ 对应子目录"
        )

    def generate(self, prompt_json: Dict[str, Any]) -> Dict[str, Any]:
        """
        输入结构化 JSON，执行生成。

        返回：
        - prompt_id
        - positive / negative（最终发送给 ComfyUI 的文本）
        - width / height
        - images: [{filename, url, file_path, base64, mime_type, markdown}]
        """
        # 预检查：模型文件
        models_ok, models_msg = self.check_models()
        if not models_ok:
            raise RuntimeError(models_msg)
        
        prompt = self._inject(prompt_json)
        prompt_id = self.queue_prompt(prompt)
        history_item = self.wait_history(prompt_id)
        images = self._extract_images(prompt_id, history_item)
        images = self._download_images(images)

        # 构建图片信息（包含多种格式）
        images_data = []
        for im in images:
            mime_type = self._get_mime_type(im.filename)
            b64 = base64.b64encode(im.content).decode("ascii") if im.content else None
            
            img_info = {
                "filename": im.filename,
                "subfolder": im.subfolder,
                "type": im.folder_type,
                # URL 格式
                "url": im.view_url,
                "view_url": im.view_url,  # 兼容旧字段
                # 本地路径
                "file_path": im.saved_path,
                "saved_path": im.saved_path,  # 兼容旧字段
                # Base64 格式（用于 MCP / Gemini / 嵌入）
                "base64": b64,
                "mime_type": mime_type,
                # Data URL（可直接用于 <img src> 或 markdown）
                "data_url": f"data:{mime_type};base64,{b64}" if b64 else None,
                # Markdown 格式（AI 可直接输出）
                "markdown": f"![{im.filename}]({im.view_url})",
            }
            images_data.append(img_info)

        # 回显最终参数（便于调试）
        actual_seed = int(prompt["19"]["inputs"]["seed"])
        actual_width = int(prompt["28"]["inputs"]["width"])
        actual_height = int(prompt["28"]["inputs"]["height"])

        result = {
            "success": True,
            "prompt_id": prompt_id,
            "positive": prompt["11"]["inputs"]["text"],
            "negative": prompt["12"]["inputs"]["text"],
            "seed": actual_seed,
            "width": actual_width,
            "height": actual_height,
            "images": images_data,
        }

        # 记录到历史
        record = self.history.add(
            params=prompt_json,
            positive_text=result["positive"],
            negative_text=result["negative"],
            prompt_id=prompt_id,
            seed=actual_seed,
            width=actual_width,
            height=actual_height,
        )
        result["history_id"] = record.id

        return result