"""FootnoteLookup — finds footnote definitions the deterministic tiers missed.

A table can carry a reference marker whose footnote text sits beyond the
correction agent's crop and defeats the reading-order text scan (a text
layer that dropped the line, a definition set in a layout the leading-token
parse cannot see). This agent is the last tier: it gets the unresolved
markers and the rendered page images following the table, and returns the
definitions it can actually see printed — each as the marker plus its text.

Demand-driven by design: callers invoke it only for markers the pure
resolution (`quber.core.fusion.footnotes`) left unresolved, so a table whose
footnotes sit under it costs no model call. A failed call returns None and
the markers stay unresolved in the exception record; the lookup must never
take a run down with it.
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Protocol, Sequence, runtime_checkable

from loguru import logger
from pydantic import BaseModel, Field

from quber.agents.factory import supports_sampling_temperature
from quber.agents.langsmith_tracer import LangSmithTracer, usage_metadata_from
from quber.agents.llm_client import FootnoteDef
from quber.settings import DEFAULT_LLM_MODEL, get_settings


class FootnoteLookupReport(BaseModel):
    """The definitions the agent found printed on the given pages."""

    found: List[FootnoteDef] = Field(
        default_factory=list,
        description=(
            "One entry per requested marker whose printed definition is visible on the "
            "pages: the marker as requested, and the definition text as printed. Omit "
            "any marker whose definition you cannot see — never invent one."
        ),
    )


LOOKUP_PROMPT = """\
You find the printed FOOTNOTE DEFINITIONS for specific reference markers.

A table carries the footnote reference markers listed below, but their
definitions were not printed directly under the table. You are given the
rendered images of the pages that follow it. Somewhere on them there may be
footnote lines — each opening with its own marker (a superscript or
parenthesized number, letter, or symbol) followed by its text.

For EVERY requested marker, look for a printed footnote line whose own
marker matches it. Report each one you find: the marker exactly as
requested, and the definition text exactly as printed (without the leading
marker). Omit markers whose definition you cannot see on these pages —
return them missing rather than guessed. Never compose, summarize, or
invent text; only transcribe what is printed.

What is NOT a footnote definition:
- A row of a table. If the marker's number appears inside a table row —
  'Commitments and contingencies (Note 16)' — that row is table content,
  not a footnote line. Do not transcribe it.
- A cross-reference to a named section of the document ('Note 16',
  '(Note 8)', 'Schedule II', 'Addendum 3'). A section is a separate part
  of the document, not a footnote printed for this table; a marker whose
  only appearance is such a reference is reported missing.
- A section heading, a caption, or body prose that merely contains the
  number.
"""


@runtime_checkable
class FootnoteLookup(Protocol):
    """Finds printed footnote definitions for unresolved markers."""

    async def lookup(
        self, markers: Sequence[str], page_pngs: Sequence[bytes], page: int
    ) -> Optional[FootnoteLookupReport]: ...


class PydanticAIFootnoteLookup:
    """FootnoteLookup backed by pydantic-ai with an Anthropic vision model.

    Mirrors `PydanticAIStatusInspector`: same auth resolution
    (`ANTHROPIC_AUTH_TOKEN` preferred over `ANTHROPIC_API_KEY`), greedy
    decoding so the same pages yield the same findings run to run.
    """

    DEFAULT_TEMPERATURE = 0.0

    def __init__(
        self,
        model: Optional[str] = None,
        auth_token: Optional[str] = None,
        api_key: Optional[str] = None,
    ) -> None:
        from pydantic_ai import Agent
        from pydantic_ai.models.anthropic import AnthropicModel
        from pydantic_ai.providers.anthropic import AnthropicProvider
        from pydantic_ai.settings import ModelSettings

        llm_settings = get_settings().llm
        model = model or llm_settings.model or DEFAULT_LLM_MODEL
        resolved_auth = auth_token or llm_settings.anthropic_auth_token
        resolved_key = api_key or llm_settings.anthropic_api_key

        if resolved_auth:
            from quber.agents._oauth_gate import make_oauth_anthropic_model

            anth_model = make_oauth_anthropic_model(model, resolved_auth)
        elif resolved_key:
            provider = AnthropicProvider(api_key=resolved_key)
            anth_model = AnthropicModel(model, provider=provider)
        else:
            raise RuntimeError(
                "PydanticAIFootnoteLookup: neither ANTHROPIC_AUTH_TOKEN nor "
                "ANTHROPIC_API_KEY is set. Provide one via env or constructor."
            )

        self.model = model
        self.agent = Agent(
            anth_model,
            output_type=FootnoteLookupReport,
            system_prompt=LOOKUP_PROMPT,
            model_settings=ModelSettings(temperature=self.DEFAULT_TEMPERATURE)
            if supports_sampling_temperature(model)
            else None,
        )
        self.tracer = LangSmithTracer(run_name="quber-footnote-lookup")

    def trace_inputs(self, user_text: str) -> Dict[str, Any]:
        return {
            "messages": [
                {"role": "system", "content": LOOKUP_PROMPT},
                {"role": "user", "content": [{"type": "text", "text": user_text}]},
            ]
        }

    async def lookup(
        self, markers: Sequence[str], page_pngs: Sequence[bytes], page: int
    ) -> Optional[FootnoteLookupReport]:
        from pydantic_ai import BinaryContent

        user_text = "Find the printed footnote definitions for these markers: " + ", ".join(
            repr(m) for m in markers
        )
        inputs = self.trace_inputs(user_text)
        try:
            async with self.tracer.llm_run("footnote_lookup", inputs, model=self.model) as run:
                content: List[Any] = [user_text]
                content.extend(BinaryContent(data=png, media_type="image/png") for png in page_pngs)
                result = await self.agent.run(content)
                report: FootnoteLookupReport = result.output
                run.outputs = {
                    "messages": [{"role": "assistant", "content": report.model_dump_json()}],
                    "usage_metadata": usage_metadata_from(result.usage),
                }
            return report
        except Exception as exc:
            logger.warning(
                "page {}: the footnote lookup call failed; the markers stay unresolved "
                "in the exception record: {}",
                page,
                exc,
            )
            return None


class MockFootnoteLookup:
    """Returns canned findings and records what it was asked. For tests."""

    def __init__(self, found: Optional[List[FootnoteDef]] = None) -> None:
        self.found = found or []
        self.calls: List[Dict[str, Any]] = []

    async def lookup(
        self, markers: Sequence[str], page_pngs: Sequence[bytes], page: int
    ) -> Optional[FootnoteLookupReport]:
        self.calls.append({"markers": list(markers), "pages": len(page_pngs), "page": page})
        return FootnoteLookupReport(found=self.found)