"""Typed contracts for LLM-owned nutrition coaching proposals."""

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import json
import re
from typing import Final, Iterable, Literal, Mapping, cast


Decision = Literal["maintain", "adjust", "observe", "safety_hold"]
Confidence = Literal["low", "medium", "high"]

COACH_V2_RESPONSE_KEYS: Final = frozenset(
    {
        "schema_version",
        "customer_key",
        "revision_binding_digest",
        "decision",
        "confidence",
        "evidence_ids",
        "interpretation",
        "recommendation_unit_system",
        "recommendation",
        "next_checkin_focus_ids",
        "customer_draft",
    }
)
_COACH_REVIEW_KEYS: Final = frozenset(
    {
        "schema_version",
        "current_targets",
        "proposed_targets",
        "decision",
        "confidence",
        "evidence_ids",
        "next_checkin_focus_ids",
        "interpretation",
        "warnings",
        "facts",
        "owner_review_notes",
        "revision_binding_digest",
    }
)
_COACH_ARTIFACT_KEYS: Final = frozenset(
    {
        "schema_version",
        "revision_binding_digest",
        "raw_coach_output",
        "raw_coach_sha256",
        "accepted_coach_output",
        "raw_polish_output",
        "accepted_polish_output",
        "polish_valid",
    }
)
_COACH_ARTIFACT_V3_KEYS: Final = (
    _COACH_ARTIFACT_KEYS | {"provider_generation_audits"}
)
_PROVIDER_GENERATION_AUDIT_KEYS: Final = frozenset(
    {
        "status",
        "finish_reason",
        "incomplete_reason",
        "input_tokens",
        "output_tokens",
        "total_tokens",
    }
)
_PROVIDER_GENERATION_REASON: Final = re.compile(r"[a-z][a-z0-9_.-]{0,63}")
_TARGET_KEYS: Final = frozenset({"calories", "protein_g", "carbs_g", "fat_g"})


def _exact_object(
    value: object,
    keys: frozenset[str],
) -> dict[str, object] | None:
    if type(value) is not dict or set(value) != keys:
        return None
    return cast(dict[str, object], value)


@dataclass(frozen=True, slots=True)
class NutritionTargets:
    calories: int
    protein_g: int
    carbs_g: int
    fat_g: int

    def as_dict(self) -> dict[str, int]:
        return {
            "calories": self.calories,
            "protein_g": self.protein_g,
            "carbs_g": self.carbs_g,
            "fat_g": self.fat_g,
        }


@dataclass(frozen=True, slots=True)
class ValidatedCoachProposal:
    decision: Decision
    confidence: Confidence
    evidence_ids: tuple[str, ...]
    interpretation: str
    recommendation: NutritionTargets
    next_checkin_focus_ids: tuple[str, ...]
    customer_draft: str
    warnings: tuple[str, ...]
    owner_review_notes: tuple[str, ...]


@dataclass(frozen=True, slots=True)
class CoachReview:
    schema_version: str
    current_targets: NutritionTargets
    proposed_targets: NutritionTargets
    decision: Decision
    confidence: Confidence
    evidence_ids: tuple[str, ...]
    next_checkin_focus_ids: tuple[str, ...]
    interpretation: str
    warnings: tuple[str, ...]
    facts: tuple[tuple[str, str], ...]
    owner_review_notes: tuple[str, ...]
    revision_binding_digest: str


def coach_review_to_dict(review: CoachReview) -> dict[str, object]:
    return {
        "schema_version": review.schema_version,
        "current_targets": review.current_targets.as_dict(),
        "proposed_targets": review.proposed_targets.as_dict(),
        "decision": review.decision,
        "confidence": review.confidence,
        "evidence_ids": list(review.evidence_ids),
        "next_checkin_focus_ids": list(review.next_checkin_focus_ids),
        "interpretation": review.interpretation,
        "warnings": list(review.warnings),
        "facts": [list(item) for item in review.facts],
        "owner_review_notes": list(review.owner_review_notes),
        "revision_binding_digest": review.revision_binding_digest,
    }


def coach_review_from_dict(value: object) -> CoachReview | None:
    record = _exact_object(value, _COACH_REVIEW_KEYS)
    if record is None:
        return None
    current = _review_targets(record.get("current_targets"))
    proposed = _review_targets(record.get("proposed_targets"))
    decision = record.get("decision")
    confidence = record.get("confidence")
    evidence_ids = _review_strings(record.get("evidence_ids"), 16)
    focus_ids = _review_strings(record.get("next_checkin_focus_ids"), 16)
    interpretation = record.get("interpretation")
    warnings = _review_strings(record.get("warnings"), 8)
    owner_review_notes = _review_strings(record.get("owner_review_notes"), 8)
    facts_value = record.get("facts")
    facts: list[tuple[str, str]] = []
    if type(facts_value) is not list or len(facts_value) > 32:
        return None
    for item in facts_value:
        if type(item) is not list or len(item) != 2:
            return None
        name, fact = item
        if (
            type(name) is not str
            or not name
            or len(name) > 500
            or type(fact) is not str
            or not fact
            or len(fact) > 500
        ):
            return None
        facts.append((name, fact))
    binding = record.get("revision_binding_digest")
    if (
        record.get("schema_version") != "nutrition-coach-review-v3"
        or current is None
        or proposed is None
        or type(decision) is not str
        or decision not in {"maintain", "adjust", "observe", "safety_hold"}
        or ((decision == "adjust") != (proposed != current))
        or type(confidence) is not str
        or confidence not in {"low", "medium", "high"}
        or not evidence_ids
        or not focus_ids
        or type(interpretation) is not str
        or not _bounded_text(interpretation, 1200)
        or warnings is None
        or owner_review_notes is None
        or any(
            term in line
            for line in owner_review_notes
            for term in ("kcal", "칼로리", "단백질", "탄수화물", "지방")
        )
        or type(binding) is not str
        or re.fullmatch(r"[0-9a-f]{64}", binding) is None
    ):
        return None
    return CoachReview(
        "nutrition-coach-review-v3",
        current,
        proposed,
        cast(Decision, decision),
        cast(Confidence, confidence),
        evidence_ids,
        focus_ids,
        interpretation,
        warnings,
        tuple(facts),
        owner_review_notes,
        binding,
    )


def coach_artifacts_from_dict(
    value: object,
    *,
    revision_binding_digest: str,
    accepted_polish_output: str,
) -> dict[str, object] | None:
    if type(value) is not dict:
        return None
    record = cast(dict[str, object], value)
    schema_version = record.get("schema_version")
    is_v2 = (
        schema_version == "nutrition-coach-artifacts-v2"
        and set(record) == _COACH_ARTIFACT_KEYS
    )
    is_v3 = (
        schema_version == "nutrition-coach-artifacts-v3"
        and set(record) == _COACH_ARTIFACT_V3_KEYS
    )
    if not is_v2 and not is_v3:
        return None
    raw_coach = record.get("raw_coach_output")
    raw_coach_sha256 = record.get("raw_coach_sha256")
    accepted_coach = record.get("accepted_coach_output")
    raw_polish = record.get("raw_polish_output")
    accepted_polish = record.get("accepted_polish_output")
    if (
        record.get("revision_binding_digest") != revision_binding_digest
        or type(raw_coach) is not str
        or not _bounded_text(raw_coach, 16_384)
        or type(raw_coach_sha256) is not str
        or raw_coach_sha256 != hashlib.sha256(raw_coach.encode()).hexdigest()
        or type(accepted_coach) is not str
        or not _bounded_text(accepted_coach, 8000)
        or (
            raw_polish is not None
            and (type(raw_polish) is not str or not _bounded_text(raw_polish, 16_384))
        )
        or type(accepted_polish) is not str
        or accepted_polish != accepted_polish_output
        or not _bounded_text(accepted_polish, 8000)
        or type(record.get("polish_valid")) is not bool
        or (
            is_v3
            and not _provider_generation_audits_valid(
                record.get("provider_generation_audits")
            )
        )
    ):
        return None
    return dict(record)


def _provider_generation_audits_valid(value: object) -> bool:
    if type(value) is not list or not 1 <= len(value) <= 2:
        return False
    for value_audit in value:
        audit = _exact_object(value_audit, _PROVIDER_GENERATION_AUDIT_KEYS)
        if audit is None:
            return False
        status = audit["status"]
        finish_reason = audit["finish_reason"]
        incomplete_reason = audit["incomplete_reason"]
        input_tokens = audit["input_tokens"]
        output_tokens = audit["output_tokens"]
        total_tokens = audit["total_tokens"]
        if (
            type(status) is not str
            or type(finish_reason) is not str
            or (incomplete_reason is not None and type(incomplete_reason) is not str)
            or type(input_tokens) is not int
            or type(output_tokens) is not int
            or type(total_tokens) is not int
            or input_tokens < 0
            or output_tokens < 0
            or total_tokens < input_tokens + output_tokens
        ):
            return False
        if status == "completed":
            if finish_reason != "stop" or incomplete_reason is not None:
                return False
        elif status == "incomplete":
            if (
                finish_reason not in {"length", "incomplete"}
                or type(incomplete_reason) is not str
                or _PROVIDER_GENERATION_REASON.fullmatch(incomplete_reason) is None
                or (
                    finish_reason == "length"
                    and incomplete_reason != "max_output_tokens"
                )
                or (
                    finish_reason == "incomplete"
                    and incomplete_reason == "max_output_tokens"
                )
            ):
                return False
        elif status == "failed":
            if finish_reason != "error" or incomplete_reason is not None:
                return False
        else:
            return False
    return True


def targets_from_plan(value: object) -> NutritionTargets | None:
    if not isinstance(value, Mapping):
        return None
    plan = cast(Mapping[str, object], value)
    calories = plan.get("calories_kcal", plan.get("calories"))
    protein = plan.get("protein_g")
    carbs = plan.get("carbohydrate_g", plan.get("carbs_g"))
    fat = plan.get("fat_g")
    if (
        type(calories) is not int
        or type(protein) is not int
        or type(carbs) is not int
        or type(fat) is not int
    ):
        return None
    return NutritionTargets(calories, protein, carbs, fat)


def valid_nutrition_targets(value: object) -> bool:
    targets = value if isinstance(value, NutritionTargets) else None
    return (
        targets is not None
        and min(targets.as_dict().values()) > 0
        and _energy_consistent(targets)
    )


def _request_id_enum(values: Iterable[str], label: str) -> tuple[str, ...]:
    ids = tuple(values)
    if (
        not ids
        or len(set(ids)) != len(ids)
        or any(
            type(value) is not str
            or value.strip() != value
            or not 1 <= len(value) <= 128
            for value in ids
        )
    ):
        raise ValueError(f"Coach request {label} are invalid")
    return ids


def coach_v2_response_schema(
    *,
    evidence_ids: Iterable[str] | None = None,
    focus_ids: Iterable[str] | None = None,
) -> dict[str, object]:
    """Return the generic or request-scoped strict Coach V2 JSON Schema."""
    if (evidence_ids is None) != (focus_ids is None):
        raise ValueError("Coach request ID constraints must be provided together")
    scoped_evidence = (
        None
        if evidence_ids is None
        else _request_id_enum(evidence_ids, "evidence IDs")
    )
    scoped_focus = (
        None if focus_ids is None else _request_id_enum(focus_ids, "focus IDs")
    )
    evidence_items: dict[str, object] = {"type": "string"}
    focus_items: dict[str, object] = {"type": "string"}
    interpretation: dict[str, object] = {"type": "string"}
    customer_draft: dict[str, object] = {"type": "string"}
    if scoped_evidence is not None and scoped_focus is not None:
        evidence_items["enum"] = list(scoped_evidence)
        focus_items["enum"] = list(scoped_focus)
        interpretation["pattern"] = "^[^0-9]*$"
        customer_draft["pattern"] = "^[^0-9]*$"
    return {
        "type": "object",
        "additionalProperties": False,
        "required": sorted(COACH_V2_RESPONSE_KEYS),
        "properties": {
            "schema_version": {
                "type": "string",
                "enum": ["nutrition-coach-response-v2"],
            },
            "customer_key": {"type": "string"},
            "revision_binding_digest": {
                "type": "string",
                "pattern": "^[a-f0-9]{64}$",
            },
            "decision": {
                "type": "string",
                "enum": ["maintain", "adjust", "observe", "safety_hold"],
            },
            "confidence": {
                "type": "string",
                "enum": ["low", "medium", "high"],
            },
            "evidence_ids": {
                "type": "array",
                "minItems": 1,
                "items": evidence_items,
            },
            "interpretation": interpretation,
            "recommendation_unit_system": {
                "type": "string",
                "enum": ["kcal_and_grams"],
            },
            "recommendation": {
                "type": "object",
                "additionalProperties": False,
                "required": ["calories", "protein_g", "carbs_g", "fat_g"],
                "properties": {
                    "calories": {"type": "integer", "minimum": 1},
                    "protein_g": {"type": "integer", "minimum": 1},
                    "carbs_g": {"type": "integer", "minimum": 1},
                    "fat_g": {"type": "integer", "minimum": 1},
                },
            },
            "next_checkin_focus_ids": {
                "type": "array",
                "minItems": 1,
                "items": focus_items,
            },
            "customer_draft": customer_draft,
        },
    }


def coach_v2_instruction_constraints(
    *,
    evidence_ids: Iterable[str],
    focus_ids: Iterable[str],
) -> str:
    """Render exact request constraints for every provider instruction."""
    scoped_evidence = _request_id_enum(evidence_ids, "evidence IDs")
    scoped_focus = _request_id_enum(focus_ids, "focus IDs")
    evidence_json = json.dumps(list(scoped_evidence), ensure_ascii=False, separators=(",", ":"))
    focus_json = json.dumps(list(scoped_focus), ensure_ascii=False, separators=(",", ":"))
    return (
        f"Exact allowed evidence_ids: {evidence_json}. "
        f"Exact allowed next_checkin_focus_ids: {focus_json}. "
        "Numeric grounding rule: interpretation and customer_draft MUST contain "
        "no ASCII numerals 0-9; put numeric calorie and macro targets only in "
        "the structured recommendation fields. Rebuild and validate every field "
        "against all response_schema constraints; never copy an unoffered ID."
    )


def coach_v2_response_schema_for_request(model_input: str) -> dict[str, object] | None:
    """Derive and authenticate the request-scoped schema at provider dispatch."""
    try:
        request = json.loads(model_input)
    except (TypeError, json.JSONDecodeError):
        return None
    if type(request) is not dict or request.get("schema_version") != "nutrition-coach-request-v2":
        return None
    evidence = request.get("evidence")
    observations = request.get("observations")
    if type(evidence) is not list or type(observations) is not list:
        return None
    try:
        evidence_values = tuple(
            cast(str, item["id"])
            for item in evidence
            if type(item) is dict and set(item) == {"id", "text"}
        )
        focus_values = tuple(
            cast(str, item["id"])
            for item in observations
            if type(item) is dict and set(item) == {"id", "text"}
        )
        if len(evidence_values) != len(evidence) or len(focus_values) != len(observations):
            return None
        expected = coach_v2_response_schema(
            evidence_ids=evidence_values,
            focus_ids=focus_values,
        )
    except (KeyError, TypeError, ValueError):
        return None
    return expected if request.get("response_schema") == expected else None


def _targets(value: object) -> NutritionTargets | None:
    record = _exact_object(value, _TARGET_KEYS)
    if record is None:
        return None
    calories = record["calories"]
    protein = record["protein_g"]
    carbs = record["carbs_g"]
    fat = record["fat_g"]
    if (
        type(calories) is not int
        or calories <= 0
        or type(protein) is not int
        or protein <= 0
        or type(carbs) is not int
        or carbs <= 0
        or type(fat) is not int
        or fat <= 0
    ):
        return None
    return NutritionTargets(calories, protein, carbs, fat)


def _bounded_text(value: object, maximum: int) -> bool:
    return type(value) is str and bool(value.strip()) and len(value) <= maximum


def _energy_consistent(targets: NutritionTargets) -> bool:
    calculated = 4 * targets.protein_g + 4 * targets.carbs_g + 9 * targets.fat_g
    return abs(calculated - targets.calories) <= 100


def _review_targets(value: object) -> NutritionTargets | None:
    if type(value) is not dict or set(value) != _TARGET_KEYS:
        return None
    targets = _targets(value)
    return targets if valid_nutrition_targets(targets) else None


def _review_strings(value: object, maximum: int) -> tuple[str, ...] | None:
    if type(value) is not list or len(value) > maximum:
        return None
    strings: list[str] = []
    for item in value:
        if type(item) is not str or not item or len(item) > 500:
            return None
        strings.append(item)
    return tuple(strings)