# sovereign_freeze_bridge.py
from __future__ import annotations

import json
import os
import uuid
from dataclasses import dataclass, asdict
from datetime import datetime, timezone
from typing import Any, Dict, Optional, List, Tuple


FREEZE_CASES_FILE = os.getenv("SOV_FREEZE_CASES_FILE", "sovereign_freeze_cases.json")
FREEZE_INDEX_FILE = os.getenv("SOV_FREEZE_INDEX_FILE", "sovereign_freeze_index.json")
FREEZE_AUDIT_FILE = os.getenv("SOV_FREEZE_AUDIT_FILE", "sovereign_freeze_audit.jsonl")


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


def _safe_read_json(path: str, default: Any) -> Any:
    try:
        if not os.path.exists(path):
            return default
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return default


def _safe_write_json(path: str, data: Any) -> None:
    directory = os.path.dirname(path)
    if directory:
        os.makedirs(directory, exist_ok=True)
    tmp = f"{path}.tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2, sort_keys=True)
    os.replace(tmp, path)


def _append_jsonl(path: str, row: Dict[str, Any]) -> None:
    directory = os.path.dirname(path)
    if directory:
        os.makedirs(directory, exist_ok=True)
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n")


def _norm_text(v: Any) -> str:
    return str(v or "").strip()


def _norm_tags(v: Any) -> List[str]:
    if isinstance(v, list):
        return [str(x).strip().lower() for x in v if str(x).strip()]
    if isinstance(v, str):
        return [x.strip().lower() for x in v.split(",") if x.strip()]
    return []


def _bucket_key(payload: Dict[str, Any]) -> str:
    engine = _norm_text(payload.get("engine") or payload.get("engine_name") or "unknown_engine")
    parent = _norm_text(payload.get("parent_model") or payload.get("parent_model_id") or payload.get("agent_id") or "unknown_parent")
    version = _norm_text(payload.get("model_version") or payload.get("build_id") or "unknown_version")
    env = _norm_text(payload.get("deployment_env") or payload.get("environment") or "unspecified")
    tags = ",".join(sorted(_norm_tags(payload.get("data_tags"))))
    return f"{engine}|{parent}|{version}|{env}|{tags}"


@dataclass
class FreezeCase:
    freeze_case_id: str
    created_at: str
    updated_at: str
    status: str                  # frozen / released / escalated_block
    engine: str
    parent_model: str
    model_version: str
    deployment_env: str
    data_tags: List[str]
    risk_level: str
    decision_id: str
    reason: str
    source: str                  # authority_gate / cognition / zero_layer / runtime
    authority_decision: str      # FREEZE / BLOCK / ALLOW
    freeze_token: str
    review_required: bool
    released_at: Optional[str] = None
    released_by: Optional[str] = None
    release_note: Optional[str] = None
    escalated_at: Optional[str] = None
    escalated_by: Optional[str] = None
    escalation_note: Optional[str] = None
    audit_trail: Optional[List[Dict[str, Any]]] = None


class SovereignFreezeBridge:
    """
    Bridges FREEZE decisions into a bank-grade freeze case lifecycle.

    What it does:
    - when decision == FREEZE -> create a persistent freeze case
    - prevents execution while case is still frozen
    - if case gets released -> execution may continue
    - if case gets escalated -> treat as BLOCK
    - writes dedicated freeze audit trail
    """

    def __init__(
        self,
        cases_file: str = FREEZE_CASES_FILE,
        index_file: str = FREEZE_INDEX_FILE,
        audit_file: str = FREEZE_AUDIT_FILE,
    ) -> None:
        self.cases_file = cases_file
        self.index_file = index_file
        self.audit_file = audit_file

    # -------------------------
    # storage
    # -------------------------
    def _load_cases(self) -> Dict[str, Dict[str, Any]]:
        obj = _safe_read_json(self.cases_file, {})
        return obj if isinstance(obj, dict) else {}

    def _save_cases(self, obj: Dict[str, Dict[str, Any]]) -> None:
        _safe_write_json(self.cases_file, obj)

    def _load_index(self) -> Dict[str, str]:
        obj = _safe_read_json(self.index_file, {})
        return obj if isinstance(obj, dict) else {}

    def _save_index(self, obj: Dict[str, str]) -> None:
        _safe_write_json(self.index_file, obj)

    def _audit(self, event_type: str, freeze_case_id: str, payload: Dict[str, Any]) -> None:
        _append_jsonl(
            self.audit_file,
            {
                "ts": _utc_now(),
                "event_type": event_type,
                "freeze_case_id": freeze_case_id,
                "payload": payload,
            },
        )

    # -------------------------
    # helpers
    # -------------------------
    def _extract_reason(self, decision_bundle: Dict[str, Any]) -> str:
        reason = decision_bundle.get("reason")
        if isinstance(reason, str) and reason.strip():
            return reason.strip()

        reasons = decision_bundle.get("reasons")
        if isinstance(reasons, list) and reasons:
            return " | ".join(str(x) for x in reasons[:6])

        authority = decision_bundle.get("authority_gate") or decision_bundle.get("authority") or {}
        if isinstance(authority, dict):
            r = authority.get("reason")
            if isinstance(r, str) and r.strip():
                return r.strip()

        return "freeze_without_explicit_reason"

    def _extract_decision(self, decision_bundle: Dict[str, Any]) -> str:
        for key in ("decision", "action", "outcome", "authority_decision"):
            val = decision_bundle.get(key)
            if isinstance(val, str) and val.strip():
                return val.strip().upper()
        auth = decision_bundle.get("authority_gate") or {}
        if isinstance(auth, dict):
            val = auth.get("decision")
            if isinstance(val, str) and val.strip():
                return val.strip().upper()
        return "UNKNOWN"

    def _extract_source(self, decision_bundle: Dict[str, Any]) -> str:
        if "authority_gate" in decision_bundle:
            return "authority_gate"
        if "zero_layer" in decision_bundle:
            return "zero_layer"
        if "cognition" in decision_bundle:
            return "cognition"
        return "runtime"

    def _extract_decision_id(self, decision_bundle: Dict[str, Any]) -> str:
        for key in ("decision_id", "event_id", "audit_id"):
            v = decision_bundle.get(key)
            if isinstance(v, str) and v.strip():
                return v.strip()
        auth = decision_bundle.get("authority_gate") or {}
        if isinstance(auth, dict):
            v = auth.get("decision_id")
            if isinstance(v, str) and v.strip():
                return v.strip()
        return str(uuid.uuid4())

    def _append_case_event(self, case_obj: Dict[str, Any], event_type: str, actor: str, note: str, extra: Optional[Dict[str, Any]] = None) -> None:
        trail = case_obj.get("audit_trail")
        if not isinstance(trail, list):
            trail = []
        trail.append(
            {
                "ts": _utc_now(),
                "event_type": event_type,
                "actor": actor,
                "note": note,
                "extra": extra or {},
            }
        )
        case_obj["audit_trail"] = trail
        case_obj["updated_at"] = _utc_now()

    # -------------------------
    # public operations
    # -------------------------
    def create_or_link_freeze_case(
        self,
        payload: Dict[str, Any],
        decision_bundle: Dict[str, Any],
    ) -> Dict[str, Any]:
        decision = self._extract_decision(decision_bundle)
        if decision != "FREEZE":
            return {
                "freeze_created": False,
                "freeze_linked": False,
                "freeze_case_id": None,
                "status": None,
                "reason": "decision_is_not_freeze",
            }

        cases = self._load_cases()
        index = self._load_index()

        bucket = _bucket_key(payload)
        existing_case_id = index.get(bucket)

        if existing_case_id and existing_case_id in cases:
            case_obj = cases[existing_case_id]
            status = case_obj.get("status", "frozen")
            self._append_case_event(
                case_obj,
                event_type="freeze_reaffirmed",
                actor="system",
                note="Freeze re-asserted for existing case.",
                extra={"decision_id": self._extract_decision_id(decision_bundle)},
            )
            cases[existing_case_id] = case_obj
            self._save_cases(cases)
            self._audit("freeze_reaffirmed", existing_case_id, {"bucket": bucket, "status": status})
            return {
                "freeze_created": False,
                "freeze_linked": True,
                "freeze_case_id": existing_case_id,
                "status": status,
                "reason": "existing_case_reused",
            }

        freeze_case_id = str(uuid.uuid4())
        freeze_token = str(uuid.uuid4())

        case = FreezeCase(
            freeze_case_id=freeze_case_id,
            created_at=_utc_now(),
            updated_at=_utc_now(),
            status="frozen",
            engine=_norm_text(payload.get("engine") or payload.get("engine_name") or "unknown_engine"),
            parent_model=_norm_text(payload.get("parent_model") or payload.get("parent_model_id") or payload.get("agent_id") or "unknown_parent"),
            model_version=_norm_text(payload.get("model_version") or payload.get("build_id") or "unknown_version"),
            deployment_env=_norm_text(payload.get("deployment_env") or payload.get("environment") or "unspecified"),
            data_tags=_norm_tags(payload.get("data_tags")),
            risk_level=_norm_text(payload.get("risk_level") or "unspecified").lower(),
            decision_id=self._extract_decision_id(decision_bundle),
            reason=self._extract_reason(decision_bundle),
            source=self._extract_source(decision_bundle),
            authority_decision="FREEZE",
            freeze_token=freeze_token,
            review_required=True,
            audit_trail=[],
        )

        case_obj = asdict(case)
        self._append_case_event(
            case_obj,
            event_type="freeze_created",
            actor="system",
            note="Freeze case created from runtime decision.",
            extra={"bucket": bucket},
        )

        cases[freeze_case_id] = case_obj
        index[bucket] = freeze_case_id

        self._save_cases(cases)
        self._save_index(index)
        self._audit("freeze_created", freeze_case_id, {"bucket": bucket, "engine": case.engine})

        return {
            "freeze_created": True,
            "freeze_linked": False,
            "freeze_case_id": freeze_case_id,
            "freeze_token": freeze_token,
            "status": "frozen",
            "reason": case.reason,
        }

    def get_case(self, freeze_case_id: str) -> Optional[Dict[str, Any]]:
        cases = self._load_cases()
        obj = cases.get(freeze_case_id)
        return obj if isinstance(obj, dict) else None

    def get_case_by_payload(self, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        index = self._load_index()
        cases = self._load_cases()
        case_id = index.get(_bucket_key(payload))
        if not case_id:
            return None
        obj = cases.get(case_id)
        return obj if isinstance(obj, dict) else None

    def effective_state(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        """
        Returns the effective enforcement state for the payload:
        - no_case
        - frozen
        - released
        - escalated_block
        """
        case_obj = self.get_case_by_payload(payload)
        if not case_obj:
            return {
                "state": "no_case",
                "allow_execution": True,
                "freeze_case_id": None,
                "reason": "no_freeze_case",
            }

        status = case_obj.get("status", "frozen")
        case_id = case_obj.get("freeze_case_id")

        if status == "released":
            return {
                "state": "released",
                "allow_execution": True,
                "freeze_case_id": case_id,
                "reason": "freeze_released",
            }

        if status == "escalated_block":
            return {
                "state": "escalated_block",
                "allow_execution": False,
                "freeze_case_id": case_id,
                "reason": "freeze_escalated_to_block",
            }

        return {
            "state": "frozen",
            "allow_execution": False,
            "freeze_case_id": case_id,
            "reason": "freeze_active",
        }

    def can_execute(self, payload: Dict[str, Any]) -> Tuple[bool, Dict[str, Any]]:
        state = self.effective_state(payload)
        return bool(state.get("allow_execution")), state

    def release_case(self, freeze_case_id: str, released_by: str, note: str = "") -> Dict[str, Any]:
        cases = self._load_cases()
        obj = cases.get(freeze_case_id)
        if not isinstance(obj, dict):
            return {"ok": False, "error": "freeze_case_not_found"}

        status = obj.get("status")
        if status == "escalated_block":
            return {"ok": False, "error": "case_already_escalated_block"}
        if status == "released":
            return {"ok": True, "freeze_case_id": freeze_case_id, "status": "released", "idempotent": True}

        obj["status"] = "released"
        obj["released_at"] = _utc_now()
        obj["released_by"] = _norm_text(released_by) or "unknown_reviewer"
        obj["release_note"] = _norm_text(note)

        self._append_case_event(
            obj,
            event_type="freeze_released",
            actor=obj["released_by"],
            note=obj["release_note"] or "Freeze released.",
        )

        cases[freeze_case_id] = obj
        self._save_cases(cases)
        self._audit("freeze_released", freeze_case_id, {"released_by": obj["released_by"]})

        return {
            "ok": True,
            "freeze_case_id": freeze_case_id,
            "status": "released",
            "released_at": obj["released_at"],
        }

    def escalate_case(self, freeze_case_id: str, escalated_by: str, note: str = "") -> Dict[str, Any]:
        cases = self._load_cases()
        obj = cases.get(freeze_case_id)
        if not isinstance(obj, dict):
            return {"ok": False, "error": "freeze_case_not_found"}

        if obj.get("status") == "escalated_block":
            return {"ok": True, "freeze_case_id": freeze_case_id, "status": "escalated_block", "idempotent": True}

        obj["status"] = "escalated_block"
        obj["authority_decision"] = "BLOCK"
        obj["escalated_at"] = _utc_now()
        obj["escalated_by"] = _norm_text(escalated_by) or "unknown_reviewer"
        obj["escalation_note"] = _norm_text(note)

        self._append_case_event(
            obj,
            event_type="freeze_escalated_block",
            actor=obj["escalated_by"],
            note=obj["escalation_note"] or "Freeze escalated to block.",
        )

        cases[freeze_case_id] = obj
        self._save_cases(cases)
        self._audit("freeze_escalated_block", freeze_case_id, {"escalated_by": obj["escalated_by"]})

        return {
            "ok": True,
            "freeze_case_id": freeze_case_id,
            "status": "escalated_block",
            "escalated_at": obj["escalated_at"],
        }

    def register_decision(self, payload: Dict[str, Any], decision_bundle: Dict[str, Any]) -> Dict[str, Any]:
        """
        Main bridge entrypoint.
        Call this right after authority/runtime decision is made.
        """
        decision = self._extract_decision(decision_bundle)

        if decision == "FREEZE":
            out = self.create_or_link_freeze_case(payload, decision_bundle)
            return {
                "bridge_status": "freeze_registered",
                "decision": "FREEZE",
                "freeze": out,
            }

        # For non-freeze decisions, still check if an existing freeze case blocks execution
        exec_ok, state = self.can_execute(payload)

        if not exec_ok:
            forced = "BLOCK" if state.get("state") == "escalated_block" else "FREEZE"
            return {
                "bridge_status": "existing_freeze_overrides_execution",
                "decision": forced,
                "freeze": state,
            }

        return {
            "bridge_status": "no_freeze_intervention",
            "decision": decision,
            "freeze": state,
        }


if __name__ == "__main__":
    bridge = SovereignFreezeBridge()

    sample_payload = {
        "engine": "banking_runtime_v1",
        "parent_model": "agent_alpha",
        "model_version": "v1",
        "deployment_env": "production",
        "data_tags": ["pii", "payments", "customer_chat"],
        "risk_level": "high",
    }

    sample_decision = {
        "decision_id": str(uuid.uuid4()),
        "decision": "FREEZE",
        "reason": "suspicious_cross_account_instruction",
        "authority_gate": {"decision": "FREEZE", "reason": "authority_freeze"},
    }

    print(json.dumps(bridge.register_decision(sample_payload, sample_decision), indent=2, ensure_ascii=False))