# Copyright (C) 2026 GoodData Corporation

"""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

RuleOrigin = Literal["embedded", "repo", "system"]


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

    query: str
    score_threshold: float
    rules_searched: int
    rules_returned: int
    low_confidence: bool = False
    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 = "Required: report selected rules concisely with provenance (id, origin, selection reasons)."
    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]


class RuleContentOmitted(TypedDict):
    kind: Literal["omitted"]
    text: str


RuleContent = RuleContentFull | RuleContentChunks | RuleContentOmitted


class CompactRulePayload(TypedDict):
    id: str
    origin: RuleOrigin
    editable: bool
    content: RuleContent


class ReportingRulePayload(TypedDict):
    id: str
    origin: RuleOrigin
    editable: bool
    reasons: list[str]


class RulesContext(TypedDict):
    low_confidence: bool
    rules: list[CompactRulePayload]


class RulesReporting(TypedDict):
    required: bool
    guidance: str
    selected: list[ReportingRulePayload]
    skipped_count: NotRequired[int]


class RulesDelta(TypedDict):
    cursor: str
    mode: NotRequired[Literal["full", "delta_empty"]]


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]
    origin: RuleOrigin
    editable: bool
    content: RuleContent


class RulesResponseCompactOk(TypedDict):
    ok: Literal[True]
    context: NotRequired[RulesContext]
    reporting: NotRequired[RulesReporting]
    delta: NotRequired[RulesDelta]
    rule_digests: NotRequired[list[RuleDigest]]
    message: NotRequired[str]


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


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


RulesResponseOk = RulesResponseCompactOk | RulesResponseDebugOk
RulesResponse = RulesResponseOk | RulesResponseErr


def _rule_content(sr: ScoredRule) -> RuleContent:
    is_large_rule = len(sr.rule.content) > LARGE_RULE_THRESHOLD
    has_chunks = bool(sr.matched_chunks)

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


def ordered_selection_reasons(sr: ScoredRule) -> list[str]:
    reasons_order = ("always", "score", "required", "injected")
    ordered = [name for name in reasons_order if name in sr.selection_reasons]
    if ordered:
        return ordered

    fallback: list[str] = []
    if sr.rule.always_apply:
        fallback.append("always")
    if sr.score > 0:
        fallback.append("score")
    if sr.injected:
        fallback.append("injected")
    return fallback


def _rule_payload(sr: ScoredRule) -> RulePayload:
    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": ordered_selection_reasons(sr),
            "required_by_rule_ids": sorted(set(sr.required_by_rule_ids)),
            "origin": cast(RuleOrigin, sr.rule.origin),
            "editable": sr.rule.editable,
            "content": _rule_content(sr),
        },
    )


def _compact_rule_payload(sr: ScoredRule) -> CompactRulePayload:
    return {
        "id": sr.rule.rule_id,
        "origin": cast(RuleOrigin, sr.rule.origin),
        "editable": sr.rule.editable,
        "content": _rule_content(sr),
    }


def _reporting_payload(sr: ScoredRule) -> ReportingRulePayload:
    return {
        "id": sr.rule.rule_id,
        "origin": cast(RuleOrigin, sr.rule.origin),
        "editable": sr.rule.editable,
        "reasons": ordered_selection_reasons(sr),
    }


def format_rules_response(
    scored_rules: list[ScoredRule],
    metadata: RetrievalMetadata | None = None,
    *,
    debug: bool = False,
) -> RulesResponse:
    """Format scored rules into a structured JSON-compatible response."""
    if not scored_rules:
        if debug:
            return cast(
                RulesResponseDebugOk,
                {
                    "ok": True,
                    "rules": [],
                    "rule_digests": [],
                    "metadata": None if metadata is None else metadata.__dict__,
                    "message": "No relevant rules found for this query.",
                },
            )
        return cast(
            RulesResponseCompactOk,
            {
                "ok": True,
                "context": {
                    "low_confidence": bool(metadata.low_confidence) if metadata is not None else False,
                    "rules": [],
                },
                "reporting": {
                    "required": bool(metadata.reporting_required) if metadata is not None else True,
                    "guidance": metadata.reminder if metadata is not None else "",
                    "selected": [],
                },
                "rule_digests": [],
                "message": "No relevant rules found for this query.",
            },
        )

    if debug:
        return cast(
            RulesResponseDebugOk,
            {
                "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],
            },
        )

    return cast(
        RulesResponseCompactOk,
        {
            "ok": True,
            "context": {
                "low_confidence": bool(metadata.low_confidence) if metadata is not None else False,
                "rules": [_compact_rule_payload(sr) for sr in scored_rules],
            },
            "reporting": {
                "required": bool(metadata.reporting_required) if metadata is not None else True,
                "guidance": metadata.reminder if metadata is not None else "",
                "selected": [_reporting_payload(sr) for sr in scored_rules],
            },
            "rule_digests": [{"rule_id": sr.rule.rule_id, "sha256": sr.rule.rule_hash} for sr in scored_rules],
        },
    )
