# sovereign_infinity_layer.py
# Sovereign-∞ Security & Trust Layer
# این لایه روی Sovereign Core سوار می‌شود بدون اینکه قابلیت‌های قبلی را بشکند.

from dataclasses import dataclass, field
from typing import Any, Dict, Optional, Callable, List
import time

# ---------------------------------------------------------------------
# Context اصلی که بین همه لایه‌ها رد می‌شود
# ---------------------------------------------------------------------
@dataclass
class SovereignContext:
    request_id: str
    model_name: str
    engine_name: str
    environment: str
    user_id: Optional[str] = None
    tenant_id: Optional[str] = None
    prompt: Optional[str] = None
    response: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: int = field(default_factory=lambda: int(time.time()))


# ---------------------------------------------------------------------
# Trust Score Engine – «تراست» در Sovereign
# ---------------------------------------------------------------------
class TrustScoreEngine:
    """
    یک موتور خیلی ساده ولی قابل‌گسترش برای تبدیل سیگنال‌ها به Trust Score.
    ورودی: سیگنال‌ها (lineage_ok, fp_match, cve_score, anomalies, policy_violations, ...)
    خروجی: عدد 0 تا 100
    """
    def __init__(self):
        # اگر خواستی بعداً این weights را از config / policy بخوانیم
        self.weights = {
            "lineage_ok": 25,
            "fingerprint_ok": 25,
            "no_critical_cve": 20,
            "low_anomaly": 15,
            "no_policy_violation": 15,
        }

    def calculate(self, signals: Dict[str, Any]) -> int:
        score = 0

        if signals.get("lineage_ok", False):
            score += self.weights["lineage_ok"]

        if signals.get("fingerprint_ok", False):
            score += self.weights["fingerprint_ok"]

        if signals.get("no_critical_cve", True):
            score += self.weights["no_critical_cve"]

        anomaly_score = signals.get("anomaly_score", 0.0)  # 0 تا 1
        if anomaly_score <= 0.2:
            score += self.weights["low_anomaly"]

        if signals.get("policy_violations", 0) == 0:
            score += self.weights["no_policy_violation"]

        # محدود به 0…100
        score = max(0, min(100, score))
        return score


# ---------------------------------------------------------------------
# Self-Sealing Engine – واکنش به رویدادهای خطرناک
# ---------------------------------------------------------------------
class SelfSealingEngine:
    """
    مسئول این است که اگر رویداد خطرناک رخ داد، سیستم را «خودکار ایمن» کند:
    - تریگر kill-switch
    - بستن session
    - قفل کردن deployment
    """

    def __init__(
        self,
        kill_switch_callback: Optional[Callable[[SovereignContext, Dict[str, Any]], None]] = None,
        freeze_env_callback: Optional[Callable[[SovereignContext, Dict[str, Any]], None]] = None,
    ):
        self.kill_switch_callback = kill_switch_callback
        self.freeze_env_callback = freeze_env_callback

    def handle_event(self, ctx: SovereignContext, event: Dict[str, Any]) -> None:
        """
        event می‌تواند شامل:
        {
            "severity": "low|medium|high|critical",
            "reason": "...",
            "signals": {...}
        }
        """
        severity = event.get("severity", "low")

        # اگر سطح «high» یا «critical» بود، وارد حالت self-sealing شویم
        if severity in ("high", "critical"):
            if self.freeze_env_callback:
                self.freeze_env_callback(ctx, event)
            if self.kill_switch_callback:
                self.kill_switch_callback(ctx, event)


# ---------------------------------------------------------------------
# Infinity GuardRail – pre + post guard برای Sovereign
# ---------------------------------------------------------------------
class InfinityGuardRail:
    """
    لایه اصلی کنترل قبل و بعد از کال مدل.
    - pre_guard: بررسی policy, environment binding, lineage, fingerprint
    - post_guard: بررسی خروجی، anomaly، risk
    """

    def __init__(
        self,
        # توابعی که از Core به ما پاس می‌شوند:
        core_check_env: Callable[[SovereignContext], bool],
        core_check_lineage: Callable[[SovereignContext], bool],
        core_check_fingerprint: Callable[[SovereignContext], bool],
        core_policy_evaluate: Callable[[SovereignContext], Dict[str, Any]],
        core_log_audit: Callable[[SovereignContext, Dict[str, Any]], None],
        anomaly_detector: Optional[Callable[[SovereignContext], float]] = None,
        cve_resolver: Optional[Callable[[SovereignContext], Dict[str, Any]]] = None,
        self_sealing: Optional[SelfSealingEngine] = None,
        trust_engine: Optional[TrustScoreEngine] = None,
    ):
        self.core_check_env = core_check_env
        self.core_check_lineage = core_check_lineage
        self.core_check_fingerprint = core_check_fingerprint
        self.core_policy_evaluate = core_policy_evaluate
        self.core_log_audit = core_log_audit

        self.anomaly_detector = anomaly_detector
        self.cve_resolver = cve_resolver
        self.self_sealing = self_sealing or SelfSealingEngine()
        self.trust_engine = trust_engine or TrustScoreEngine()

    # ----------------- PRE -----------------
    def pre_guard(self, ctx: SovereignContext) -> Dict[str, Any]:
        """
        قبل از کال مدل، همه چیز را چک می‌کند.
        اگر چیزی مشکل جدی داشته باشد، می‌تواند اجازه ندهد ادامه دهیم.
        """
        lineage_ok = self.core_check_lineage(ctx)
        fp_ok = self.core_check_fingerprint(ctx)
        env_ok = self.core_check_env(ctx)
        policy_res = self.core_policy_evaluate(ctx)

        policy_allowed = policy_res.get("allowed", True)
        policy_violations = policy_res.get("violations", [])

        signals = {
            "lineage_ok": lineage_ok,
            "fingerprint_ok": fp_ok,
            "env_ok": env_ok,
            "policy_violations": len(policy_violations),
        }

        cve_info = {}
        if self.cve_resolver:
            cve_info = self.cve_resolver(ctx)
            # فرض: cve_info مثلا {"has_critical": False}
            signals["no_critical_cve"] = not cve_info.get("has_critical", False)
        else:
            signals["no_critical_cve"] = True

        # اگر env یا lineage یا policy مشکل جدی دارند → رویداد خطرناک
        if not env_ok or not lineage_ok or not fp_ok or not policy_allowed:
            event = {
                "severity": "critical",
                "reason": "pre_guard_failed",
                "signals": signals,
                "policy_violations": policy_violations,
            }
            self.self_sealing.handle_event(ctx, event)

            self.core_log_audit(ctx, {
                "type": "pre_guard_block",
                "severity": "critical",
                "signals": signals,
                "policy_violations": policy_violations,
            })

            return {
                "allowed": False,
                "trust_score": 0,
                "signals": signals,
                "policy_violations": policy_violations,
                "cve": cve_info,
            }

        # اگر همه چیز OK بود، فعلاً فقط لاگ می‌زنیم
        self.core_log_audit(ctx, {
            "type": "pre_guard_pass",
            "signals": signals,
            "policy_violations": policy_violations,
            "cve": cve_info,
        })

        return {
            "allowed": True,
            "signals": signals,
            "policy_violations": policy_violations,
            "cve": cve_info,
        }

    # ----------------- POST -----------------
    def post_guard(self, ctx: SovereignContext) -> Dict[str, Any]:
        """
        بعد از کال مدل، خروجی را برای ریسک، انحراف رفتار و ... بررسی می‌کند.
        """
        anomaly_score = 0.0
        if self.anomaly_detector:
            anomaly_score = self.anomaly_detector(ctx)

        signals = {
            "anomaly_score": anomaly_score,
        }

        # Trust score نهایی بر اساس ترکیب pre + post
        trust_signals = {
            "lineage_ok": ctx.metadata.get("lineage_ok", True),
            "fingerprint_ok": ctx.metadata.get("fingerprint_ok", True),
            "no_critical_cve": ctx.metadata.get("no_critical_cve", True),
            "anomaly_score": anomaly_score,
            "policy_violations": ctx.metadata.get("policy_violations", 0),
        }

        trust_score = self.trust_engine.calculate(trust_signals)

        # اگر anomaly شدید باشد → self-sealing
        severity = "low"
        if anomaly_score > 0.8:
            severity = "critical"
        elif anomaly_score > 0.5:
            severity = "high"
        elif anomaly_score > 0.3:
            severity = "medium"

        if severity in ("high", "critical"):
            event = {
                "severity": severity,
                "reason": "high_anomaly_response",
                "signals": trust_signals,
            }
            self.self_sealing.handle_event(ctx, event)

        self.core_log_audit(ctx, {
            "type": "post_guard",
            "severity": severity,
            "anomaly_score": anomaly_score,
            "trust_score": trust_score,
            "signals": trust_signals,
        })

        return {
            "trust_score": trust_score,
            "severity": severity,
            "anomaly_score": anomaly_score,
        }


# ---------------------------------------------------------------------
# Orchestrator – یک entrypoint ساده برای استفاده
# ---------------------------------------------------------------------
class SovereignInfinityOrchestrator:
    """
    این کلاس کاری می‌کند که تو فقط یک تابع مدل بدهی
    و Sovereign-∞ همه pre + post + audit + trust را انجام دهد.
    """

    def __init__(self, guard: InfinityGuardRail):
        self.guard = guard

    def secure_model_call(
        self,
        ctx: SovereignContext,
        model_call_fn: Callable[[SovereignContext], str],
    ) -> Dict[str, Any]:
        """
        مدل را فقط از این مسیر صدا بزن.
        - اول pre_guard
        - اگر اوکی بود → کال مدل
        - بعد post_guard
        """
        pre = self.guard.pre_guard(ctx)

        # ذخیره برخی سیگنال‌ها داخل متادیتای context
        ctx.metadata["lineage_ok"] = pre["signals"].get("lineage_ok", True)
        ctx.metadata["fingerprint_ok"] = pre["signals"].get("fingerprint_ok", True)
        ctx.metadata["no_critical_cve"] = pre["signals"].get("no_critical_cve", True)
        ctx.metadata["policy_violations"] = pre.get("policy_violations", 0)

        if not pre.get("allowed", True):
            return {
                "allowed": False,
                "reason": "blocked_by_pre_guard",
                "pre": pre,
                "response": None,
            }

        # اگر مجاز بود، مدل را صدا بزنیم
        response = model_call_fn(ctx)
        ctx.response = response

        post = self.guard.post_guard(ctx)

        return {
            "allowed": True,
            "response": response,
            "pre": pre,
            "post": post,
        }