"""AI 生图 API：生成结果存储 + 手表拉取端点 + 共享生图流程。

生图走两条既有链路（与聊天一致，前端无需新协议）：
1. 阻塞式：POST /v1/xtc/chat —— 生图模型直接返回带 image_url 的 XTC 响应
   （适合非手表端 / 生成快的模型）
2. 伪流式：POST /v1/xtc/chat/pseudo/start —— 后台生图，轮询到 done 后拉取
   （手表默认路径，防 fetch.fetch 30s 超时）

本模块提供：
- store_image() / get_image()：处理后图片的内存存储（TTL + 容量上限，
  与 TTS 会话一致；手表生成后立即拉取，无需跨重启持久化）
- run_image_generation()：完整生图流程（调上游 → 手表化处理 → 存储 → 组装 XTC 字段），
  供 chat / pseudo_stream 两个入口复用
- GET /v1/xtc/image/gen/{image_id}：手表拉取处理后的 JPEG（capability URL，
  与 TTS audio 一致：image_id 为 128bit 随机 hex 不可猜测 + TTL，免 access_key，
  因为手表 request.download 不便携带自定义 header）
"""
from __future__ import annotations

import logging
import re
import time
import uuid
from typing import Optional

from fastapi import APIRouter, Response

from ..adapters import imagegen as imagegen_adapter
from ..errors import HttpError
from ..media.imagegen import DEFAULT_JPEG_QUALITY, DEFAULT_MAX_DIM, process_image_for_watch
from ..providers.policy import is_image_model
from ..services.config_store import load_config
from ..services.prompt_optimizer import optimize_image_prompt
from ._common import CORS_HEADERS, ok_with_cors, select_provider_and_key

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/v1/xtc/image", tags=["imagegen"])

# 图片会话存活 30 分钟（生图比 TTS 慢，手表断网重试窗口也要留足）
_IMAGE_TTL = 1800
# 内存上限：处理后图片典型 50~100KB，32 张约 2~3MB，HF Space 内存可控
_IMAGE_MAX = 32
_images: dict[str, dict] = {}


def _cleanup_images() -> None:
    """清理过期图片 + 超限时淘汰最老的（与 TTS 缓存同策略）。"""
    now = time.time()
    expired = [k for k, v in _images.items() if now - v["created_at"] > _IMAGE_TTL]
    for k in expired:
        _images.pop(k, None)
    if len(_images) > _IMAGE_MAX:
        for k, _ in sorted(_images.items(), key=lambda kv: kv[1]["created_at"])[
            : len(_images) - _IMAGE_MAX
        ]:
            _images.pop(k, None)


def store_image(image_bytes: bytes, meta: Optional[dict] = None) -> str:
    """存入处理后图片，返回 image_id（随机 hex，用于 capability URL）。"""
    _cleanup_images()
    image_id = uuid.uuid4().hex
    _images[image_id] = {
        "bytes": image_bytes,
        "meta": meta or {},
        "created_at": time.time(),
    }
    return image_id


def get_image(image_id: str) -> Optional[dict]:
    entry = _images.get(image_id)
    if not entry:
        return None
    if time.time() - entry["created_at"] > _IMAGE_TTL:
        _images.pop(image_id, None)
        return None
    return entry


def image_url_for(image_id: str) -> str:
    """手表端可拉取的相对 URL（前端拼 baseUrl）。"""
    return f"/v1/xtc/image/gen/{image_id}"


@router.get("/gen/{image_id}")
async def download_generated_image(image_id: str) -> Response:
    """手表拉取处理后的生图 JPEG（capability URL，无需鉴权，见模块头注释）。"""
    entry = get_image(image_id)
    if not entry:
        raise HttpError("image not found or expired", status=404, code="not_found")
    meta = entry.get("meta") or {}
    return Response(
        content=entry["bytes"],
        media_type="image/jpeg",
        headers={
            **CORS_HEADERS,
            "Cache-Control": "no-store",
            "x-xtc-image-width": str(meta.get("width") or ""),
            "x-xtc-image-height": str(meta.get("height") or ""),
        },
    )


# ===== 共享生图流程（chat / pseudo 复用）=====

# 附件块：normalize_chat_body 会把文本类附件追加成独立的一条 user 消息，
# 形如 "[文件名]\n<内容>"（多个用空行连接）。生图 prompt 必须能把它和用户自己打的字区分开，
# 否则「用户打了提示词 + 顺手带了个文件」时，生图/优化拿到的是文件内容而不是那句话。
_ATTACHMENT_HEAD_RE = re.compile(r"^\[[^\]\n]{1,120}\]\n")


def _split_attachment(text: str) -> tuple[bool, str]:
    """返回 (是否附件块, 去掉文件名头的内容)。"""
    m = _ATTACHMENT_HEAD_RE.match(text or "")
    if not m:
        return False, text or ""
    return True, (text or "")[m.end():]


def extract_image_prompt(messages: list) -> str:
    """从 messages 提取生图 prompt。

    取「最近一条用户自己打的文字」；如果还有附件块，把附件内容作为参考附在末尾
    （截断到 600 字）。只有附件没有文字时，附件内容就是 prompt（此时它就是用户
    表达要画什么的方式）。
    """
    typed = ""
    attach = ""
    for msg in reversed(messages or []):
        if not isinstance(msg, dict) or msg.get("role") != "user":
            continue
        content = msg.get("content")
        text = ""
        if isinstance(content, str):
            text = content.strip()
        elif isinstance(content, list):
            parts = [
                str(p.get("text") or "").strip()
                for p in content
                if isinstance(p, dict) and p.get("type") == "text"
            ]
            text = "\n".join(x for x in parts if x).strip()
        if not text:
            continue
        is_attach, body = _split_attachment(text)
        if is_attach:
            if not attach:
                attach = body.strip()
            continue
        if not typed:
            typed = text
            break
    if typed and attach:
        return (typed + "\n\n参考内容：\n" + attach[:600]).strip()
    if typed:
        return typed
    return attach[:2000]


async def run_image_generation(
    *,
    provider,
    api_key: str,
    clean_model: str,
    messages: list,
    size: Optional[str] = None,
    access_key: Optional[str] = None,
    optimize: bool = False,
    optimizer_provider: Optional[str] = None,
    optimizer_model: Optional[str] = None,
    optimize_context: Optional[str] = None,
) -> dict:
    """完整生图流程：调上游 → 手表化处理 → 存储 → 返回 XTC 响应附加字段。

    optimize=True 时先调用「优化模型」把用户提示词润色为专业生图提示词，
    再送上游生成；优化模型缺失/调用失败自动回退原始 prompt（不阻断生图），
    并通过 optimize_warning 反馈前端。

    返回 dict 可直接并入 XTC 响应 / 伪流式会话结果：
    {image, image_id, image_url, image_width, image_height, text, usage,
     optimized, optimized_prompt?, original_prompt?, optimize_warning?}
    """
    if provider.type == "gemini":
        raise HttpError(
            "Image generation is only supported on openai-compatible providers",
            status=400,
            code="bad_request",
        )
    prompt = extract_image_prompt(messages)
    if not prompt:
        raise HttpError("prompt is required for image generation", status=400, code="bad_request")

    gen_prompt = prompt
    optimized_prompt: Optional[str] = None
    optimize_warning = ""
    if optimize:
        try:
            clean_opt = await _optimize_prompt(
                access_key=access_key,
                optimizer_provider=optimizer_provider,
                optimizer_model=optimizer_model,
                prompt_text=prompt,
                context_text=optimize_context or "",
            )
            if clean_opt:
                optimized_prompt = clean_opt
                gen_prompt = clean_opt
        except HttpError as e:
            optimize_warning = "提示词优化失败：" + str(e.message or e.code or "unknown")
            logger.warning("[imagegen] optimize failed (fallback raw prompt) provider=%s model=%s: %s",
                           provider.id, clean_model, e.message)
        except Exception as e:
            optimize_warning = "提示词优化异常：" + str(e)
            logger.warning("[imagegen] optimize error (fallback raw prompt) provider=%s model=%s",
                           provider.id, clean_model, exc_info=True)

    gen = await imagegen_adapter.generate_image(
        base_url=provider.base_url or "",
        api_key=api_key,
        model=clean_model,
        prompt=gen_prompt,
        n=1,
        size=size,
    )
    processed = process_image_for_watch(gen["image_bytes"])
    image_id = store_image(
        processed["image_bytes"],
        meta={
            "width": processed["width"],
            "height": processed["height"],
            "orig_width": processed["orig_width"],
            "orig_height": processed["orig_height"],
        },
    )
    revised = str(gen.get("revised_prompt") or "").strip()
    logger.info(
        "[imagegen] done provider=%s model=%s %dx%d->%dx%d %dbytes optimize=%s",
        provider.id, clean_model,
        processed["orig_width"], processed["orig_height"],
        processed["width"], processed["height"], len(processed["image_bytes"]),
        "on" if optimized_prompt else "off",
    )
    out = {
        "image": True,
        "image_id": image_id,
        "image_url": image_url_for(image_id),
        "image_width": processed["width"],
        "image_height": processed["height"],
        # 手表消息列表需要一条占位文本；revised_prompt 有值时附带给前端展示
        "revised_prompt": revised,
        "text": "图片已生成",
        "usage": gen.get("usage") or {},
    }
    if optimize:
        out["optimized"] = bool(optimized_prompt)
        if optimized_prompt:
            out["optimized_prompt"] = optimized_prompt
            out["original_prompt"] = prompt
        if optimize_warning:
            out["optimize_warning"] = optimize_warning
    return out


async def _optimize_prompt(
    *,
    access_key: Optional[str],
    optimizer_provider: Optional[str],
    optimizer_model: Optional[str],
    prompt_text: str,
    context_text: str,
) -> str:
    """把原始提示词交给用户指定的聊天模型（或最近用的聊天模型）润色。

    返回优化后的提示词；模型/厂商缺失时抛 HttpError 由调用方回退原始 prompt。
    """
    opt_model = str(optimizer_model or "").strip()
    if not opt_model:
        raise HttpError("optimizer model not specified", status=400, code="bad_request")
    config = await load_config()
    opt_provider, opt_key, clean_opt_model = select_provider_and_key(
        config=config,
        provider_id=(str(optimizer_provider or "").strip() or None),
        model=opt_model,
        fallback_model=None,
    )
    return await optimize_image_prompt(
        provider=opt_provider,
        api_key=opt_key,
        clean_model=clean_opt_model,
        prompt_text=prompt_text,
        context_text=context_text,
        access_key=access_key,
    )


def provider_supports_image(provider, model: str) -> bool:
    """select_provider_and_key 之后调用：判断是否走生图流程。"""
    return is_image_model(provider, model)


__all__ = [
    "router",
    "DEFAULT_MAX_DIM",
    "DEFAULT_JPEG_QUALITY",
    "store_image",
    "get_image",
    "image_url_for",
    "extract_image_prompt",
    "run_image_generation",
    "provider_supports_image",
]