"""Compile module envelopes into canonical runtime policy artifacts."""

from __future__ import annotations

import hashlib
import json
import os
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field

from .modules import (
    ModuleCompileError,
    ModuleEnvelope,
    ModuleValidationError,
    compute_revision_hash,
    get_module_definition,
    validate_module_envelope,
)
from .policy import AllowRule, ApprovalRule, AutonomyPolicy, DenyRule, RuleMatch, compute_policy_hash

try:  # Optional dependency in some environments.
    import yaml
except Exception as exc:  # pragma: no cover - handled at runtime
    yaml = None  # type: ignore
    _yaml_import_error = exc
else:  # pragma: no cover - no error path when yaml present
    _yaml_import_error = None


def _ensure_yaml() -> None:
    if yaml is None:  # pragma: no cover - raised only when dependency missing
        raise RuntimeError("PyYAML is required for runtime policy compilation.") from _yaml_import_error


def _canonical_json(payload: Any) -> str:
    return json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=True)


def _utc_now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _sanitize_rule_fragment(value: str) -> str:
    return value.replace(" ", "_")


def _module_rule_name(module_id: str, write_class_key: str, tier: Literal["green", "yellow", "red"]) -> str:
    return f"module:{_sanitize_rule_fragment(module_id)}:{_sanitize_rule_fragment(write_class_key)}:{tier}"


def _module_rule_tags(
    module_id: str,
    write_class_key: str,
    tier: Literal["green", "yellow", "red"],
) -> list[str]:
    return [
        f"module:{module_id}",
        f"write_class:{write_class_key}",
        f"risk_tier:{tier}",
    ]


def _editable_envelope_payload(envelope: ModuleEnvelope) -> dict[str, Any]:
    return {
        "enabled": envelope.enabled,
        "schedule_ref": envelope.schedule_ref,
        "exception_routing": envelope.exception_routing,
        "risk_tiers": envelope.risk_tiers.model_dump(exclude_none=False),
        "config": envelope.config,
    }


def _flatten(value: Any, prefix: str = "") -> dict[str, Any]:
    output: dict[str, Any] = {}
    if isinstance(value, dict):
        if not value:
            output[prefix] = {}
            return output
        for key in sorted(value.keys()):
            key_path = f"{prefix}.{key}" if prefix else str(key)
            output.update(_flatten(value[key], key_path))
        return output
    if isinstance(value, list):
        if not value:
            output[prefix] = []
            return output
        for index, item in enumerate(value):
            idx_path = f"{prefix}[{index}]"
            output.update(_flatten(item, idx_path))
        return output
    output[prefix] = value
    return output


class SimulationDecision(BaseModel):
    model_config = ConfigDict(extra="forbid")

    sample_index: int
    risk_score: int
    risk_tier: Literal["green", "yellow", "red"]
    decision: Literal["allow", "approval_required", "deny"]
    recommended_action: Literal["allow", "approval_required", "deny"]
    reason_code: str
    decision_path: list[str] = Field(default_factory=list)


class SimulationResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    module_id: str
    sample_source: Literal["user"] = "user"
    no_write: bool = True
    policy_hash: str
    decisions: list[SimulationDecision] = Field(default_factory=list)


class DiffEntry(BaseModel):
    model_config = ConfigDict(extra="forbid")

    path: str
    change_type: Literal["added", "removed", "changed"]
    old_value: Any = None
    new_value: Any = None


class DiffResult(BaseModel):
    model_config = ConfigDict(extra="forbid")

    module_id: str
    old_revision_hash: str
    new_revision_hash: str
    changes: list[DiffEntry] = Field(default_factory=list)
    policy_preview: dict[str, list[str]] = Field(default_factory=dict)


@dataclass(frozen=True)
class CompiledPolicyResult:
    policy: AutonomyPolicy
    policy_hash: str
    runtime_policy_path: Path
    runtime_meta_path: Path
    compiled_at: str
    revision_counter: int


def _normalize_module(
    module_payload: ModuleEnvelope | dict[str, Any],
    *,
    schedule_keys: set[str] | None,
) -> ModuleEnvelope:
    envelope = (
        module_payload
        if isinstance(module_payload, ModuleEnvelope)
        else ModuleEnvelope.model_validate(module_payload)
    )
    return validate_module_envelope(
        envelope,
        schedule_keys=schedule_keys,
        routing_targets=None,
        require_full_config=False,
        allow_unknown_config=True,
    )


def _write_yaml_atomic(path: Path, payload: dict[str, Any]) -> None:
    _ensure_yaml()
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_suffix(f"{path.suffix}.tmp")
    with open(tmp_path, "w", encoding="utf-8") as handle:
        yaml.safe_dump(
            payload,
            handle,
            sort_keys=False,
            default_flow_style=False,
            allow_unicode=False,
        )
    os.replace(str(tmp_path), str(path))


def _write_json_atomic(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_suffix(f"{path.suffix}.tmp")
    tmp_path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
    os.replace(str(tmp_path), str(path))


def _next_revision_counter(meta_path: Path) -> int:
    if not meta_path.exists():
        return 1
    try:
        payload = json.loads(meta_path.read_text(encoding="utf-8"))
    except Exception:
        return 1
    previous = payload.get("revision_counter")
    if isinstance(previous, int) and previous > 0:
        return previous + 1
    return 1


def compile_runtime_policy(
    state_dir: Path | str,
    base_policy: AutonomyPolicy,
    modules: list[ModuleEnvelope | dict[str, Any]],
) -> CompiledPolicyResult:
    """Compile enabled modules into canonical runtime policy artifacts."""
    state_path = Path(state_dir).resolve()
    schedule_keys = set(base_policy.schedules.keys())
    try:
        normalized_modules = [_normalize_module(module, schedule_keys=schedule_keys) for module in modules]
    except ModuleValidationError as exc:
        raise ModuleCompileError(str(exc)) from exc
    enabled_modules = [module for module in normalized_modules if module.enabled]
    enabled_modules.sort(key=lambda item: item.module_id)

    allow_overlays: list[AllowRule] = []
    approval_overlays: list[ApprovalRule] = []
    deny_overlays: list[DenyRule] = []
    owners_by_write_class: dict[str, str] = {}

    for module in enabled_modules:
        if module.schedule_ref and module.schedule_ref not in schedule_keys:
            raise ModuleCompileError(
                f"module {module.module_id} references unknown schedule_ref '{module.schedule_ref}'"
            )
        definition = get_module_definition(module.module_id)
        write_class_keys = sorted(definition.write_class_keys)
        for write_class_key in write_class_keys:
            existing_owner = owners_by_write_class.get(write_class_key)
            if existing_owner and existing_owner != module.module_id:
                raise ModuleCompileError(
                    "duplicate write-class key detected for enabled modules: "
                    f"{write_class_key} ({existing_owner}, {module.module_id})"
                )
            owners_by_write_class[write_class_key] = module.module_id

            allow_overlays.append(
                AllowRule(
                    name=_module_rule_name(module.module_id, write_class_key, "green"),
                    match=RuleMatch(tags=_module_rule_tags(module.module_id, write_class_key, "green")),
                    auto_execute=True,
                )
            )
            approval_overlays.append(
                ApprovalRule(
                    name=_module_rule_name(module.module_id, write_class_key, "yellow"),
                    match=RuleMatch(tags=_module_rule_tags(module.module_id, write_class_key, "yellow")),
                    approvers=[],
                )
            )
            deny_overlays.append(
                DenyRule(
                    name=_module_rule_name(module.module_id, write_class_key, "red"),
                    match=RuleMatch(tags=_module_rule_tags(module.module_id, write_class_key, "red")),
                    reason_code=f"module:{module.module_id}:red_tier",
                )
            )

    compiled = base_policy.model_copy(deep=True)
    compiled = compiled.model_copy(
        update={
            "allow_rules": list(compiled.allow_rules) + allow_overlays,
            "approval_rules": list(compiled.approval_rules) + approval_overlays,
            "deny_rules": list(compiled.deny_rules) + deny_overlays,
        }
    )
    policy_hash = compute_policy_hash(compiled)
    compiled_at = _utc_now_iso()

    runtime_policy_path = state_path / "policy_runtime.yaml"
    runtime_meta_path = state_path / "policy_runtime.meta.json"
    revision_counter = _next_revision_counter(runtime_meta_path)
    _write_yaml_atomic(
        runtime_policy_path,
        compiled.model_dump(by_alias=True, exclude_none=True),
    )
    _write_json_atomic(
        runtime_meta_path,
        {
            "policy_hash": policy_hash,
            "compiled_at": compiled_at,
            "revision_counter": revision_counter,
        },
    )
    return CompiledPolicyResult(
        policy=compiled,
        policy_hash=policy_hash,
        runtime_policy_path=runtime_policy_path,
        runtime_meta_path=runtime_meta_path,
        compiled_at=compiled_at,
        revision_counter=revision_counter,
    )


def _resolve_risk_score(sample: dict[str, Any]) -> int:
    for key in ("risk_score", "score"):
        raw = sample.get(key)
        if isinstance(raw, (int, float)):
            bounded = int(raw)
            return max(0, min(100, bounded))
    digest = hashlib.sha256(_canonical_json(sample).encode("utf-8")).hexdigest()
    return int(digest[:8], 16) % 101


def simulate_module_on_samples(
    module_id: str,
    envelope: ModuleEnvelope | dict[str, Any],
    samples: list[dict[str, Any]],
    policy: AutonomyPolicy,
) -> SimulationResult:
    """Run deterministic no-write simulation for module candidate actions."""
    if not isinstance(samples, list):
        raise ModuleValidationError("samples must be a list")

    normalized = _normalize_module(envelope, schedule_keys=set(policy.schedules.keys()))
    if normalized.module_id != module_id:
        raise ModuleValidationError(
            f"module_id mismatch between path ({module_id}) and envelope ({normalized.module_id})"
        )

    decisions: list[SimulationDecision] = []
    for index, sample in enumerate(samples):
        if not isinstance(sample, dict):
            raise ModuleValidationError(f"sample at index {index} must be an object")

        score = _resolve_risk_score(sample)
        decision_path = [
            f"module_enabled:{'1' if normalized.enabled else '0'}",
            f"risk_score:{score}",
        ]

        if not normalized.enabled:
            decisions.append(
                SimulationDecision(
                    sample_index=index,
                    risk_score=score,
                    risk_tier="red",
                    decision="deny",
                    recommended_action="deny",
                    reason_code=f"module:{module_id}:disabled",
                    decision_path=decision_path + ["decision:deny:module_disabled"],
                )
            )
            continue

        risk_tiers = normalized.risk_tiers
        if score <= risk_tiers.green_max_score:
            tier: Literal["green", "yellow", "red"] = "green"
            decision: Literal["allow", "approval_required", "deny"] = "allow"
            reason = f"module:{module_id}:green_tier"
        elif score <= risk_tiers.yellow_max_score:
            tier = "yellow"
            decision = "approval_required"
            reason = f"module:{module_id}:yellow_tier"
        else:
            tier = "red"
            decision = "deny"
            reason = f"module:{module_id}:red_tier"

        decisions.append(
            SimulationDecision(
                sample_index=index,
                risk_score=score,
                risk_tier=tier,
                decision=decision,
                recommended_action=decision,
                reason_code=reason,
                decision_path=decision_path + [f"risk_tier:{tier}", f"decision:{decision}:{reason}"],
            )
        )

    return SimulationResult(
        module_id=module_id,
        sample_source="user",
        no_write=True,
        policy_hash=compute_policy_hash(policy),
        decisions=decisions,
    )


def _module_policy_preview(envelope: ModuleEnvelope) -> list[str]:
    if not envelope.enabled:
        return []
    definition = get_module_definition(envelope.module_id)
    names: list[str] = []
    for write_class_key in sorted(definition.write_class_keys):
        names.extend(
            [
                _module_rule_name(envelope.module_id, write_class_key, "green"),
                _module_rule_name(envelope.module_id, write_class_key, "yellow"),
                _module_rule_name(envelope.module_id, write_class_key, "red"),
            ]
        )
    return names


def diff_module_revision(
    module_id: str,
    old_env: ModuleEnvelope | dict[str, Any],
    new_env: ModuleEnvelope | dict[str, Any],
) -> DiffResult:
    """Return semantic envelope diff and policy preview deltas."""
    old_normalized = _normalize_module(old_env, schedule_keys=None)
    new_normalized = _normalize_module(new_env, schedule_keys=None)
    if old_normalized.module_id != module_id or new_normalized.module_id != module_id:
        raise ModuleValidationError("diff module_id does not match one or both envelope payloads")

    old_payload = _editable_envelope_payload(old_normalized)
    new_payload = _editable_envelope_payload(new_normalized)
    old_flat = _flatten(old_payload)
    new_flat = _flatten(new_payload)

    missing = object()
    changes: list[DiffEntry] = []
    for path in sorted(set(old_flat.keys()) | set(new_flat.keys())):
        old_value = old_flat.get(path, missing)
        new_value = new_flat.get(path, missing)
        if old_value is missing:
            changes.append(
                DiffEntry(
                    path=path,
                    change_type="added",
                    old_value=None,
                    new_value=new_value,
                )
            )
            continue
        if new_value is missing:
            changes.append(
                DiffEntry(
                    path=path,
                    change_type="removed",
                    old_value=old_value,
                    new_value=None,
                )
            )
            continue
        if old_value != new_value:
            changes.append(
                DiffEntry(
                    path=path,
                    change_type="changed",
                    old_value=old_value,
                    new_value=new_value,
                )
            )

    old_rules = _module_policy_preview(old_normalized)
    new_rules = _module_policy_preview(new_normalized)
    old_set = set(old_rules)
    new_set = set(new_rules)

    return DiffResult(
        module_id=module_id,
        # Always compute from editable fields so client-provided revision tokens
        # cannot mask actual content changes in diff previews.
        old_revision_hash=compute_revision_hash(old_normalized),
        new_revision_hash=compute_revision_hash(new_normalized),
        changes=changes,
        policy_preview={
            "added_rules": sorted(new_set - old_set),
            "removed_rules": sorted(old_set - new_set),
        },
    )


__all__ = [
    "CompiledPolicyResult",
    "SimulationDecision",
    "SimulationResult",
    "DiffEntry",
    "DiffResult",
    "compile_runtime_policy",
    "simulate_module_on_samples",
    "diff_module_revision",
]
