from __future__ import annotations

import asyncio
import base64
import io
import math
import random
from pathlib import Path
from typing import TYPE_CHECKING, Any, Awaitable, Callable, Protocol

from openai import APIConnectionError, APIStatusError, APITimeoutError, AsyncOpenAI
from PIL import Image

from aicd.config import AppConfig, LLMConfig, effective_llm_config

if TYPE_CHECKING:
    from aicd.session_log import SessionLogger

_VISION_MIME_BY_EXT: dict[str, str] = {
    ".png": "image/png",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".webp": "image/webp",
    ".gif": "image/gif",
}


def _lanczos() -> int:
    try:
        return Image.Resampling.LANCZOS  # type: ignore[attr-defined]
    except AttributeError:
        return Image.LANCZOS


def _is_retryable_llm_error(exc: BaseException) -> bool:
    if isinstance(exc, (APITimeoutError, APIConnectionError)):
        return True
    if isinstance(exc, APIStatusError):
        return exc.status_code in (429, 502, 503)
    return False


def image_file_to_data_url_with_meta(
    path: Path,
    *,
    max_long_edge: int = 1920,
    compress_if_source_ge_bytes: int = 1024 * 1024,
    max_encoded_bytes: int = 1024 * 1024,
    hard_max_source_bytes: int = 20 * 1024 * 1024,
) -> tuple[str, dict[str, int]]:
    """
    返回 ``(data_url, meta)``。``meta`` 含 **source_***（磁盘原图像素）与 **sent_***（实际送入多模态 API 的 PNG 像素），
    供 **image_crop** 将模型在「缩小图」上给出的 bbox 映射回原图。
    """
    suf = path.suffix.lower()
    if suf not in _VISION_MIME_BY_EXT:
        raise ValueError(
            f"unsupported image extension {path.suffix!r} (use png/jpg/jpeg/webp/gif)"
        )
    raw = path.read_bytes()
    if len(raw) > hard_max_source_bytes:
        raise ValueError(
            f"image file too large (>{hard_max_source_bytes // (1024 * 1024)}MB): {path}"
        )
    im = Image.open(io.BytesIO(raw))
    im.load()
    if im.mode not in ("RGB", "RGBA"):
        im = im.convert("RGB")
    source_width, source_height = im.size
    w, h = source_width, source_height
    long_edge = max(w, h)
    scale = 1.0
    if long_edge > max_long_edge:
        scale = min(scale, max_long_edge / long_edge)
    if len(raw) >= compress_if_source_ge_bytes:
        scale = min(scale, math.sqrt(compress_if_source_ge_bytes / len(raw)))
    nw, nh = max(1, int(w * scale)), max(1, int(h * scale))
    if nw != w or nh != h:
        im = im.resize((nw, nh), _lanczos())

    min_side_floor = 128
    for _ in range(16):
        buf = io.BytesIO()
        im.save(buf, format="PNG", optimize=True)
        data = buf.getvalue()
        if len(data) <= max_encoded_bytes or min(im.size) <= min_side_floor:
            break
        nw2 = max(1, int(im.width * 0.85))
        nh2 = max(1, int(im.height * 0.85))
        if nw2 == im.width and nh2 == im.height:
            break
        im = im.resize((nw2, nh2), _lanczos())

    b64 = base64.standard_b64encode(data).decode("ascii")
    meta = {
        "source_width": source_width,
        "source_height": source_height,
        "sent_width": im.width,
        "sent_height": im.height,
    }
    return "data:image/png;base64," + b64, meta


def image_file_to_data_url(
    path: Path,
    *,
    max_long_edge: int = 1920,
    compress_if_source_ge_bytes: int = 1024 * 1024,
    max_encoded_bytes: int = 1024 * 1024,
    hard_max_source_bytes: int = 20 * 1024 * 1024,
) -> str:
    """
    将本地图片转为 ``data:image/...;base64,...``，供 OpenAI 兼容多模态接口使用。

    - 任一边长大于 ``max_long_edge``（默认 1920）时按长边缩放到不超过该值。
    - 源文件不小于 ``compress_if_source_ge_bytes``（默认 1MiB）时参与缩放（按像素与体积启发式），
      且编码为 PNG 后若仍大于 ``max_encoded_bytes``（默认 1MiB）则迭代缩小直至满足或达到最小边 128。
    """
    url, _ = image_file_to_data_url_with_meta(
        path,
        max_long_edge=max_long_edge,
        compress_if_source_ge_bytes=compress_if_source_ge_bytes,
        max_encoded_bytes=max_encoded_bytes,
        hard_max_source_bytes=hard_max_source_bytes,
    )
    return url


class LLMClient(Protocol):
    async def chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
    ) -> Any:
        ...


def create_llm_client(app_cfg: AppConfig) -> "OpenAICompatibleClient":
    """持有 ``AppConfig`` 引用，便于运行时调整 ``cfg.llm`` 采样参数与模型名。"""
    return OpenAICompatibleClient(app_cfg)


class OpenAICompatibleClient:
    """三种协议均走 OpenAI SDK；按当前 ``cfg.llm`` 组装 ``chat.completions`` 参数。"""

    def __init__(self, app_cfg: AppConfig) -> None:
        self._app_cfg = app_cfg
        self._client: AsyncOpenAI | None = None
        self._client_sig: tuple[str, str] | None = None
        self._session_log: SessionLogger | None = None

    def attach_session_log(self, slog: SessionLogger | None) -> None:
        """记录 LLM 重试等事件到会话日志（可选）。"""
        self._session_log = slog

    def _get_client(self) -> AsyncOpenAI:
        eff = effective_llm_config(self._app_cfg.llm)
        sig = (eff.base_url, eff.api_key or "")
        if self._client is None or self._client_sig != sig:
            kwargs: dict[str, Any] = {"api_key": eff.api_key or "dummy"}
            if eff.base_url:
                kwargs["base_url"] = eff.base_url
            self._client = AsyncOpenAI(**kwargs)
            self._client_sig = sig
        return self._client

    def _completion_kwargs(
        self, llm: LLMConfig, eff: LLMConfig
    ) -> dict[str, Any]:
        body: dict[str, Any] = {
            "model": eff.model,
            "timeout": eff.timeout_sec,
        }
        if llm.temperature is not None:
            body["temperature"] = float(llm.temperature)
        if llm.top_p is not None:
            body["top_p"] = float(llm.top_p)
        if llm.frequency_penalty is not None:
            body["frequency_penalty"] = float(llm.frequency_penalty)
        if llm.presence_penalty is not None:
            body["presence_penalty"] = float(llm.presence_penalty)
        if llm.max_output_tokens is not None:
            body["max_tokens"] = int(llm.max_output_tokens)
        extra: dict[str, Any] = {}
        if llm.extra_body:
            extra.update(llm.extra_body)
        if llm.repetition_penalty is not None:
            extra.setdefault("repetition_penalty", float(llm.repetition_penalty))
        if extra:
            body["extra_body"] = extra
        return body

    def _log_retry(
        self, operation: str, attempt_index: int, total_attempts: int, err: str, delay: float
    ) -> None:
        slog = self._session_log
        if slog:
            slog.write_event(
                "llm",
                "OpenAICompatibleClient",
                operation,
                f"retry {attempt_index}/{total_attempts} sleep={delay:.2f}s",
                err[:2000],
            )

    async def _call_chat_completions_with_retry(
        self,
        operation: str,
        call: Callable[[], Awaitable[Any]],
    ) -> Any:
        llm = self._app_cfg.llm
        max_extra = max(0, min(20, int(llm.max_retries)))
        total = max_extra + 1
        base = float(llm.retry_base_delay_sec)
        if base <= 0:
            base = 0.1
        last: BaseException | None = None
        for attempt in range(total):
            try:
                return await call()
            except Exception as e:  # noqa: BLE001
                last = e
                if attempt >= total - 1 or not _is_retryable_llm_error(e):
                    raise
                delay = base * (2**attempt) + random.uniform(
                    0, max(0.05, base * 0.25)
                )
                self._log_retry(
                    operation,
                    attempt + 1,
                    total,
                    repr(e),
                    delay,
                )
                await asyncio.sleep(delay)
        assert last is not None
        raise last

    async def chat(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None = None,
    ) -> Any:
        client = self._get_client()
        eff = effective_llm_config(self._app_cfg.llm)
        llm = self._app_cfg.llm
        kw = self._completion_kwargs(llm, eff)

        async def _once() -> Any:
            return await client.chat.completions.create(
                messages=messages,
                tools=tools or [],
                **kw,
            )

        return await self._call_chat_completions_with_retry("chat.completions", _once)

    async def chat_with_images(
        self, user_prompt: str, image_paths: list[Path]
    ) -> dict[str, Any]:
        """
        单次用户消息中含多张图片（OpenAI 兼容 ``image_url`` data URL）。
        模型名：``llm.vision_model`` 或回退 ``llm.model``。

        返回 ``{"analysis": str, "per_image": [{path, source_width, source_height, sent_width, sent_height}, ...]}``。
        **sent_*** 为实际送入 API 的像素尺寸；插入 Word 等应仍使用 **path** 指向的**原始文件**，
        若模型按像素报 bbox，应对 **原图** 调用 **image_crop** 并传入 **bbox_reference_width/height = sent_***。
        """
        if not image_paths:
            raise ValueError("image_paths must be non-empty")
        if len(image_paths) > 10:
            raise ValueError("at most 10 images per vision request")
        eff = effective_llm_config(self._app_cfg.llm)
        llm = self._app_cfg.llm
        model = (llm.vision_model or eff.model).strip()
        content: list[dict[str, Any]] = [{"type": "text", "text": user_prompt}]
        per_image: list[dict[str, Any]] = []
        for p in image_paths:
            pr = p.resolve()
            url, meta = image_file_to_data_url_with_meta(pr)
            per_image.append(
                {
                    "path": str(pr),
                    "source_width": meta["source_width"],
                    "source_height": meta["source_height"],
                    "sent_width": meta["sent_width"],
                    "sent_height": meta["sent_height"],
                }
            )
            content.append({"type": "image_url", "image_url": {"url": url}})
        messages: list[dict[str, Any]] = [{"role": "user", "content": content}]
        client = self._get_client()
        kw = self._completion_kwargs(llm, eff)
        kw["model"] = model

        async def _once() -> Any:
            return await client.chat.completions.create(messages=messages, **kw)

        resp = await self._call_chat_completions_with_retry(
            "chat.completions.vision",
            _once,
        )
        msg = resp.choices[0].message
        text = (msg.content or "").strip()
        return {"analysis": text, "per_image": per_image}
