"""Response formatting for rule retrieval."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal, NotRequired, TypedDict, cast

from aida_mcp.tools.rules.embeddings.chunker import MIN_CHUNK_THRESHOLD
from aida_mcp.tools.rules.scoring import ScoredRule

# Rules larger than this threshold will show matched chunks instead of full content
LARGE_RULE_THRESHOLD = MIN_CHUNK_THRESHOLD


@dataclass
class RetrievalMetadata:
    """Metadata about the retrieval operation."""

    query: str
    score_threshold: float
    rules_searched: int
    rules_returned: int
    low_confidence: bool = False
    # Convenience fields for UIs/agents that want a compact summary without loading rule payloads.
    selected_rule_ids: list[str] | None = None
    selected_rule_labels: list[str] | None = None
    skipped_rule_ids: list[str] | None = None
    reporting_required: bool = True
    reminder: str = (
        "Recommended: show a compact summary of selected/skipped rule IDs "
        "(from rule_digests and skipped_rule_digests). Full rule details are optional."
    )
    # When False, retrieval used BM25-only (embedding model not loaded). Clients can suggest warmup_embeddings.
    semantic_available: bool = True
    bm25_only_hint: str | None = None


class RuleContentChunk(TypedDict):
    chunk_index: int
    total_chunks: int
    text: str


class RuleContentFull(TypedDict):
    kind: Literal["full"]
    text: str


class RuleContentChunks(TypedDict):
    kind: Literal["chunks"]
    chunks: list[RuleContentChunk]


RuleContent = RuleContentFull | RuleContentChunks


class RuleDigest(TypedDict):
    rule_id: str
    sha256: str


class RulePayload(TypedDict):
    path: str
    rule_id: str
    sha256: str
    description: str
    always_apply: bool
    required: list[str]
    related: list[str]
    injected: bool
    score: float
    bm25_score: float
    embedding_score: float
    matched_tokens: list[str]
    selection_reasons: list[str]
    required_by_rule_ids: list[str]
    content: RuleContent


class RulesResponseOk(TypedDict):
    ok: Literal[True]
    metadata: dict[str, object] | None
    rule_digests: list[RuleDigest]
    skipped_rule_digests: NotRequired[list[RuleDigest]]
    rules: list[RulePayload]
    message: NotRequired[str]


class RulesResponseErr(TypedDict):
    ok: Literal[False]
    error: str
    message: str


RulesResponse = RulesResponseOk | RulesResponseErr


def _rule_payload(sr: ScoredRule) -> RulePayload:
    """
    Build structured payload for a single rule.

    NOTE: We intentionally do NOT return megabytes of content for large rules. For large rules we return only
    the matched chunks (same behavior as the markdown formatter).
    """
    is_large_rule = len(sr.rule.content) > LARGE_RULE_THRESHOLD
    has_chunks = bool(sr.matched_chunks)

    if is_large_rule and has_chunks:
        content: RuleContent = {
            "kind": "chunks",
            "chunks": [
                {
                    "chunk_index": c.chunk_index,
                    "total_chunks": c.total_chunks,
                    "text": c.text,
                }
                for c in sr.matched_chunks
            ],
        }
    else:
        content = {"kind": "full", "text": sr.rule.content}

    return cast(
        RulePayload,
        {
            "path": sr.rule.path,
            "rule_id": sr.rule.rule_id,
            "sha256": sr.rule.rule_hash,
            "description": sr.rule.description,
            "always_apply": sr.rule.always_apply,
            "required": sr.rule.required,
            "related": sr.rule.related,
            "injected": sr.injected,
            "score": round(sr.score, 3),
            "bm25_score": round(sr.bm25_score, 3),
            "embedding_score": round(sr.embedding_score, 3),
            "matched_tokens": sorted(sr.matched_tokens),
            "selection_reasons": sorted(set(sr.selection_reasons)),
            "required_by_rule_ids": sorted(set(sr.required_by_rule_ids)),
            "content": content,
        },
    )


def format_rules_response(
    scored_rules: list[ScoredRule],
    metadata: RetrievalMetadata | None = None,
) -> RulesResponse:
    """
    Format scored rules into a structured JSON-compatible response.

    AIDA is experimental; this structure is optimized for LLMs and for observing Cursor's context shrinking.
    """
    if not scored_rules:
        return cast(
            RulesResponseOk,
            {
                "ok": True,
                "rules": [],
                "rule_digests": [],
                "metadata": None if metadata is None else metadata.__dict__,
                "message": "No relevant rules found for this query.",
            },
        )

    return cast(
        RulesResponseOk,
        {
            "ok": True,
            "metadata": None if metadata is None else metadata.__dict__,
            "rule_digests": [{"rule_id": sr.rule.rule_id, "sha256": sr.rule.rule_hash} for sr in scored_rules],
            "skipped_rule_digests": [],
            "rules": [_rule_payload(sr) for sr in scored_rules],
        },
    )
