"""
ABOS AI Worker – OpenRouter LLM Client
========================================
Thin, async HTTP wrapper around the OpenRouter chat-completions API.

Design decisions
----------------
* Uses httpx (async) for non-blocking I/O — fits naturally into FastAPI's
  async event loop without spinning up a thread pool executor.
* Strict JSON mode is requested via the `response_format` parameter when the
  model supports it; a fallback regex-strip is applied for models that don't.
* The OpenRouter key is passed in per-request and is NEVER stored or logged.
"""

from __future__ import annotations

import json
import logging
import re
from typing import Any

import httpx

from app.config import settings

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# JSON extraction helpers
# ---------------------------------------------------------------------------

_JSON_BLOCK_RE = re.compile(r"```(?:json)?\s*([\s\S]*?)\s*```", re.IGNORECASE)


def _strip_markdown_json(text: str) -> str:
    """Remove markdown fences that some models wrap JSON output in."""
    match = _JSON_BLOCK_RE.search(text)
    if match:
        return match.group(1).strip()
    return text.strip()


# ---------------------------------------------------------------------------
# Main client function
# ---------------------------------------------------------------------------

async def call_openrouter(
    *,
    openrouter_key: str,
    prompt: str,
    document_text: str,
    model: str | None = None,
) -> tuple[dict[str, Any], str]:
    """
    Send *document_text* + *prompt* to OpenRouter and return parsed JSON.

    Parameters
    ----------
    openrouter_key:
        Bearer token for the OpenRouter API (supplied per-request by Spring Boot).
    prompt:
        Extraction instruction for the LLM.
    document_text:
        Full text extracted from the PDF, passed as user context.
    model:
        Optional model override. Falls back to ``settings.openrouter_default_model``.

    Returns
    -------
    (data_points, model_used)
        data_points  – parsed dict of AI-extracted data.
        model_used   – the model identifier that was actually used.

    Raises
    ------
    httpx.HTTPStatusError
        If OpenRouter returns a non-2xx response.
    ValueError
        If the LLM response cannot be parsed as JSON.
    """
    resolved_model: str = model or settings.openrouter_default_model

    system_message = (
        "You are a precise data extraction assistant. "
        "You MUST respond with ONLY a valid JSON object — no prose, "
        "no markdown fences, no explanation. "
        "If a value is not found in the document, use null for that key."
    )

    user_message = (
        f"{prompt}\n\n"
        "--- DOCUMENT TEXT START ---\n"
        f"{document_text}\n"
        "--- DOCUMENT TEXT END ---"
    )

    payload: dict[str, Any] = {
        "model": resolved_model,
        "messages": [
            {"role": "system", "content": system_message},
            {"role": "user", "content": user_message},
        ],
        "temperature": 0.0,   # Deterministic output for data extraction
        "max_tokens": 4096,
        # Ask for JSON output when the API supports it
        "response_format": {"type": "json_object"},
    }

    headers = {
        "Authorization": f"Bearer {openrouter_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://abos.erp",          # Required by OpenRouter
        "X-Title": "ABOS AI Worker",                  # Shown in OpenRouter dashboard
    }

    logger.info("Calling OpenRouter model=%s …", resolved_model)

    async with httpx.AsyncClient(timeout=settings.openrouter_timeout) as client:
        response = await client.post(
            f"{settings.openrouter_base_url}/chat/completions",
            headers=headers,
            json=payload,
        )

    # Raise for HTTP-level errors (4xx / 5xx)
    try:
        response.raise_for_status()
    except httpx.HTTPStatusError as exc:
        logger.error(
            "OpenRouter returned HTTP %d: %s",
            exc.response.status_code,
            exc.response.text,
        )
        raise

    body = response.json()

    # ── Extract the assistant message ─────────────────────────────────────
    try:
        raw_content: str = body["choices"][0]["message"]["content"]
    except (KeyError, IndexError) as exc:
        raise ValueError(
            f"Unexpected OpenRouter response structure: {body}"
        ) from exc

    logger.debug("Raw LLM response: %s", raw_content[:500])

    # ── Parse JSON from the response ─────────────────────────────────────
    cleaned = _strip_markdown_json(raw_content)
    try:
        data_points: dict[str, Any] = json.loads(cleaned)
    except json.JSONDecodeError as exc:
        raise ValueError(
            f"LLM did not return valid JSON. Raw content: {raw_content[:300]}"
        ) from exc

    # Determine the model that was actually used (OpenRouter may auto-route)
    model_used: str = body.get("model", resolved_model)

    logger.info("OpenRouter extraction complete. model_used=%s", model_used)
    return data_points, model_used


async def call_openrouter_text(
    *,
    openrouter_key: str,
    system_message: str,
    user_message: str,
    model: str | None = None,
) -> tuple[str, str]:
    """
    Plain single-shot prose completion — no JSON-object response format, no
    document-text wrapping. Used for free-text generation (e.g. drafting/polishing a
    description) where the layered extraction-with-fallback strategy in
    call_openrouter_layered would be the wrong tool (that's for recovering missing
    *structured fields*, not for writing prose).
    """
    resolved_model: str = model or settings.openrouter_default_model

    payload: dict[str, Any] = {
        "model": resolved_model,
        "messages": [
            {"role": "system", "content": system_message},
            {"role": "user", "content": user_message},
        ],
        "temperature": 0.4,
        "max_tokens": 600,
    }

    headers = {
        "Authorization": f"Bearer {openrouter_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://abos.erp",
        "X-Title": "ABOS AI Worker",
    }

    logger.info("Calling OpenRouter (text) model=%s …", resolved_model)

    async with httpx.AsyncClient(timeout=settings.openrouter_timeout) as client:
        response = await client.post(
            f"{settings.openrouter_base_url}/chat/completions",
            headers=headers,
            json=payload,
        )

    try:
        response.raise_for_status()
    except httpx.HTTPStatusError as exc:
        logger.error("OpenRouter returned HTTP %d: %s", exc.response.status_code, exc.response.text)
        raise

    body = response.json()
    try:
        content: str = body["choices"][0]["message"]["content"]
    except (KeyError, IndexError) as exc:
        raise ValueError(f"Unexpected OpenRouter response structure: {body}") from exc

    model_used: str = body.get("model", resolved_model)
    return content.strip(), model_used


# ---------------------------------------------------------------------------
# Multi-layer extraction strategy
# ---------------------------------------------------------------------------

# Statement-level fields worth a fallback retry if the primary pass leaves them null.
# (Deliberately excludes `transactions`, handled separately since it's a list, not a scalar.)
_GAP_FILL_FIELDS: tuple[str, ...] = (
    "accountHolder",
    "accountNumber",
    "statementPeriod",
    "openingBalance",
    "closingBalance",
    "currency",
)


def _find_gaps(data: dict[str, Any]) -> list[str]:
    """Return the names of required fields still missing/null/empty, transactions included."""
    gaps = [key for key in _GAP_FILL_FIELDS if not data.get(key)]
    if not data.get("transactions"):
        gaps.append("transactions")
    return gaps


async def call_openrouter_layered(
    *,
    openrouter_key: str,
    prompt: str,
    document_text: str,
    model: str | None = None,
    fallback_models: list[str] | None = None,
) -> tuple[dict[str, Any], str, list[str]]:
    """
    Multi-layer extraction: run the primary model, then — only for whatever fields are still
    missing afterward — retry against each fallback model in turn, merging in the first
    non-empty value found for each gap. Stops as soon as no gaps remain or the fallback chain
    is exhausted.

    This exists because a single LLM pass can plausibly nail the transaction table but miss a
    header field (or vice versa) depending on how the PDF's text layer is laid out; retrying the
    *whole* document against a second model for just the missing pieces recovers those fields
    without discarding what the primary pass already got right.

    Returns
    -------
    (data_points, primary_model_used, layers_used)
        data_points        – merged dict of extracted fields (primary values take precedence).
        primary_model_used – the model identifier the primary pass actually used.
        layers_used        – every model identifier that contributed at least one field,
                              in the order they were tried (always includes the primary).
    """
    data, primary_model_used = await call_openrouter(
        openrouter_key=openrouter_key,
        prompt=prompt,
        document_text=document_text,
        model=model,
    )
    layers_used = [primary_model_used]
    gaps = _find_gaps(data)

    candidates = [m for m in (fallback_models or []) if m and m != primary_model_used]
    for fallback_model in candidates:
        if not gaps:
            break
        gap_prompt = (
            f"{prompt}\n\n"
            "IMPORTANT: A previous extraction pass on this exact document left the following "
            f"field(s) empty or missing: {', '.join(gaps)}. Re-read the document text carefully "
            "and return the COMPLETE JSON object again, but pay special attention to recovering "
            "these specific fields — only use null for a field if it is genuinely absent from "
            "the document."
        )
        try:
            fallback_data, fallback_model_used = await call_openrouter(
                openrouter_key=openrouter_key,
                prompt=gap_prompt,
                document_text=document_text,
                model=fallback_model,
            )
        except Exception as exc:  # noqa: BLE001 — a failed fallback layer must not fail the request
            logger.warning("Fallback layer model=%s failed, skipping: %s", fallback_model, exc)
            continue

        filled_any = False
        for key in _GAP_FILL_FIELDS:
            if not data.get(key) and fallback_data.get(key):
                data[key] = fallback_data[key]
                filled_any = True
        if not data.get("transactions") and fallback_data.get("transactions"):
            data["transactions"] = fallback_data["transactions"]
            filled_any = True

        if filled_any:
            layers_used.append(fallback_model_used)
        gaps = _find_gaps(data)

    if len(layers_used) > 1:
        logger.info(
            "Layered extraction used %d layer(s): %s (remaining gaps: %s)",
            len(layers_used), layers_used, gaps or "none",
        )

    return data, primary_model_used, layers_used