"""
TableClassifier Protocol — decides whether a Camelot candidate is an
actual table.

Each candidate (its extracted markdown) gets a binary is_table
decision plus a short reason. In the Camelot LLM pipeline
(`quber.core.extractors.camelot.llm`), rejected candidates are dropped
before unification and correction, so the correction stage only sees
inputs that are actually tables — keeps stream-flavor noise (footnote
lists, narrative paragraphs, TOC entries, layout artifacts) out of the
correction step. In the dual flow (`quber.core.extractors.dual.camelot`)
the verdict is stored as the candidate's classification and nothing is
dropped.

Kept separate from `LLMClient` so the classifier can use a different
(cheaper/faster) model than the structure-correction agent, and so
retry semantics don't bleed across roles.
"""

from __future__ import annotations

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

from loguru import logger
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 ClassifierResult(BaseModel):
    is_table: bool = Field(
        description=(
            "True if the candidate is an actual data table (a structured grid with "
            "a header row identifying columns and at least one row of discrete data "
            "values). False for footnote lists, narrative paragraphs, Table-of-Contents "
            "entries, single-row layout artifacts, or any block that lacks a "
            "header/data structure."
        ),
    )
    reason: str = Field(
        default="",
        description=("One short sentence justifying the decision; what made it a table or not."),
    )


# Bias toward strict — false negatives (missing a borderline case) are
# preferred over false positives (shipping a non-table as a table).
# Schema is embedded so the model is the single source of truth for the
# response shape; same pattern as the unifier and detector prompts.
CLASSIFY_TABLE_PROMPT = f"""\
You are classifying a candidate text block extracted from a PDF page.

Decide whether the candidate is an *actual data table*:

- *Yes* if it has a header row identifying columns AND at least one row
  of discrete data values arranged in a rectangular grid.
- *No* otherwise. NOT a table: footnote lists, Table-of-Contents
  entries, narrative paragraphs, bullet lists formatted as cells,
  single-row layout artifacts (decorative banners, section dividers),
  callout boxes, or any block where the cells contain flowing prose
  rather than discrete values.

Bias toward strict: when the block is borderline, return False. Missing
a borderline real table is cheaper than polluting downstream with a
fake one.

Return a JSON object conforming exactly to this schema:

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

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


@runtime_checkable
class TableClassifier(Protocol):
    async def classify(self, markdown: str) -> ClassifierResult: ...


class PydanticAIClassifier:
    """TableClassifier backed by pydantic-ai with an Anthropic model.

    Auth resolution mirrors `PydanticAIClient`: `ANTHROPIC_AUTH_TOKEN`
    (OAuth subscription) preferred over `ANTHROPIC_API_KEY`. Model
    defaults to Haiku 4.5 — classification is a simple judgement and
    doesn't benefit from a larger model.
    """

    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(
                "PydanticAIClassifier: 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=ClassifierResult,
            system_prompt=CLASSIFY_TABLE_PROMPT,
        )

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

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

    async def classify(self, markdown: str) -> ClassifierResult:
        inputs = self.trace_inputs(markdown)
        try:
            async with self.tracer.llm_run("classify_table", inputs, model=self.model) as run:
                result = await self.agent.run(markdown)
                output = result.output
                run.outputs = {
                    "messages": [{"role": "assistant", "content": output.model_dump_json()}],
                    "is_table": output.is_table,
                    "reason": output.reason,
                    "usage_metadata": usage_metadata_from(result.usage),
                }
            return output
        except Exception as exc:
            logger.warning("classify: LLM call failed; default-accepting candidate: {}", exc)
            return ClassifierResult(is_table=True, reason="default-accept on classifier failure")


class MockClassifier:
    """Returns a canned ClassifierResult. For tests."""

    def __init__(self, result: Optional[ClassifierResult] = None) -> None:
        self.result = result or ClassifierResult(is_table=True, reason="mock")

    async def classify(self, markdown: str) -> ClassifierResult:
        _ = markdown
        return self.result


ClassifierBackend = Literal["api", "mock"]


def get_classifier(backend: Optional[ClassifierBackend] = None) -> TableClassifier:
    selected = backend or get_settings().llm.classifier_backend
    if selected == "api":
        return PydanticAIClassifier()
    if selected == "mock":
        return MockClassifier()
    raise ValueError(f"Unknown QUBER_CLASSIFIER_BACKEND: {selected!r}. Expected api|mock.")