from __future__ import annotations

from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Any
import json

try:
    import yaml
except Exception:  # pragma: no cover
    yaml = None


DEFAULT_POLICY_RELATIVE_PATH = Path("configs/sponsor_model_policy.yaml")
DEFAULT_WAIVER_RELATIVE_PATH = Path("configs/sponsor_model_waiver.yaml")


@dataclass(frozen=True)
class SponsorRequirement:
    scope: str
    key: str
    component: str
    model_id: str
    subkey: str | None = None


@dataclass(frozen=True)
class SponsorMismatch:
    scope: str
    key: str
    expected_component: str
    expected_model_id: str
    actual_component: str | None
    actual_model_id: str | None
    problem: str
    subkey: str | None = None


@dataclass(frozen=True)
class SponsorWaiver:
    path: Path
    reason: str
    date: str
    approved_by: str


@dataclass(frozen=True)
class SponsorPolicyCheckResult:
    ok: bool
    requirements: tuple[SponsorRequirement, ...]
    mismatches: tuple[SponsorMismatch, ...]
    waiver: SponsorWaiver | None


def _load_data(path: Path) -> Any:
    text = path.read_text(encoding="utf-8")
    if path.suffix.lower() == ".json":
        return json.loads(text)
    if yaml is not None:
        return yaml.safe_load(text)
    return json.loads(text)


def load_policy(path: str | Path) -> dict[str, Any]:
    path = Path(path)
    data = _load_data(path)
    if not isinstance(data, dict):
        raise ValueError(f"sponsor policy must be a mapping, got {type(data)!r}")
    return data


def load_registry(path: str | Path) -> dict[str, Any]:
    path = Path(path)
    data = _load_data(path)
    if not isinstance(data, dict):
        raise ValueError(f"model registry must be a mapping, got {type(data)!r}")
    return data


def _require_mapping(data: Any, *, label: str) -> dict[str, Any]:
    if not isinstance(data, dict):
        raise ValueError(f"{label} must be a mapping, got {type(data)!r}")
    return data


def _iter_requirements(policy: dict[str, Any]) -> tuple[SponsorRequirement, ...]:
    required = policy.get("required_models", policy)
    required = _require_mapping(required, label="sponsor policy.required_models")
    requirements: list[SponsorRequirement] = []

    shared = required.get("shared", {})
    shared = _require_mapping(shared, label="sponsor policy.required_models.shared")
    for key, spec in shared.items():
        spec = _require_mapping(spec, label=f"sponsor policy.required_models.shared.{key}")
        component = str(spec.get("component", key))
        model_id = spec.get("model_id")
        if not isinstance(model_id, str) or not model_id.strip():
            raise ValueError(f"sponsor policy.required_models.shared.{key}.model_id must be a non-empty string")
        requirements.append(SponsorRequirement(scope="shared", key=key, component=component, model_id=model_id))

    projects = required.get("projects", {})
    projects = _require_mapping(projects, label="sponsor policy.required_models.projects")
    for project_key, spec in projects.items():
        spec = _require_mapping(spec, label=f"sponsor policy.required_models.projects.{project_key}")
        if "model_id" in spec:
            component = str(spec.get("component", ""))
            model_id = spec.get("model_id")
            if not component:
                raise ValueError(f"sponsor policy.required_models.projects.{project_key}.component must be a non-empty string")
            if not isinstance(model_id, str) or not model_id.strip():
                raise ValueError(f"sponsor policy.required_models.projects.{project_key}.model_id must be a non-empty string")
            requirements.append(
                SponsorRequirement(scope="projects", key=project_key, component=component, model_id=model_id)
            )
            continue

        for component_key, component_spec in spec.items():
            component_spec = _require_mapping(
                component_spec,
                label=f"sponsor policy.required_models.projects.{project_key}.{component_key}",
            )
            component = str(component_spec.get("component", component_key))
            model_id = component_spec.get("model_id")
            if not component:
                raise ValueError(
                    f"sponsor policy.required_models.projects.{project_key}.{component_key}.component must be a non-empty string"
                )
            if not isinstance(model_id, str) or not model_id.strip():
                raise ValueError(
                    f"sponsor policy.required_models.projects.{project_key}.{component_key}.model_id must be a non-empty string"
                )
            requirements.append(
                SponsorRequirement(
                    scope="projects",
                    key=project_key,
                    component=component,
                    model_id=model_id,
                    subkey=component_key,
                )
            )

    return tuple(requirements)


def _validate_waiver(path: Path, waiver_data: Any) -> SponsorWaiver:
    data = _require_mapping(waiver_data, label="sponsor waiver")
    reason = data.get("reason")
    approved_by = data.get("approved_by")
    waiver_date = data.get("date")
    for field_name, field_value in (("reason", reason), ("approved_by", approved_by), ("date", waiver_date)):
        if not isinstance(field_value, str) or not field_value.strip():
            raise ValueError(f"sponsor waiver.{field_name} must be a non-empty string")
    try:
        date.fromisoformat(waiver_date)
    except ValueError as exc:
        raise ValueError("sponsor waiver.date must use ISO format YYYY-MM-DD") from exc
    return SponsorWaiver(path=path, reason=reason.strip(), date=waiver_date.strip(), approved_by=approved_by.strip())


def load_waiver(path: str | Path | None) -> SponsorWaiver | None:
    if path is None:
        return None
    path = Path(path)
    if not path.exists():
        return None
    return _validate_waiver(path, _load_data(path))


def _registry_entry_for_requirement(registry: dict[str, Any], requirement: SponsorRequirement) -> dict[str, Any] | None:
    if requirement.scope == "shared":
        shared = registry.get("shared", {})
        if not isinstance(shared, dict):
            raise ValueError("model registry.shared must be a mapping")
        entry = shared.get(requirement.key)
        return entry if isinstance(entry, dict) else None
    if requirement.scope == "projects":
        projects = registry.get("projects", {})
        if not isinstance(projects, dict):
            raise ValueError("model registry.projects must be a mapping")
        project_entry = projects.get(requirement.key)
        if not isinstance(project_entry, dict):
            return None
        if requirement.subkey is None:
            return project_entry
        entry = project_entry.get(requirement.subkey)
        return entry if isinstance(entry, dict) else None
    raise ValueError(f"unsupported sponsor requirement scope: {requirement.scope!r}")


def check_sponsor_policy(
    model_registry_path: str | Path,
    policy_path: str | Path,
    waiver_path: str | Path | None = None,
) -> SponsorPolicyCheckResult:
    registry = load_registry(model_registry_path)
    policy = load_policy(policy_path)
    requirements = _iter_requirements(policy)
    waiver = load_waiver(waiver_path)

    mismatches: list[SponsorMismatch] = []
    for requirement in requirements:
        entry = _registry_entry_for_requirement(registry, requirement)
        if entry is None:
            mismatches.append(
                SponsorMismatch(
                    scope=requirement.scope,
                    key=requirement.key,
                    expected_component=requirement.component,
                    expected_model_id=requirement.model_id,
                    actual_component=None,
                    actual_model_id=None,
                    problem="missing registry entry",
                    subkey=requirement.subkey,
                )
            )
            continue
        actual_component = entry.get("component")
        actual_model_id = entry.get("model_id")
        if actual_component != requirement.component or actual_model_id != requirement.model_id:
            if actual_component != requirement.component and actual_model_id != requirement.model_id:
                problem = "component and model_id mismatch"
            elif actual_component != requirement.component:
                problem = "component mismatch"
            else:
                problem = "model_id mismatch"
            mismatches.append(
                SponsorMismatch(
                    scope=requirement.scope,
                    key=requirement.key,
                    expected_component=requirement.component,
                    expected_model_id=requirement.model_id,
                    actual_component=actual_component if isinstance(actual_component, str) else None,
                    actual_model_id=actual_model_id if isinstance(actual_model_id, str) else None,
                    problem=problem,
                    subkey=requirement.subkey,
                )
            )

    ok = not mismatches or waiver is not None
    return SponsorPolicyCheckResult(ok=ok, requirements=requirements, mismatches=tuple(mismatches), waiver=waiver)


def format_sponsor_policy_result(result: SponsorPolicyCheckResult) -> tuple[str, str, int]:
    if result.mismatches and result.waiver is None:
        lines = ["ERROR: sponsor model mismatch detected:"]
        for mismatch in result.mismatches:
            actual_component = mismatch.actual_component or "<missing>"
            actual_model_id = mismatch.actual_model_id or "<missing>"
            path = f"{mismatch.scope}.{mismatch.key}"
            if mismatch.subkey is not None:
                path = f"{path}.{mismatch.subkey}"
            lines.append(
                f"- {path}: expected component={mismatch.expected_component!r}, "
                f"model_id={mismatch.expected_model_id!r}; got component={actual_component!r}, "
                f"model_id={actual_model_id!r} ({mismatch.problem})"
            )
        lines.append("Fix configs/model_registry.yaml or provide configs/sponsor_model_waiver.yaml to justify the exception.")
        return ("", "\n".join(lines), 1)

    if result.mismatches and result.waiver is not None:
        warning_lines = [
            "WARNING: sponsor model mismatch waived.",
            f"- waiver file: {result.waiver.path}",
            f"- approved_by: {result.waiver.approved_by}",
            f"- date: {result.waiver.date}",
            f"- reason: {result.waiver.reason}",
        ]
        stdout = "Sponsor model policy check passed with waiver."
        return (stdout, "\n".join(warning_lines), 0)

    stdout = f"Sponsor model policy check passed: {len(result.requirements)} required sponsor model(s) aligned."
    return (stdout, "", 0)