"""
Risk Quantification Agent — ZIA reference implementation.

Reads asset + vulnerability data from the SIF graph, computes risk scores
using the RiskPropagationEngine, and writes RiskNodes back to the graph.

This agent wraps the QBER risk quantification logic as a ZIA BaseAgent.
The QBER engine (PyMC-based probabilistic modeling) will be integrated in Phase 3.
"""

from __future__ import annotations

from uuid import uuid4

from zia.core.runtime.agent import AgentContext, AgentResult, BaseAgent
from zia.core.runtime.registry import register_agent
from zia.sif.risk.propagation import RiskInputs, RiskPropagationEngine
from zia.sif.schema.nodes import RiskNode

try:
    from zia.sif.graph.adapter import KuzuAdapter
except Exception:
    KuzuAdapter = None  # type: ignore[assignment,misc]


@register_agent(
    domain="risk_quant",
    description="Computes risk scores for all assets using the ZIA risk propagation engine",
    version="1.0.0",
)
class RiskQuantAgent(BaseAgent):
    """
    Computes and persists risk scores for all assets in a tenant's SIF graph.

    For each asset:
    1. Loads linked vulnerabilities and controls from the graph
    2. Computes node_risk using RiskPropagationEngine
    3. Writes a RiskNode back to the SIF

    QBER integration point: replace _compute_eal() stub with actual PyMC model.
    """

    def __init__(self, adapter: object | None = None) -> None:
        self._adapter = adapter
        self._engine = RiskPropagationEngine()

    def execute(self, context: AgentContext) -> AgentResult:
        tenant_id = context.tenant_id
        scored: list[dict] = []
        errors: list[str] = []

        # Load all assets for this tenant (graceful when graph is unavailable)
        assets = (
            self._adapter.get_nodes(tenant_id=tenant_id, node_type="asset")
            if self._adapter is not None else []
        )

        for asset in assets:
            try:
                risk_output = self._score_asset(asset, tenant_id)
                risk_node = RiskNode(
                    node_id=f"risk-{asset['node_id']}",
                    risk_type="cyber",
                    likelihood=risk_output.raw_score,
                    impact=float(asset.get("criticality_score", 5.0)),
                    risk_score=risk_output.risk_score,
                    eal=self._compute_eal_stub(risk_output.risk_score),
                    source=context.agent_id,
                )
                if self._adapter is not None:
                    self._adapter.upsert_node(tenant_id, risk_node)
                scored.append({
                    "asset_id": asset["node_id"],
                    "risk_score": risk_output.risk_score,
                    "risk_level": risk_output.risk_level,
                })
            except Exception as e:
                errors.append(f"Failed to score asset {asset.get('node_id')}: {e}")

        if errors and not scored:
            return AgentResult.fail(context, errors=errors)

        return AgentResult.ok(
            context,
            output={
                "assets_scored": len(scored),
                "results": scored,
                "errors": errors,
            },
        )

    def _score_asset(self, asset: dict, tenant_id: str) -> object:
        """Compute risk for a single asset dict."""
        criticality = asset.get("criticality", "medium")
        exposure = asset.get("exposure_level", "internal")

        # Load worst-case vulnerability exploitability for this asset
        vulns = self._adapter.get_nodes(tenant_id=tenant_id, node_type="vulnerability") if self._adapter else []
        max_exploitability = max(
            (float(v.get("exploitability", 0.5)) for v in vulns), default=0.5
        )

        # Load best control effectiveness
        controls = self._adapter.get_nodes(tenant_id=tenant_id, node_type="control") if self._adapter else []
        max_control_eff = max(
            (float(c.get("effectiveness", 0.5)) for c in controls), default=0.0
        )

        inputs = RiskInputs(
            base_risk=RiskPropagationEngine.criticality_to_base_risk(criticality),
            exposure_factor=RiskPropagationEngine.exposure_to_factor(exposure),
            exploitability=max_exploitability,
            control_effectiveness=max_control_eff,
            privilege_amplifier=1.0,  # Will be computed from IdentityNodes in Phase 3
        )
        return RiskPropagationEngine.compute(inputs)

    def _compute_eal_stub(self, risk_score: float) -> float:
        """
        Stub for Expected Annual Loss computation.
        TODO: Replace with QBER PyMC probabilistic model in Phase 3.
        """
        # Simple placeholder: EAL scales with risk score (in USD thousands)
        return round(risk_score * 50_000, 2)
