"""
CoverAttributor Protocol — decides scan breadth for a document nobody typed.

One narrow question, answered from the cover page and the page count: is this
a presentation-style document or a prose-and-table filing? The answer selects
how many pages the figure scan reads — a presentation scans every page, a
filing scans only the pages the parse nominates. Nothing else hangs off it:
the attributor never writes the document's filing type, which is identity
metadata the user owns, and it never runs when a filing type was supplied.

The caller owns the failure policy. An attribution that raises must fall back
to the cheap path and say so where the run is recorded, because the point of
this agent is that a missing type can never buy the expensive path by
accident.
"""

from __future__ import annotations

import json
from typing import Any, Dict, Literal, Optional, Protocol, runtime_checkable

from pydantic import BaseModel, Field

from quber.agents.langsmith_tracer import LangSmithTracer, usage_metadata_from
from quber.settings import DEFAULT_LLM_MODEL, get_settings


class CoverAttribution(BaseModel):
    kind: Literal["presentation", "filing"] = Field(
        description=(
            "'presentation' for a slide-deck-style document: an earnings deck, "
            "supplemental package, or investor presentation, where content is "
            "laid out visually page by page. 'filing' for a prose-and-table "
            "document: a 10-Q, 10-K, prospectus, or agreement, read as "
            "continuous text."
        ),
    )
    reason: str = Field(
        default="",
        description="One short sentence naming what on the cover decided it.",
    )


ATTRIBUTE_COVER_PROMPT = f"""\
You are shown the first page of a financial document and told its page count.

Decide which of two kinds the document is:

- *presentation*: a slide-deck-style document — an earnings presentation,
  quarterly supplemental package, or investor deck. Typical covers carry a
  company name and quarter as a designed title page, often landscape, often
  image-heavy, with little continuous prose.
- *filing*: a prose-and-table document — a Form 10-Q or 10-K, prospectus, or
  agreement. Typical covers carry an SEC form header, registrant details, and
  dense typeset text, and the document runs long.

Answer from what the cover shows. The page count is context, not the answer:
long presentations and short filings both exist.

Return a JSON object conforming exactly to this schema:

{json.dumps(CoverAttribution.model_json_schema(), indent=2)}

Return only the JSON object — no prose, no markdown code fences.
"""


@runtime_checkable
class CoverAttributor(Protocol):
    async def attribute(self, cover_png: bytes, page_count: int) -> CoverAttribution: ...


class PydanticAICoverAttributor:
    """CoverAttributor backed by pydantic-ai with an Anthropic model.

    Auth resolution mirrors the other agents: `ANTHROPIC_AUTH_TOKEN` preferred
    over `ANTHROPIC_API_KEY`. A failure raises to the caller — this agent's
    contract is that the caller falls back to the cheap path and records why,
    so a silent default here would hide exactly what must be surfaced.
    """

    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

        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(
                "PydanticAICoverAttributor: 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=CoverAttribution,
            system_prompt=ATTRIBUTE_COVER_PROMPT,
        )
        self.tracer = LangSmithTracer(run_name="quber-cover-attribution")

    def trace_inputs(self, page_count: int) -> Dict[str, Any]:
        return {
            "messages": [
                {"role": "system", "content": ATTRIBUTE_COVER_PROMPT},
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": f"The document is {page_count} pages long."},
                        {"type": "text", "text": "<cover page image>"},
                    ],
                },
            ]
        }

    async def attribute(self, cover_png: bytes, page_count: int) -> CoverAttribution:
        from pydantic_ai import BinaryContent

        image = BinaryContent(data=cover_png, media_type="image/png")
        inputs = self.trace_inputs(page_count)
        async with self.tracer.llm_run("attribute_cover", inputs, model=self.model) as run:
            result = await self.agent.run([f"The document is {page_count} pages long.", image])
            output = result.output
            run.outputs = {
                "messages": [{"role": "assistant", "content": output.model_dump_json()}],
                "kind": output.kind,
                "reason": output.reason,
                "usage_metadata": usage_metadata_from(result.usage),
            }
        return output


class MockCoverAttributor:
    """Returns a canned CoverAttribution. For tests."""

    def __init__(self, result: Optional[CoverAttribution] = None) -> None:
        self.result = result or CoverAttribution(kind="presentation", reason="mock")

    async def attribute(self, cover_png: bytes, page_count: int) -> CoverAttribution:
        _ = cover_png, page_count
        return self.result


def get_cover_attributor() -> CoverAttributor:
    return PydanticAICoverAttributor()