"""The two agents behind the figure-value reconciliation pass.

One page produces two independent readings of the values its figures print.
The scan's prose descriptions already state chart titles and labeled values;
a text-only agent lifts them into label/series/value entries. A vision agent
reads the page image alongside the parse's positioned fragments and reports
each printed value with the ids of the fragments whose text contains it, so
the deterministic tie in ``quber.core.figures.values`` can ground agreement
in a box on the page.

Both agents draw the same line: a value is something the figure states — a
bar label, a donut segment value or percentage, a donut center total, a
line-point label. Axis scale labels and gridline numbers are the figure's
furniture and are never values. And neither agent invents: the parser
outputs nothing the prose does not state, the reader nothing the image does
not show, and a printed value no fragment contains is reported with empty
fragment ids rather than a fabricated citation.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional, Sequence

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.core.figures.values import LocalValueReader, PageFragment, ParsedValue, ReadValue, ScanValueParser
from quber.settings import DEFAULT_LLM_MODEL, get_settings


class ParsedValues(BaseModel):
    values: List[ParsedValue] = Field(
        default_factory=list, description="Every labeled value the prose states; empty if none"
    )


class ReadValues(BaseModel):
    values: List[ReadValue] = Field(
        default_factory=list, description="Every value printed inside a figure on the page; empty if none"
    )


SCAN_VALUE_PROMPT = f"""\
You are given a page scan's prose descriptions of the figures on one page of
a financial document. The descriptions state chart titles and the labeled
values each chart prints.

Extract EVERY labeled numeric value the prose states: bar labels, donut
segment values and their percentages, donut center totals, line-point
labels. Report each as a label/series/value entry. A donut segment stated
with both a value and a percentage is two entries, one per stated form.

The value is copied EXACTLY as the prose states it — keep the printed form,
including any '$', '%', commas, and parentheses.

Axis scale labels and gridline numbers are furniture, not values. When the
prose mentions an axis range or the numbers running along an axis, output
nothing for them.

Output nothing that is not stated in the prose.

Return a JSON object conforming exactly to this schema:

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

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


FIGURE_VALUE_READER_PROMPT = f"""\
You are shown one page of a financial document as an image, plus a listing
of the positioned text fragments a parser extracted from that page — each
with an id, its text, and its box in the page's normalized frame.

Read EVERY numeric value printed inside a chart or figure on the page image:
bar labels, donut segment values and percentages, donut center totals,
line-point labels. For each, report:

- chart_title: the title of the chart it is printed in
- label: what the value is labeled as on the chart
- series: the series it belongs to, when the chart plots more than one
- value: EXACTLY as printed, including any '$', '%', commas, and parentheses
- fragment_ids: the id(s) of the fragment(s) whose text contains the value,
  plus the fragment carrying its label when the label is printed separately

Cite only fragments whose text actually contains the value or its label.
When a printed value has no fragment containing it, report the value with
EMPTY fragment_ids — never cite a fragment that does not print it.

Axis scale labels and gridline numbers are furniture — skip them.

Report nothing that is not visible in the image.

Return a JSON object conforming exactly to this schema:

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

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


def figure_text_block(figure_texts: Sequence[str]) -> str:
    """The scan's descriptions, numbered the way one page presents its figures."""
    return "\n\n".join(f"FIGURE {i}:\n{text}" for i, text in enumerate(figure_texts))


def fragment_listing(fragments: Sequence[PageFragment]) -> str:
    """The page's fragments, one per line: id, text, box rounded to 3 decimals."""
    lines = []
    for frag in fragments:
        b = frag.box
        lines.append(
            f"{frag.id}: {frag.text} "
            f"box=({b['left']:.3f}, {b['top']:.3f}, {b['right']:.3f}, {b['bottom']:.3f})"
        )
    return "\n".join(lines)


class PydanticAIScanValueParser:
    """ScanValueParser backed by pydantic-ai with an Anthropic model.

    Auth resolution mirrors the other agents: `ANTHROPIC_AUTH_TOKEN` (OAuth
    subscription) preferred over `ANTHROPIC_API_KEY`.
    """

    # The same prose states the same values between runs. Greedy decoding
    # removes the sampling variance.
    DEFAULT_TEMPERATURE = 0.0

    def __init__(
        self,
        model: Optional[str] = None,
        auth_token: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: Optional[float] = DEFAULT_TEMPERATURE,
    ) -> 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
        if temperature is not None and not supports_sampling_temperature(model):
            temperature = None
        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(
                "PydanticAIScanValueParser: neither ANTHROPIC_AUTH_TOKEN nor "
                "ANTHROPIC_API_KEY is set. Provide one via env or constructor."
            )

        self.model = model
        self.temperature = temperature
        model_settings = None if temperature is None else ModelSettings(temperature=temperature)
        self.agent = Agent(
            anth_model,
            output_type=ParsedValues,
            system_prompt=SCAN_VALUE_PROMPT,
            model_settings=model_settings,
        )

        # LangSmith tracer is a no-op if `TRACE_TO_LANGSMITH` is unset.
        self.tracer = LangSmithTracer(run_name="quber-figure-scan-values")

    def trace_inputs(self, figure_texts: Sequence[str], page: int) -> Dict[str, Any]:
        return {
            "messages": [
                {"role": "system", "content": SCAN_VALUE_PROMPT},
                {
                    "role": "user",
                    "content": f"page {page}\n\n{figure_text_block(figure_texts)}",
                },
            ]
        }

    async def parse(self, figure_texts: Sequence[str], page: int) -> List[ParsedValue]:
        described = f"Figure descriptions read on page {page}:\n\n{figure_text_block(figure_texts)}"
        async with self.tracer.llm_run(
            "figure_scan_values", self.trace_inputs(figure_texts, page), model=self.model
        ) as run:
            result = await self.agent.run(described)
            output = result.output
            run.outputs = {
                "messages": [{"role": "assistant", "content": output.model_dump_json()}],
                "values": len(output.values),
                "usage_metadata": usage_metadata_from(result.usage),
            }
        return list(output.values)


class PydanticAIFigureValueReader:
    """LocalValueReader backed by pydantic-ai with an Anthropic vision model.

    Auth resolution mirrors the other agents: `ANTHROPIC_AUTH_TOKEN` (OAuth
    subscription) preferred over `ANTHROPIC_API_KEY`.
    """

    # The same page prints the same values between runs. Greedy decoding
    # removes the sampling variance.
    DEFAULT_TEMPERATURE = 0.0

    def __init__(
        self,
        model: Optional[str] = None,
        auth_token: Optional[str] = None,
        api_key: Optional[str] = None,
        temperature: Optional[float] = DEFAULT_TEMPERATURE,
    ) -> 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.figure_values_model or llm_settings.model or DEFAULT_LLM_MODEL
        if temperature is not None and not supports_sampling_temperature(model):
            temperature = None
        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(
                "PydanticAIFigureValueReader: neither ANTHROPIC_AUTH_TOKEN nor "
                "ANTHROPIC_API_KEY is set. Provide one via env or constructor."
            )

        self.model = model
        self.temperature = temperature
        model_settings = None if temperature is None else ModelSettings(temperature=temperature)
        self.agent = Agent(
            anth_model,
            output_type=ReadValues,
            system_prompt=FIGURE_VALUE_READER_PROMPT,
            model_settings=model_settings,
        )

        # LangSmith tracer is a no-op if `TRACE_TO_LANGSMITH` is unset.
        self.tracer = LangSmithTracer(run_name="quber-figure-value-reader")

    def trace_inputs(self, fragments: Sequence[PageFragment], page: int) -> Dict[str, Any]:
        return {
            "messages": [
                {"role": "system", "content": FIGURE_VALUE_READER_PROMPT},
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": f"page {page}\n{fragment_listing(fragments)}"},
                        {"type": "image", "text": "[page image attached]"},
                    ],
                },
            ]
        }

    async def read(self, image_path: Path, fragments: Sequence[PageFragment], page: int) -> List[ReadValue]:
        from pydantic_ai import BinaryContent

        image = BinaryContent(data=image_path.read_bytes(), media_type="image/png")
        listed = f"Fragments extracted from page {page}:\n{fragment_listing(fragments)}"
        async with self.tracer.llm_run(
            "figure_value_read", self.trace_inputs(fragments, page), model=self.model
        ) as run:
            result = await self.agent.run([listed, image])
            output = result.output
            run.outputs = {
                "messages": [{"role": "assistant", "content": output.model_dump_json()}],
                "values": len(output.values),
                "usage_metadata": usage_metadata_from(result.usage),
            }
        return list(output.values)


FigureValuesBackend = Literal["api", "off"]


def get_scan_value_parser(backend: Optional[FigureValuesBackend] = None) -> Optional[ScanValueParser]:
    """The scan-value parser for this run, or None when the pass is switched off."""
    selected = backend or get_settings().llm.figure_values_backend
    if selected == "api":
        return PydanticAIScanValueParser()
    if selected == "off":
        return None
    raise ValueError(f"Unknown QUBER_FIGURE_VALUES_BACKEND: {selected!r}. Expected api|off.")


def get_figure_value_reader(backend: Optional[FigureValuesBackend] = None) -> Optional[LocalValueReader]:
    """The local figure-value reader for this run, or None when the pass is switched off."""
    selected = backend or get_settings().llm.figure_values_backend
    if selected == "api":
        return PydanticAIFigureValueReader()
    if selected == "off":
        return None
    raise ValueError(f"Unknown QUBER_FIGURE_VALUES_BACKEND: {selected!r}. Expected api|off.")