"""Landing.AI ADE parse client.

`parse_document` submits one PDF or image through the async Parse Jobs API,
polls the job to completion, and returns the full parse response as a dict.
Polling is the only completion signal the platform offers; there is no
webhook. Results over the inline size limit arrive as a presigned URL that
expires shortly after the poll that returned it, so the download happens
immediately.

Two API generations are in play. The generally-available dpt-2 model line
rides the v1 jobs API; the preview dpt-3 line rides the v2 jobs API. The
model name picks the route.

`custom_prompts` passes per-chunk-type instruction to the v1 model line, keyed
by chunk type (`figure`, `text`, ...). It cannot change the shape of what comes
back — an instruction to emit a chart as an HTML table returns the same prose
description — so it is for content-level instruction only, and it is the
caller's to supply.

A parse that reports failed pages raises instead of returning: a silently
missing page in a financial document is a liability, so a partial parse
never looks like a success.
"""

from __future__ import annotations

import json
import time
import urllib.request
from pathlib import Path
from typing import Any, Optional

from loguru import logger

from quber.settings import get_settings

#: Generally-available default. The dpt-3 preview names route to the v2 API.
DEFAULT_MODEL = "dpt-2"

POLL_SECONDS = 5.0
TIMEOUT_SECONDS = 1800.0


def parse_document(
    source: Path,
    model: str = DEFAULT_MODEL,
    custom_prompts: Optional[dict[str, str]] = None,
    poll_seconds: float = POLL_SECONDS,
) -> dict[str, Any]:
    """Parse one local PDF or image through ADE; return the raw parse response.

    `custom_prompts` reaches the v1 model line only; the v2 API takes no such
    argument, so it is dropped for a dpt-3 model. `poll_seconds` sets how often
    the job is polled — a single-page submission finishes in seconds and wants
    a shorter interval than a whole document.
    """
    client = build_client()
    if is_v2_model(model):
        jobs = client.v2.parse_jobs
        job = jobs.create(document=source, model=model)
    else:
        jobs = client.parse_jobs
        extra = {"custom_prompts": custom_prompts} if custom_prompts else {}
        job = jobs.create(document=source, model=model, **extra)
    logger.info(
        "ADE parse job {} submitted: {} ({} bytes, model={})",
        job.job_id,
        source.name,
        source.stat().st_size,
        model,
    )
    return _await_result(jobs, job.job_id, poll_seconds)


def is_v2_model(model: str) -> bool:
    """True for the dpt-3 model line, which lives on the v2 parse API."""
    return model.startswith("dpt-3")


def build_client() -> Any:
    """An authenticated SDK client, keyed from settings (ADE_API_KEY)."""
    from landingai_ade import LandingAIADE

    key = get_settings().landing.api_key
    if not key:
        raise RuntimeError("No Landing.AI key: set ADE_API_KEY in the environment or .env")
    return LandingAIADE(apikey=key)


def _await_result(jobs: Any, job_id: str, poll_seconds: float = POLL_SECONDS) -> dict[str, Any]:
    """Poll one parse job to completion and return its response dict."""
    t0 = time.time()
    last = None
    while True:
        r = jobs.get(job_id)
        state = (r.status, r.progress)
        if state != last:
            logger.info(
                "ADE job {}: status={} progress={:.0%} elapsed={:.0f}s",
                job_id,
                r.status,
                r.progress or 0.0,
                time.time() - t0,
            )
            last = state
        if r.status in ("completed", "failed", "cancelled"):
            break
        if time.time() - t0 > TIMEOUT_SECONDS:
            raise RuntimeError(f"ADE job {job_id} timed out after {TIMEOUT_SECONDS:.0f}s")
        time.sleep(poll_seconds)

    if r.status != "completed":
        # v1 jobs report a failure_reason string; v2 jobs report an error object.
        reason = getattr(r, "failure_reason", None) or getattr(r, "error", None)
        raise RuntimeError(f"ADE job {job_id} {r.status}: {reason}")

    # The completed payload is `data` on a v1 job and `result` on a v2 job.
    payload = getattr(r, "data", None) or getattr(r, "result", None)
    output_url = getattr(r, "output_url", None) or (getattr(r, "raw", None) or {}).get("output_url")
    if payload is not None:
        doc = json.loads(payload.model_dump_json())
    elif output_url:
        logger.info("ADE job {}: result over inline limit, downloading from output_url", job_id)
        with urllib.request.urlopen(output_url) as resp:
            doc = json.loads(resp.read())
    else:
        raise RuntimeError(f"ADE job {job_id} completed with neither inline payload nor output_url")

    meta = doc.get("metadata") or {}
    failed = meta.get("failed_pages") or []
    if failed:
        raise RuntimeError(f"ADE parse failed on pages {failed}; refusing to return a partial parse")
    # v1 metadata reports credit_usage and version; v2 reports
    # billing.total_credits and model_version.
    credits = meta.get("credit_usage")
    if credits is None:
        credits = (meta.get("billing") or {}).get("total_credits")
    logger.info(
        "ADE parse complete: pages={} credits={} version={}",
        meta.get("page_count"),
        credits,
        meta.get("version") or meta.get("model_version"),
    )
    return doc