"""Focused adjudication prompt building for Intelligence Pipeline v3.

Builds adjudication packets containing only relevant chunks and candidates,
uses strict JSON Schema with temperature zero, and enforces a bounded output
budget (max 1536 tokens for decisions, not summaries).
"""

from __future__ import annotations

from typing import Any

from pydantic import BaseModel, Field

from services.intelligence_pipeline_v3.adjudication.schemas import (
    AdjudicationCandidate,
    AdjudicationQuestion,
    ConflictDescription,
    EvidencePacket,
    QuestionCode,
)
from services.intelligence_pipeline_v3.segmenter.models import DocumentChunk

# --- Constants ---

MAX_OUTPUT_TOKENS: int = 1536
"""Maximum tokens for adjudication decisions output. Bounded to prevent
long summaries — the adjudicator produces decisions, not narratives."""

TEMPERATURE: float = 0.0
"""Temperature for adjudication requests. Zero for deterministic output."""

PROMPT_SCHEMA_VERSION: str = "1.0.0"
"""Version of the adjudication prompt schema format."""

PROVIDER_LINEAGE_KEY: str = "adjudication_v3"
"""Lineage identifier for adjudication prompts."""


# --- Models ---


class PromptMetadata(BaseModel):
    """Metadata for the adjudication prompt including version and lineage.

    Tracks prompt version, schema version, and provider lineage for
    reproducibility and auditing.
    """

    prompt_version: str = Field(
        default="1.0.0",
        description="Version of the prompt template",
    )
    schema_version: str = Field(
        default=PROMPT_SCHEMA_VERSION,
        description="Version of the JSON Schema format used",
    )
    provider_lineage: str = Field(
        default=PROVIDER_LINEAGE_KEY,
        description="Identifier for the prompt provider/pipeline stage",
    )
    max_output_tokens: int = Field(
        default=MAX_OUTPUT_TOKENS,
        description="Maximum output token budget for this prompt",
    )
    temperature: float = Field(
        default=TEMPERATURE,
        description="Generation temperature",
    )


class AdjudicationPacket(BaseModel):
    """Complete packet sent to the 9B adjudicator.

    Contains only the information relevant to resolving the specific
    ambiguity — relevant chunks, candidates, conflicts, and questions.
    """

    document_id: str = Field(description="Source document identifier")
    document_type: str = Field(description="Type of document")
    relevant_chunks: list[DocumentChunk] = Field(
        description="Only chunks relevant to the adjudication questions",
    )
    candidates: list[AdjudicationCandidate] = Field(
        description="Candidates requiring adjudication",
    )
    conflicts: list[ConflictDescription] = Field(
        default_factory=list,
        description="Conflicts between candidates",
    )
    questions: list[AdjudicationQuestion] = Field(
        description="Specific questions the adjudicator must answer",
    )
    evidence: list[EvidencePacket] = Field(
        description="Evidence spans available for reference",
    )
    metadata: PromptMetadata = Field(
        default_factory=PromptMetadata,
        description="Prompt metadata for versioning and lineage",
    )


# --- Output schema for strict JSON mode ---


def get_decision_json_schema() -> dict[str, Any]:
    """Return the strict JSON Schema for adjudication decisions.

    Used as the `response_format.json_schema.schema` payload when
    calling the 9B model with strict structured output.
    """
    return {
        "type": "object",
        "properties": {
            "decisions": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "decision_id": {"type": "string"},
                        "question_code": {
                            "type": "string",
                            "enum": [code.value for code in QuestionCode],
                        },
                        "verdict": {
                            "type": "string",
                            "enum": [
                                "accept",
                                "reject",
                                "merge",
                                "split",
                                "reattribute",
                            ],
                        },
                        "candidate_ids": {
                            "type": "array",
                            "items": {"type": "string"},
                            "minItems": 1,
                        },
                        "evidence_ids": {
                            "type": "array",
                            "items": {"type": "string"},
                            "minItems": 1,
                        },
                        "reasoning": {"type": "string"},
                        "resolved_value": {"type": "object"},
                    },
                    "required": [
                        "decision_id",
                        "question_code",
                        "verdict",
                        "candidate_ids",
                        "evidence_ids",
                        "reasoning",
                    ],
                    "additionalProperties": False,
                },
            },
        },
        "required": ["decisions"],
        "additionalProperties": False,
    }


# --- Packet builder ---


def _get_relevant_chunk_ids(
    candidates: list[AdjudicationCandidate],
    conflicts: list[ConflictDescription],
    questions: list[AdjudicationQuestion],
) -> set[str]:
    """Collect chunk IDs referenced by candidates, conflicts, and questions."""
    chunk_ids: set[str] = set()
    for candidate in candidates:
        chunk_ids.update(candidate.source_chunk_ids)
    return chunk_ids


def _filter_relevant_chunks(
    document_chunks: list[DocumentChunk],
    relevant_chunk_ids: set[str],
) -> list[DocumentChunk]:
    """Filter document chunks to include only those referenced by candidates."""
    if not relevant_chunk_ids:
        # If no specific chunks referenced, include all (fallback for
        # cases where chunk IDs weren't specified in candidates)
        return document_chunks
    return [c for c in document_chunks if c.chunk_id in relevant_chunk_ids]


def _collect_evidence_ids(
    candidates: list[AdjudicationCandidate],
    conflicts: list[ConflictDescription],
) -> set[str]:
    """Collect all evidence IDs referenced by candidates and conflicts."""
    evidence_ids: set[str] = set()
    for candidate in candidates:
        evidence_ids.update(candidate.evidence_ids)
    for conflict in conflicts:
        evidence_ids.update(conflict.evidence_ids)
    return evidence_ids


def build_adjudication_packet(
    document_id: str,
    document_type: str,
    document_chunks: list[DocumentChunk],
    candidates: list[AdjudicationCandidate],
    conflicts: list[ConflictDescription],
    questions: list[AdjudicationQuestion],
    evidence: list[EvidencePacket],
    *,
    question_codes: list[str] | None = None,
) -> AdjudicationPacket:
    """Build an adjudication packet with only relevant chunks and evidence.

    Filters document_chunks to include only those referenced by the
    candidates being adjudicated. Ensures the packet is focused and
    within the bounded context the adjudicator expects.

    Args:
        document_id: Source document identifier.
        document_type: Type of document (article, filing, transcript, etc.).
        document_chunks: All available chunks for the document.
        candidates: Candidates requiring adjudication.
        conflicts: Conflicts between candidates.
        questions: Specific questions to resolve.
        evidence: Available evidence spans.
        question_codes: Optional filter to limit questions by code.

    Returns:
        AdjudicationPacket with only relevant chunks included.
    """
    # Filter questions by code if specified
    filtered_questions = questions
    if question_codes:
        code_set = set(question_codes)
        filtered_questions = [
            q for q in questions if q.question_code.value in code_set
        ]

    # Determine which chunks are relevant
    relevant_chunk_ids = _get_relevant_chunk_ids(
        candidates, conflicts, filtered_questions
    )
    relevant_chunks = _filter_relevant_chunks(document_chunks, relevant_chunk_ids)

    # Filter evidence to only include those referenced by candidates/conflicts
    referenced_evidence_ids = _collect_evidence_ids(candidates, conflicts)
    if referenced_evidence_ids:
        relevant_evidence = [
            e for e in evidence if e.evidence_id in referenced_evidence_ids
        ]
    else:
        # Include all evidence if none specifically referenced
        relevant_evidence = evidence

    return AdjudicationPacket(
        document_id=document_id,
        document_type=document_type,
        relevant_chunks=relevant_chunks,
        candidates=candidates,
        conflicts=conflicts,
        questions=filtered_questions,
        evidence=relevant_evidence,
        metadata=PromptMetadata(),
    )


def build_request_payload(packet: AdjudicationPacket) -> dict[str, Any]:
    """Build the full inference request payload for the adjudicator.

    Returns a dict suitable for passing to the inference gateway,
    including strict JSON Schema response format and temperature zero.
    """
    system_prompt = (
        "You are a semantic adjudicator for financial document extraction. "
        "Resolve the ambiguities described in the questions using ONLY the "
        "provided evidence spans. Every decision MUST reference evidence_ids "
        "from the provided evidence. Do NOT estimate confidence, novelty, "
        "impact magnitude, or time horizon — those are computed by separate "
        "calibrated pipelines. Output valid JSON matching the required schema."
    )

    user_content = packet.model_dump_json()

    return {
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content},
        ],
        "temperature": packet.metadata.temperature,
        "max_tokens": packet.metadata.max_output_tokens,
        "response_format": {
            "type": "json_schema",
            "json_schema": {
                "name": "adjudication_response",
                "strict": True,
                "schema": get_decision_json_schema(),
            },
        },
    }