"""Strict structured judgment boundary for customer nutrition coaching."""

from __future__ import annotations

import json
from typing import Final

from .nutrition_coaching_judgment_request import (
    build_judgment_request,
    generation_request_fingerprint,
    judgment_revision_binding,
)
from .nutrition_coaching_judgment_types import (
    Confidence,
    JsonValue,
    JudgmentOption,
    NutritionJudgmentGrounding,
    ValidatedNutritionJudgment,
)
from .nutrition_coaching_proposal import ValidatedCoachProposal
from .nutrition_coaching_proposal_validation import (
    CoachProposalValidationDiagnostic,
    diagnose_coach_proposal,
    validate_coach_proposal,
)

_RESPONSE_KEYS: Final = frozenset(
    {
        "schema_version",
        "customer_key",
        "revision_binding_digest",
        "primary_observation_ids",
        "evidence_ids",
        "action_id",
        "confidence",
        "limitation_ids",
    }
)


def validate_judgment(
    raw_response: str,
    grounding: NutritionJudgmentGrounding,
) -> ValidatedNutritionJudgment | ValidatedCoachProposal | None:
    try:
        version_probe = json.loads(raw_response)
    except (TypeError, json.JSONDecodeError):
        version_probe = None
    if (
        type(version_probe) is dict
        and version_probe.get("schema_version") == "nutrition-coach-response-v2"
    ):
        return validate_coach_proposal(
            raw_response,
            customer_key=grounding.customer_key,
            revision_binding_digest=grounding.revision_binding_digest,
            evidence_ids=(item.option_id for item in grounding.evidence),
            focus_ids=(item.option_id for item in grounding.observations),
            available_text=(
                item.text for item in (*grounding.observations, *grounding.evidence)
            ),
            current_targets=grounding.current_targets,
            valid_sample_count=grounding.valid_sample_count,
            plan_adherence=grounding.plan_adherence,
            average_sleep_hours=grounding.average_sleep_hours,
            deterministic_baseline=grounding.deterministic_baseline,
            safety_held=getattr(grounding, "safety_held", False) is True,
        )
    try:
        payload = json.loads(raw_response)
    except (TypeError, json.JSONDecodeError):
        return None
    if type(payload) is not dict or set(payload) != _RESPONSE_KEYS:
        return None
    if (
        payload.get("schema_version") != "nutrition-coaching-judgment-v1"
        or payload.get("customer_key") != grounding.customer_key
        or payload.get("revision_binding_digest")
        != grounding.revision_binding_digest
    ):
        return None
    observation_ids = _strict_ids(
        payload.get("primary_observation_ids"),
        1,
        3,
    )
    evidence_ids = _strict_ids(payload.get("evidence_ids"), 1, 2)
    limitation_ids = _strict_ids(
        payload.get("limitation_ids"),
        0,
        len(grounding.limitations),
    )
    action_id = payload.get("action_id")
    confidence = payload.get("confidence")
    if (
        observation_ids is None
        or evidence_ids is None
        or limitation_ids is None
        or type(action_id) is not str
        or confidence not in {"low", "medium", "high"}
    ):
        return None
    if not _offered(observation_ids, grounding.observations):
        return None
    if not _offered(evidence_ids, grounding.evidence):
        return None
    if not _offered((action_id,), grounding.actions):
        return None
    if not _offered(limitation_ids, grounding.limitations):
        return None
    return ValidatedNutritionJudgment(
        observation_ids,
        evidence_ids,
        action_id,
        confidence,
        limitation_ids,
    )


def diagnose_coach_judgment(
    raw_response: str,
    grounding: NutritionJudgmentGrounding,
) -> tuple[ValidatedCoachProposal | None, CoachProposalValidationDiagnostic]:
    """Validate the Coach-only contract without retaining raw model output."""
    return diagnose_coach_proposal(
        raw_response,
        customer_key=grounding.customer_key,
        revision_binding_digest=grounding.revision_binding_digest,
        evidence_ids=(item.option_id for item in grounding.evidence),
        focus_ids=(item.option_id for item in grounding.observations),
        available_text=(
            item.text for item in (*grounding.observations, *grounding.evidence)
        ),
        current_targets=grounding.current_targets,
        valid_sample_count=grounding.valid_sample_count,
        plan_adherence=grounding.plan_adherence,
        average_sleep_hours=grounding.average_sleep_hours,
        deterministic_baseline=grounding.deterministic_baseline,
        safety_held=getattr(grounding, "safety_held", False) is True,
    )


def render_validated_judgment(
    judgment: ValidatedNutritionJudgment,
    grounding: NutritionJudgmentGrounding,
) -> str:
    actions = {item.option_id: item.text for item in grounding.actions}
    limitations = {
        item.option_id: item.text
        for item in grounding.limitations
    }
    result = actions[judgment.action_id]
    if judgment.limitation_ids:
        result = (
            f"{result} "
            f"{' '.join(limitations[item] for item in judgment.limitation_ids)}"
        )
    return result


def _strict_ids(
    value: JsonValue,
    minimum: int,
    maximum: int,
) -> tuple[str, ...] | None:
    if type(value) is not list or not minimum <= len(value) <= maximum:
        return None
    if (
        any(type(item) is not str or not item for item in value)
        or len(set(value)) != len(value)
    ):
        return None
    return tuple(value)


def _offered(
    ids: tuple[str, ...],
    options: tuple[JudgmentOption, ...],
) -> bool:
    allowed = {item.option_id for item in options}
    return all(item in allowed for item in ids)


__all__ = [
    "Confidence",
    "JsonValue",
    "JudgmentOption",
    "NutritionJudgmentGrounding",
    "ValidatedNutritionJudgment",
    "build_judgment_request",
    "diagnose_coach_judgment",
    "generation_request_fingerprint",
    "judgment_revision_binding",
    "render_validated_judgment",
    "validate_judgment",
]