"""
eaigins_sdk.models
──────────────────
Core data models for the Eaigins Enterprise AI Control Plane SDK.
All models are dataclasses with JSON-serialisable export methods.
"""

from __future__ import annotations

import uuid
from dataclasses import dataclass, field, asdict
from datetime import datetime
from enum import Enum
from typing import Any, Optional


# ── Enumerations ──────────────────────────────────────────────────────────────

class RiskLevel(str, Enum):
    """Ordered risk classification from the weighted scoring formula."""
    LOW      = "LOW"
    MEDIUM   = "MEDIUM"
    HIGH     = "HIGH"
    CRITICAL = "CRITICAL"


class ExecutionDecision(str, Enum):
    """Final governance verdict for an agent action."""
    ALLOW    = "ALLOW"     # Action may proceed to enterprise systems
    ESCALATE = "ESCALATE"  # Action requires human approval before execution
    BLOCK    = "BLOCK"     # Action is prevented; reason recorded in provenance


class ThreatType(str, Enum):
    """Threat categories detected by the AI Agent Firewall."""
    PROMPT_INJECTION       = "PROMPT_INJECTION"
    SENSITIVE_DATA_EXPOSURE = "SENSITIVE_DATA_EXPOSURE"
    UNAUTHORIZED_TOOL      = "UNAUTHORIZED_TOOL"
    MALICIOUS_CONTEXT      = "MALICIOUS_CONTEXT"
    POLICY_VIOLATION       = "POLICY_VIOLATION"
    NONE                   = "NONE"


class NodeType(str, Enum):
    """Provenance graph node types."""
    PROMPT          = "PROMPT"
    CONTEXT_SOURCE  = "CONTEXT_SOURCE"
    REASONING       = "REASONING"
    POLICY_EVAL     = "POLICY_EVAL"
    RISK_SCORE      = "RISK_SCORE"
    TOOL_INVOCATION = "TOOL_INVOCATION"
    ACTION          = "ACTION"
    BLOCK           = "BLOCK"


# ── Core input model ──────────────────────────────────────────────────────────

@dataclass
class AgentAction:
    """
    Represents a single AI agent action to be governed.

    This is the primary input to the control plane.  Build one before
    every tool call your agent makes, then pass it to
    ``ControlPlane.evaluate(action)``.

    Example::

        action = AgentAction(
            agent_id        = "my-agent-v1",
            prompt          = "Retrieve all customer records for analysis.",
            tool_name       = "export_customer_data",
            tool_params     = {"scope": "ALL", "include_pii": True},
            context_sources = ["crm_customer_profile", "payment_records"],
            reasoning       = "Marketing team requested full dataset.",
            metadata        = {"regulatory_framework": "GDPR"},
        )
    """
    # Auto-generated
    action_id:       str      = field(default_factory=lambda: str(uuid.uuid4())[:8])
    timestamp:       datetime = field(default_factory=datetime.now)

    # Agent identity
    agent_id:        str      = "unknown-agent"

    # The action content
    prompt:          str      = ""               # Natural language intent
    tool_name:       str      = ""               # Enterprise tool being invoked
    tool_params:     dict     = field(default_factory=dict)   # Tool arguments
    context_sources: list[str] = field(default_factory=list)  # RAG/retrieval sources
    reasoning:       str      = ""               # Agent's stated reasoning
    metadata:        dict     = field(default_factory=dict)   # Domain context

    def to_dict(self) -> dict:
        d = asdict(self)
        d["timestamp"] = self.timestamp.isoformat()
        return d


# ── Pipeline stage results ────────────────────────────────────────────────────

@dataclass
class FirewallResult:
    """Output from Stage 2: AI Agent Firewall."""
    passed:           bool       = True
    threat_type:      ThreatType = ThreatType.NONE
    threat_score:     float      = 0.0
    details:          str        = ""
    sanitized_prompt: Optional[str] = None

    def to_dict(self) -> dict:
        return {
            "passed":       self.passed,
            "threat_type":  self.threat_type.value,
            "threat_score": self.threat_score,
            "details":      self.details,
        }


@dataclass
class PolicyResult:
    """Output from Stage 3: Policy Enforcement Gateway."""
    passed:             bool       = True
    violated_policies:  list[str]  = field(default_factory=list)
    required_approvals: list[str]  = field(default_factory=list)
    details:            str        = ""

    def to_dict(self) -> dict:
        return {
            "passed":             self.passed,
            "violated_policies":  self.violated_policies,
            "required_approvals": self.required_approvals,
            "details":            self.details,
        }


@dataclass
class RiskScore:
    """Output from Stage 4: Risk Scoring Engine."""
    total_score:               float     = 0.0
    data_sensitivity_score:    float     = 0.0
    financial_impact_score:    float     = 0.0
    regulatory_exposure_score: float     = 0.0
    policy_violation_score:    float     = 0.0
    risk_level:                RiskLevel = RiskLevel.LOW
    breakdown:                 dict      = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "total_score":               self.total_score,
            "risk_level":                self.risk_level.value,
            "data_sensitivity_score":    self.data_sensitivity_score,
            "financial_impact_score":    self.financial_impact_score,
            "regulatory_exposure_score": self.regulatory_exposure_score,
            "policy_violation_score":    self.policy_violation_score,
            "breakdown":                 self.breakdown,
        }


# ── Provenance graph ──────────────────────────────────────────────────────────

@dataclass
class ProvenanceNode:
    """A single node in the decision provenance DAG."""
    node_id:   str       = field(default_factory=lambda: str(uuid.uuid4())[:8])
    node_type: NodeType  = NodeType.PROMPT
    label:     str       = ""
    data:      dict      = field(default_factory=dict)
    timestamp: datetime  = field(default_factory=datetime.now)


# ── Final output ──────────────────────────────────────────────────────────────

@dataclass
class GovernanceDecision:
    """
    The complete output of a governance evaluation.

    Returned by ``ControlPlane.evaluate()``.  Includes the verdict,
    all stage results, and a full provenance graph.
    """
    action_id:          str                    = ""
    decision:           ExecutionDecision      = ExecutionDecision.ALLOW
    risk_level:         RiskLevel              = RiskLevel.LOW
    risk_score:         float                  = 0.0
    firewall_result:    Optional[FirewallResult]  = None
    policy_result:      Optional[PolicyResult]    = None
    risk_score_detail:  Optional[RiskScore]       = None
    provenance_nodes:   list[ProvenanceNode]      = field(default_factory=list)
    provenance_edges:   list[tuple]               = field(default_factory=list)
    reason:             str                    = ""
    timestamp:          datetime               = field(default_factory=datetime.now)

    # ── Convenience properties ────────────────────────────────────────
    @property
    def allowed(self) -> bool:
        return self.decision == ExecutionDecision.ALLOW

    @property
    def blocked(self) -> bool:
        return self.decision == ExecutionDecision.BLOCK

    @property
    def escalated(self) -> bool:
        return self.decision == ExecutionDecision.ESCALATE

    def to_dict(self) -> dict:
        """Full JSON-serialisable representation for audit export."""
        return {
            "action_id":  self.action_id,
            "timestamp":  self.timestamp.isoformat(),
            "decision":   self.decision.value,
            "risk_level": self.risk_level.value,
            "risk_score": self.risk_score,
            "reason":     self.reason,
            "firewall":   self.firewall_result.to_dict()   if self.firewall_result  else None,
            "policy":     self.policy_result.to_dict()     if self.policy_result    else None,
            "risk":       self.risk_score_detail.to_dict() if self.risk_score_detail else None,
            "provenance": {
                "node_count": len(self.provenance_nodes),
                "edge_count": len(self.provenance_edges),
            },
        }

    def raise_if_blocked(self) -> "GovernanceDecision":
        """
        Raise ``GovernanceViolationError`` if the decision is BLOCK or ESCALATE.
        Useful for strict pipelines where escalation should also halt execution.

        Returns self if ALLOW so callers can chain::

            decision = cp.evaluate(action).raise_if_blocked()
        """
        if self.decision != ExecutionDecision.ALLOW:
            raise GovernanceViolationError(self)
        return self


# ── Exception ─────────────────────────────────────────────────────────────────

class GovernanceViolationError(Exception):
    """
    Raised when an agent action is BLOCKED or ESCALATED by the control plane.

    Attributes:
        decision: The full ``GovernanceDecision`` that triggered the error.

    Example::

        try:
            cp.evaluate(action).raise_if_blocked()
            # ... execute tool ...
        except GovernanceViolationError as e:
            print(e.decision.reason)
            print(e.decision.decision.value)   # "BLOCK" or "ESCALATE"
    """
    def __init__(self, decision: GovernanceDecision):
        self.decision = decision
        super().__init__(
            f"[Eaigins] {decision.decision.value}: {decision.reason}"
        )
