"""
eaigins_sdk.provenance
──────────────────────
Stage 5: Decision Provenance Engine  (Section 5 of the paper)

Records every AI decision as a directed acyclic graph (DAG), not a flat log.
Each node is a pipeline stage; each edge is a causal relationship.

    Prompt → Context(s) → Reasoning → Firewall → Policy → Risk → Decision

Enables: audit trails, incident investigation, regulatory compliance,
model debugging, governance policy verification.

Export formats: JSON (audit trail), NetworkX DiGraph (visualisation).
"""

from __future__ import annotations

import json
from datetime import datetime
from typing import Optional
from .models import (
    AgentAction, FirewallResult, PolicyResult,
    RiskScore, ExecutionDecision, GovernanceDecision,
    ProvenanceNode, NodeType, RiskLevel,
)


class ProvenanceEngine:
    """
    Builds and stores decision provenance graphs.

    Usage::

        engine = ProvenanceEngine()
        decision = engine.record(
            action, firewall_result, policy_result,
            risk_score, ExecutionDecision.BLOCK, "reason..."
        )

        # Export full audit trail as JSON
        trail = engine.export_audit_trail()

        # Get a single decision
        d = engine.get(action.action_id)

        # Render as NetworkX graph (requires networkx)
        G = engine.to_networkx(action.action_id)
    """

    def __init__(self):
        self._store: dict[str, GovernanceDecision] = {}

    # ── Public API ────────────────────────────────────────────────────────────

    def record(
        self,
        action:          AgentAction,
        firewall_result: FirewallResult,
        policy_result:   PolicyResult,
        risk_score:      RiskScore,
        final_decision:  ExecutionDecision,
        reason:          str,
    ) -> GovernanceDecision:
        """Build and persist a complete decision provenance record."""
        nodes: list[ProvenanceNode] = []
        edges: list[tuple]          = []

        # [1] Prompt node
        prompt_node = ProvenanceNode(
            node_type=NodeType.PROMPT,
            label=f"Prompt\n[{action.agent_id}]",
            data={
                "prompt_preview": action.prompt[:120] + "…" if len(action.prompt) > 120 else action.prompt,
                "agent_id":       action.agent_id,
                "timestamp":      action.timestamp.isoformat(),
            },
        )
        nodes.append(prompt_node)

        # [2] Context source nodes
        ctx_nodes: list[ProvenanceNode] = []
        for src in action.context_sources:
            ctx = ProvenanceNode(
                node_type=NodeType.CONTEXT_SOURCE,
                label=f"Context\n{src}",
                data={"source": src, "sensitivity": self._sensitivity_label(src)},
            )
            nodes.append(ctx)
            ctx_nodes.append(ctx)
            edges.append((prompt_node.node_id, ctx.node_id, "retrieves"))

        # [3] Reasoning node
        reason_node = ProvenanceNode(
            node_type=NodeType.REASONING,
            label=f"AI Reasoning\n[{action.tool_name}]",
            data={
                "tool":             action.tool_name,
                "tool_params":      action.tool_params,
                "reasoning_preview": action.reasoning[:100] + "…" if len(action.reasoning) > 100 else action.reasoning,
            },
        )
        nodes.append(reason_node)
        for ctx in ctx_nodes:
            edges.append((ctx.node_id, reason_node.node_id, "informs"))
        if not ctx_nodes:
            edges.append((prompt_node.node_id, reason_node.node_id, "produces"))

        # [4] Firewall node
        fw_status = "✓ PASS" if firewall_result.passed else f"✗ BLOCKED\n{firewall_result.threat_type.value}"
        fw_node = ProvenanceNode(
            node_type=NodeType.POLICY_EVAL,
            label=f"Firewall\n{fw_status}",
            data={
                "passed":       firewall_result.passed,
                "threat_type":  firewall_result.threat_type.value,
                "threat_score": firewall_result.threat_score,
                "details":      firewall_result.details,
            },
        )
        nodes.append(fw_node)
        edges.append((reason_node.node_id, fw_node.node_id, "inspected_by"))

        # [5] Policy node
        pol_status = "✓ PASS" if policy_result.passed else f"✗ {len(policy_result.violated_policies)} violations"
        pol_node = ProvenanceNode(
            node_type=NodeType.POLICY_EVAL,
            label=f"Policy Gate\n{pol_status}",
            data={
                "passed":             policy_result.passed,
                "violated":           policy_result.violated_policies,
                "approvals_required": policy_result.required_approvals,
                "details":            policy_result.details[:200],
            },
        )
        nodes.append(pol_node)
        edges.append((fw_node.node_id, pol_node.node_id, "evaluated_by"))

        # [6] Risk score node
        risk_node = ProvenanceNode(
            node_type=NodeType.RISK_SCORE,
            label=f"Risk Score\n{risk_score.total_score:.2f} [{risk_score.risk_level.value}]",
            data={
                "total_score": risk_score.total_score,
                "risk_level":  risk_score.risk_level.value,
                "breakdown":   risk_score.breakdown,
            },
        )
        nodes.append(risk_node)
        edges.append((pol_node.node_id, risk_node.node_id, "risk_evaluated"))

        # [7] Final decision node
        icons  = {"ALLOW": "✅", "ESCALATE": "⚠️", "BLOCK": "🛑"}
        dec_node = ProvenanceNode(
            node_type=NodeType.ACTION if final_decision == ExecutionDecision.ALLOW else NodeType.BLOCK,
            label=f"{icons.get(final_decision.value, '')} {final_decision.value}\n{action.tool_name}",
            data={
                "decision":  final_decision.value,
                "tool_name": action.tool_name,
                "reason":    reason,
                "timestamp": datetime.now().isoformat(),
            },
        )
        nodes.append(dec_node)
        edges.append((risk_node.node_id, dec_node.node_id, "decides"))

        decision = GovernanceDecision(
            action_id         = action.action_id,
            decision          = final_decision,
            risk_level        = risk_score.risk_level,
            risk_score        = risk_score.total_score,
            firewall_result   = firewall_result,
            policy_result     = policy_result,
            risk_score_detail = risk_score,
            provenance_nodes  = nodes,
            provenance_edges  = edges,
            reason            = reason,
        )
        self._store[action.action_id] = decision
        return decision

    def get(self, action_id: str) -> Optional[GovernanceDecision]:
        """Retrieve a single decision by action ID."""
        return self._store.get(action_id)

    def get_all(self) -> list[GovernanceDecision]:
        """Return all recorded decisions."""
        return list(self._store.values())

    def clear(self) -> None:
        """Clear all stored decisions."""
        self._store.clear()

    def export_audit_trail(self) -> list[dict]:
        """
        Export all decisions as a JSON-serialisable audit trail.
        Each record is suitable for regulatory review or SIEM ingestion.
        """
        return [
            {
                "action_id":           d.action_id,
                "timestamp":           d.timestamp.isoformat(),
                "decision":            d.decision.value,
                "risk_level":          d.risk_level.value,
                "risk_score":          d.risk_score,
                "firewall_passed":     d.firewall_result.passed if d.firewall_result else None,
                "firewall_threat":     d.firewall_result.threat_type.value if d.firewall_result else None,
                "policy_passed":       d.policy_result.passed if d.policy_result else None,
                "violated_policies":   d.policy_result.violated_policies if d.policy_result else [],
                "required_approvals":  d.policy_result.required_approvals if d.policy_result else [],
                "reason":              d.reason,
                "provenance_nodes":    len(d.provenance_nodes),
                "provenance_edges":    len(d.provenance_edges),
            }
            for d in self._store.values()
        ]

    def to_networkx(self, action_id: str):
        """
        Return a NetworkX DiGraph for the given decision (requires networkx).

        Example::

            import networkx as nx
            G = engine.to_networkx(action_id)
            # G.nodes / G.edges are now available for visualisation
        """
        try:
            import networkx as nx
        except ImportError:
            raise ImportError("networkx is required: pip install networkx")

        decision = self._store.get(action_id)
        if decision is None:
            raise KeyError(f"No decision found for action_id={action_id!r}")

        G = nx.DiGraph()
        for node in decision.provenance_nodes:
            G.add_node(node.node_id, label=node.label,
                       node_type=node.node_type.value, data=node.data)
        for src, dst, label in decision.provenance_edges:
            G.add_edge(src, dst, label=label)
        return G

    def export_json(self, action_id: str) -> str:
        """Return pretty-printed JSON of a single decision's provenance."""
        d = self._store.get(action_id)
        if d is None:
            raise KeyError(action_id)
        return json.dumps(d.to_dict(), indent=2)

    # ── Private ───────────────────────────────────────────────────────────────

    @staticmethod
    def _sensitivity_label(source: str) -> str:
        s = source.lower()
        if any(w in s for w in ["transaction", "customer", "account", "balance",
                                 "wire", "patient", "ehr", "psychiatric"]):
            return "HIGH"
        if any(w in s for w in ["fraud", "aml", "investigation", "compliance",
                                 "clinical", "lab", "imaging"]):
            return "MEDIUM"
        return "LOW"

    def __len__(self) -> int:
        return len(self._store)

    def __repr__(self) -> str:
        return f"ProvenanceEngine({len(self._store)} decisions)"
