"""Async client for Anthropic-compatible Messages APIs (with an
OpenAI-compatible chat/completions fallback).

Verified wire protocol (2026-07, against https://openrouter.ai/api and
https://api.anthropic.com):

  POST {base}/v1/messages          (or {base}/messages when base ends /v1)
  headers: x-api-key, authorization: Bearer, anthropic-version: 2023-06-01
  body:    {model, max_tokens, system, messages:[{role, content:[
               {type:"image", source:{type:"base64", media_type, data}},
               {type:"text", text}]}]}
  reply:   {content:[{type:"text", text}], stop_reason, ...}

OpenAI fallback (only on HTTP 404, or when LLM_API_STYLE=openai):
  POST {base}/chat/completions (base normalised to include /v1)
  headers: authorization: Bearer
  body:    {model, max_tokens, messages:[{role:"user", content:[
               {type:"text", text}, {type:"image_url",
                image_url:{url:"data:image/png;base64,..."}}]}]}
  reply:   {choices:[{message:{content}}]}

Retry policy: retry at most ``max_retries`` times on 408/409/429/5xx and
transport errors with exponential backoff + jitter, honouring Retry-After.
Never retry 400/401/403/404-style client faults. The API key is never logged
or included in error messages.
"""

from __future__ import annotations

import asyncio
import json
import random
import re
from dataclasses import dataclass

import httpx

from .config import Settings

ANTHROPIC_VERSION = "2023-06-01"
_RETRYABLE_STATUS = {408, 409, 429, 500, 502, 503, 504}


class LLMError(Exception):
    """User-safe LLM failure description (never contains credentials)."""

    def __init__(self, message: str, *, code: str = "llm_error",
                 status: int | None = None) -> None:
        super().__init__(message)
        self.code = code
        self.status = status


@dataclass(frozen=True)
class LLMResponse:
    text: str
    stop_reason: str | None
    style: str  # "anthropic" | "openai"


def user_turn(prompt: str, image_png: bytes | None = None) -> dict:
    """Anthropic-format user message; optionally with a PNG image block."""
    content: list[dict] = []
    if image_png is not None:
        import base64

        content.append({
            "type": "image",
            "source": {"type": "base64", "media_type": "image/png",
                       "data": base64.b64encode(image_png).decode("ascii")},
        })
    content.append({"type": "text", "text": prompt})
    return {"role": "user", "content": content}


def assistant_turn(text: str) -> dict:
    return {"role": "assistant", "content": [{"type": "text", "text": text}]}


def _messages_url(base: str) -> str:
    base = base.rstrip("/")
    return base + "/messages" if base.endswith("/v1") else base + "/v1/messages"


def _chat_completions_url(base: str) -> str:
    base = base.rstrip("/")
    return base + "/chat/completions" if base.endswith("/v1") else base + "/v1/chat/completions"


def _headers(settings: Settings, style: str) -> dict[str, str]:
    headers = {"content-type": "application/json"}
    key = settings.llm_api_key or ""
    if style == "anthropic":
        # Both header conventions are sent: Anthropic accepts x-api-key;
        # OpenRouter's anthropic-compatible endpoint accepts either.
        headers["x-api-key"] = key
        headers["authorization"] = f"Bearer {key}"
        headers["anthropic-version"] = ANTHROPIC_VERSION
    else:
        headers["authorization"] = f"Bearer {key}"
    if settings.llm_referer:
        headers["HTTP-Referer"] = settings.llm_referer
    if settings.llm_title:
        headers["X-Title"] = settings.llm_title
    return headers


def _anthropic_body(settings: Settings, system: str, messages: list[dict]) -> dict:
    return {
        "model": settings.llm_model,
        "max_tokens": settings.llm_max_tokens,
        "system": system,
        "messages": messages,
    }


def _openai_body(settings: Settings, system: str, messages: list[dict]) -> dict:
    converted: list[dict] = [{"role": "system", "content": system}]
    for item in messages:
        role = item.get("role", "user")
        parts = item.get("content")
        if isinstance(parts, str):
            converted.append({"role": role, "content": parts})
            continue
        out_parts: list[dict] = []
        for part in parts if isinstance(parts, list) else []:
            if not isinstance(part, dict):
                continue
            if part.get("type") == "text":
                out_parts.append({"type": "text", "text": part.get("text", "")})
            elif part.get("type") == "image":
                source = part.get("source") or {}
                url = f"data:{source.get('media_type', 'image/png')};base64,{source.get('data', '')}"
                out_parts.append({"type": "image_url", "image_url": {"url": url}})
        # Collapse text-only content to a plain string (canonical OpenAI shape).
        if out_parts and all(p["type"] == "text" for p in out_parts):
            converted.append({"role": role,
                              "content": "".join(p["text"] for p in out_parts)})
        else:
            converted.append({"role": role, "content": out_parts})
    return {
        "model": settings.llm_model,
        "max_tokens": settings.llm_max_tokens,
        "messages": converted,
    }


def _parse_anthropic(payload: dict) -> LLMResponse:
    content = payload.get("content")
    if not isinstance(content, list):
        raise LLMError("The model returned an unexpected response shape (no content).",
                       code="llm_bad_response")
    text = "".join(
        block.get("text", "") for block in content
        if isinstance(block, dict) and block.get("type") == "text"
    ).strip()
    if not text:
        raise LLMError("The model returned an empty response.", code="llm_bad_response")
    return LLMResponse(text=text, stop_reason=payload.get("stop_reason"), style="anthropic")


def _parse_openai(payload: dict) -> LLMResponse:
    try:
        choice = payload["choices"][0]
        text = (choice.get("message") or {}).get("content") or ""
    except (KeyError, IndexError, TypeError) as exc:
        raise LLMError("The model returned an unexpected response shape.",
                       code="llm_bad_response") from exc
    text = text.strip()
    if not text:
        raise LLMError("The model returned an empty response.", code="llm_bad_response")
    return LLMResponse(text=text, stop_reason=choice.get("finish_reason"), style="openai")


class LLMClient:
    """One configured provider client. Instantiate per job or share."""

    def __init__(self, settings: Settings, http: httpx.AsyncClient | None = None) -> None:
        self.settings = settings
        self._http = http

    async def complete_vision(
        self,
        *,
        system: str,
        messages: list[dict],
    ) -> LLMResponse:
        style_pref = self.settings.llm_api_style
        styles = ["anthropic", "openai"] if style_pref == "auto" else [style_pref]

        last_error: LLMError | None = None
        for index, style in enumerate(styles):
            try:
                return await self._call(style, system, messages)
            except LLMError as exc:
                last_error = exc
                # Only fall over to the other protocol when the endpoint
                # plainly does not speak it.
                if exc.status == 404 and index < len(styles) - 1:
                    continue
                raise
        raise last_error or LLMError("No LLM API style available.", code="llm_error")

    async def _call(self, style: str, system: str,
                    messages: list[dict]) -> LLMResponse:
        settings = self.settings
        if style == "anthropic":
            url = _messages_url(settings.llm_base_url)
            body = _anthropic_body(settings, system, messages)
        else:
            url = _chat_completions_url(settings.llm_base_url)
            body = _openai_body(settings, system, messages)

        timeout = httpx.Timeout(connect=10.0, read=settings.llm_timeout_s,
                                write=30.0, pool=10.0)
        close_client = False
        http = self._http
        if http is None:
            http = httpx.AsyncClient(timeout=timeout)
            close_client = True

        attempts = max(1, settings.llm_max_retries + 1)
        delay = 2.0
        try:
            for attempt in range(attempts):
                try:
                    response = await http.post(url, headers=_headers(settings, style), json=body)
                except httpx.HTTPError as exc:
                    if attempt + 1 >= attempts:
                        raise LLMError(
                            f"The model endpoint could not be reached ({type(exc).__name__}).",
                            code="llm_unreachable") from exc
                    await asyncio.sleep(delay + random.uniform(0, delay))
                    delay = min(delay * 4, 32.0)
                    continue

                if response.status_code == 404:
                    raise LLMError("The model endpoint returned 404 (unknown path/model).",
                                   code="llm_not_found", status=404)
                if response.status_code in (400, 401, 403):
                    detail = _safe_error_detail(response)
                    raise LLMError(
                        f"The model endpoint rejected the request "
                        f"(HTTP {response.status_code}). {detail}",
                        code="llm_rejected", status=response.status_code)
                if response.status_code in _RETRYABLE_STATUS:
                    if attempt + 1 >= attempts:
                        raise LLMError(
                            f"The model endpoint is unavailable "
                            f"(HTTP {response.status_code} after {attempts} attempts).",
                            code="llm_unavailable", status=response.status_code)
                    retry_after = response.headers.get("retry-after")
                    wait = delay + random.uniform(0, delay)
                    if retry_after:
                        try:
                            wait = max(wait, float(retry_after))
                        except ValueError:
                            pass
                    await asyncio.sleep(wait)
                    delay = min(delay * 4, 32.0)
                    continue
                if response.status_code >= 400:
                    raise LLMError(
                        f"The model endpoint returned HTTP {response.status_code}.",
                        code="llm_error", status=response.status_code)

                try:
                    payload = response.json()
                except json.JSONDecodeError as exc:
                    # Some upstream gateways occasionally return an HTML or
                    # plain-text error body with HTTP 200 after a long model
                    # wait. Treat that exactly like the transient transport
                    # failures above while the configured retry budget
                    # remains; never expose or log the provider body.
                    if attempt + 1 < attempts:
                        await asyncio.sleep(delay + random.uniform(0, delay))
                        delay = min(delay * 4, 32.0)
                        continue
                    raise LLMError("The model endpoint returned non-JSON.",
                                   code="llm_bad_response") from exc
                parsed = _parse_anthropic(payload) if style == "anthropic" else _parse_openai(payload)
                if parsed.stop_reason in {"max_tokens", "length"}:
                    raise LLMError(
                        "The model's reply was truncated (max_tokens reached). "
                        "Increase LLM_MAX_TOKENS or simplify the subject.",
                        code="llm_truncated")
                return parsed
        finally:
            if close_client:
                await http.aclose()
        raise LLMError("The model call failed unexpectedly.", code="llm_error")


def _safe_error_detail(response: httpx.Response) -> str:
    """Extract a short, credential-free detail string from an error body."""
    try:
        payload = response.json()
        message = payload.get("error", {})
        if isinstance(message, dict):
            message = message.get("message") or message.get("type") or ""
        if isinstance(message, str) and message:
            return message[:200]
    except Exception:
        pass
    return ""


_FENCE_RE = re.compile(r"^```(?:json)?\s*|\s*```$", re.MULTILINE)


def extract_json_object(text: str) -> dict:
    """Parse a JSON object from an LLM reply, tolerating markdown fences
    and leading/trailing prose. Raises LLMError on failure."""
    candidate = _FENCE_RE.sub("", text).strip()
    try:
        parsed = json.loads(candidate)
    except json.JSONDecodeError:
        start = candidate.find("{")
        end = candidate.rfind("}")
        if start == -1 or end == -1 or end <= start:
            raise LLMError("The model did not return a JSON object.",
                           code="llm_bad_json")
        try:
            parsed = json.loads(candidate[start:end + 1])
        except json.JSONDecodeError as exc:
            raise LLMError(f"The model returned malformed JSON ({exc.msg}).",
                           code="llm_bad_json") from exc
    if not isinstance(parsed, dict):
        raise LLMError("The model returned JSON that is not an object.",
                       code="llm_bad_json")
    return parsed