"""W3C PROV-DM export for the Vaara audit trail.

Produces PROV-JSON (W3C Submission, 2013) so any PROV-aware consumer can
ingest a Vaara audit trail without a Vaara-specific adapter. Two layers:

1. **Per-action bundle.** Full provenance of one action's lifecycle.
   Each AuditRecord becomes a PROV Activity. Tool requests, scores,
   decisions, and outcomes become Entities. The AI agent, the Vaara
   pipeline, and any human reviewer become Agents. Regulatory articles
   surface as ``aiact:satisfies`` / ``dora:satisfies`` attributes.

2. **Audit-record chain layer.** Sequence of record-level Entities,
   each ``wasDerivedFrom`` its predecessor. PROV expresses chain
   structure. Cryptographic verification stays Vaara's job.

Zero new runtime deps. See ``research/prov_dm_mapping.md``.
"""

from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Any, Iterable, Optional

from vaara.audit.trail import AuditRecord, EventType

PROV_CONTEXT = "https://www.w3.org/ns/prov-context"
VAARA_NS = "https://vaara.io/prov/"
AIACT_NS = "https://eur-lex.europa.eu/eli/reg/2024/1689/oj#"
DORA_NS = "https://eur-lex.europa.eu/eli/reg/2022/2554/oj#"

_PIPELINE_AGENT = "vaara:agent/vaara-pipeline"
_PIPELINE_EVENTS = frozenset({
    EventType.RISK_SCORED, EventType.DECISION_MADE, EventType.ACTION_BLOCKED,
    EventType.ACTION_EXECUTED, EventType.ESCALATION_SENT, EventType.OUTCOME_RECORDED,
})


def _ts(record: AuditRecord) -> str:
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(record.timestamp))


def _articles_attr(record: AuditRecord) -> dict:
    by_domain: dict[str, list[str]] = {}
    for art in record.regulatory_articles or []:
        if not isinstance(art, dict):
            continue
        article = art.get("article", "")
        if not article:
            continue
        by_domain.setdefault(art.get("domain", ""), []).append(article)
    out: dict[str, list[str]] = {}
    if "EU AI Act" in by_domain:
        out["aiact:satisfies"] = sorted(set(by_domain["EU AI Act"]))
    if "DORA" in by_domain:
        out["dora:satisfies"] = sorted(set(by_domain["DORA"]))
    return out


def _entity_for(rec: AuditRecord, action_id: str) -> tuple[Optional[str], dict]:
    """Pick the Entity (if any) generated by this event."""
    data = rec.data or {}
    et = rec.event_type
    if et == EventType.ACTION_REQUESTED:
        return f"vaara:request/{action_id}", {
            "vaara:toolName": rec.tool_name,
            "vaara:parameters": data.get("parameters", {}),
            "vaara:agentConfidence": data.get("agent_confidence"),
        }
    if et == EventType.RISK_SCORED:
        attrs: dict[str, Any] = {}
        score = data.get("point_estimate", data.get("risk_score"))
        if score is not None:
            attrs["vaara:riskScore"] = score
        if "conformal_lower" in data and "conformal_upper" in data:
            attrs["vaara:conformalInterval"] = [
                data["conformal_lower"], data["conformal_upper"],
            ]
        # Conformal calibration context. Surfacing these on the score Entity
        # lets an external auditor distinguish a wide interval from a cold
        # calibrator (low calibration_size) versus one from genuine high
        # uncertainty (large effective_alpha after FACI drift), and tells
        # them which Mondrian bucket produced the interval when in
        # class-conditional mode.
        cal_size = data.get("calibration_size")
        if cal_size is not None:
            attrs["vaara:calibrationSize"] = cal_size
        eff_alpha = data.get("effective_alpha")
        if eff_alpha is not None:
            attrs["vaara:effectiveAlpha"] = eff_alpha
        bucket = data.get("bucket_category")
        if bucket:
            attrs["vaara:bucketCategory"] = bucket
        return f"vaara:score/{action_id}", attrs
    if et in (EventType.DECISION_MADE, EventType.ACTION_BLOCKED):
        default_dec = "deny" if et == EventType.ACTION_BLOCKED else "allow"
        return f"vaara:decision/{action_id}/{rec.record_id}", {
            "vaara:decisionValue": data.get("decision", default_dec),
            "vaara:reason": data.get("reason", ""),
            "vaara:riskScore": data.get("risk_score"),
        }
    if et == EventType.ESCALATION_RESOLVED:
        return f"vaara:operatorResponse/{action_id}/{rec.record_id}", {
            "vaara:responseValue": data.get("resolution", ""),
            "vaara:justification": data.get("justification", ""),
        }
    if et == EventType.OUTCOME_RECORDED:
        return f"vaara:outcome/{action_id}/{rec.record_id}", {
            "vaara:outcomeSeverity": data.get("outcome_severity"),
        }
    if et == EventType.POLICY_OVERRIDE:
        return f"vaara:decision/{action_id}/{rec.record_id}", {
            "vaara:decisionValue": data.get("new_decision", ""),
            "vaara:reason": data.get("override_reason", ""),
        }
    return None, {}


def _per_action_bundle(records: list[AuditRecord]) -> dict[str, Any]:
    """Build one PROV bundle for a single action's lifecycle."""
    if not records:
        return {}
    action_id = records[0].action_id
    agents: dict[str, dict] = {_PIPELINE_AGENT: {"prov:type": "prov:SoftwareAgent"}}
    entities: dict[str, dict] = {}
    activities: dict[str, dict] = {}
    wgb: dict[str, dict] = {}
    used: dict[str, dict] = {}
    wib: dict[str, dict] = {}
    waw: dict[str, dict] = {}
    wat: dict[str, dict] = {}
    revisions: dict[str, dict] = {}

    counter = [0]
    def rid(prefix: str) -> str:
        counter[0] += 1
        return f"_:{prefix}{counter[0]}"

    prev_act: Optional[str] = None
    last_decision: Optional[str] = None

    for rec in records:
        act_id = f"vaara:event/{rec.record_id}"
        ai_agent = f"vaara:agent/{rec.agent_id}"
        agents.setdefault(ai_agent, {"prov:type": "prov:SoftwareAgent"})

        attrs: dict[str, Any] = {
            "prov:type": f"vaara:{rec.event_type.value}",
            "prov:startTime": _ts(rec),
        }
        attrs.update(_articles_attr(rec))
        activities[act_id] = attrs

        if prev_act is not None:
            wib[rid("wib")] = {"prov:informed": act_id, "prov:informant": prev_act}

        if rec.event_type == EventType.ACTION_REQUESTED:
            waw[rid("wa")] = {"prov:activity": act_id, "prov:agent": ai_agent}
        elif rec.event_type in _PIPELINE_EVENTS:
            waw[rid("wa")] = {"prov:activity": act_id, "prov:agent": _PIPELINE_AGENT}
        elif rec.event_type in (EventType.ESCALATION_RESOLVED, EventType.POLICY_OVERRIDE):
            who = (rec.data or {}).get("reviewer") or (rec.data or {}).get("overrider")
            if who:
                op_agent = f"vaara:agent/{who}"
                agents.setdefault(op_agent, {"prov:type": "prov:Person"})
                waw[rid("wa")] = {"prov:activity": act_id, "prov:agent": op_agent}

        ent_id, ent_attrs = _entity_for(rec, action_id)
        if ent_id is not None:
            entities[ent_id] = ent_attrs
            wgb[rid("wgb")] = {"prov:entity": ent_id, "prov:activity": act_id}
            if rec.event_type in (EventType.DECISION_MADE, EventType.ACTION_BLOCKED):
                last_decision = ent_id
            if rec.event_type == EventType.ESCALATION_RESOLVED:
                reviewer = (rec.data or {}).get("reviewer")
                if reviewer:
                    wat[rid("wat")] = {
                        "prov:entity": ent_id,
                        "prov:agent": f"vaara:agent/{reviewer}",
                    }
            if rec.event_type == EventType.POLICY_OVERRIDE and last_decision:
                revisions[rid("wrv")] = {
                    "prov:generatedEntity": ent_id,
                    "prov:usedEntity": last_decision,
                    "prov:type": "prov:Revision",
                }
                last_decision = ent_id

        if rec.event_type == EventType.RISK_SCORED:
            used[rid("u")] = {
                "prov:activity": act_id,
                "prov:entity": f"vaara:request/{action_id}",
            }
        elif rec.event_type in (EventType.DECISION_MADE, EventType.ACTION_BLOCKED):
            used[rid("u")] = {
                "prov:activity": act_id,
                "prov:entity": f"vaara:score/{action_id}",
            }

        prev_act = act_id

    bundle: dict[str, Any] = {}
    for key, val in (
        ("agent", agents), ("entity", entities), ("activity", activities),
        ("wasGeneratedBy", wgb), ("used", used), ("wasInformedBy", wib),
        ("wasAssociatedWith", waw), ("wasAttributedTo", wat),
        ("wasDerivedFrom", revisions),
    ):
        if val:
            bundle[key] = val
    return bundle


def _chain_layer(records: list[AuditRecord]) -> tuple[dict, dict]:
    """Build the per-record audit chain.

    Each record becomes a Bundle-typed Entity carrying its hash. Edges are
    derived from ``previous_hash`` rather than iteration order so a filtered
    or partially loaded slice never asserts false lineage. Records whose
    ``previous_hash`` is empty or missing from the slice produce no edge.
    """
    entities: dict[str, dict] = {}
    rels: dict[str, dict] = {}
    by_hash: dict[str, str] = {}
    for rec in records:
        ent_id = f"vaara:auditRecord/{rec.record_id}"
        entities[ent_id] = {
            "prov:type": "prov:Bundle",
            "vaara:eventType": rec.event_type.value,
            "vaara:actionId": rec.action_id,
            "vaara:recordHash": rec.record_hash,
            "vaara:previousHash": rec.previous_hash,
        }
        if rec.record_hash:
            by_hash[rec.record_hash] = ent_id
    counter = 0
    for rec in records:
        if rec.previous_hash and rec.previous_hash in by_hash:
            counter += 1
            rels[f"_:chain{counter}"] = {
                "prov:generatedEntity": f"vaara:auditRecord/{rec.record_id}",
                "prov:usedEntity": by_hash[rec.previous_hash],
                "prov:type": "prov:Revision",
            }
    return entities, rels


def audit_to_prov_json(
    records: Iterable[AuditRecord],
    action_id: Optional[str] = None,
    include_chain: bool = True,
) -> dict[str, Any]:
    """Serialise audit records to a PROV-JSON document.

    Args:
        records: Iterable of AuditRecord. Within an action, chronological
            order is required; cross-action ordering does not matter.
        action_id: If given, only that action's bundle is emitted.
        include_chain: If True, append the audit-record chain layer as a
            top-level bundle named ``vaara:auditChain``.
    """
    rec_list = list(records)
    if action_id is not None:
        rec_list = [r for r in rec_list if r.action_id == action_id]

    by_action: dict[str, list[AuditRecord]] = {}
    for r in rec_list:
        by_action.setdefault(r.action_id, []).append(r)

    bundles: dict[str, Any] = {}
    for aid, recs in by_action.items():
        recs_sorted = sorted(recs, key=lambda r: r.timestamp)
        bundle = _per_action_bundle(recs_sorted)
        if bundle:
            bundles[f"vaara:action/{aid}"] = bundle

    doc: dict[str, Any] = {
        "@context": PROV_CONTEXT,
        "prefix": {"vaara": VAARA_NS, "aiact": AIACT_NS, "dora": DORA_NS},
    }
    if bundles:
        doc["bundle"] = bundles

    if include_chain:
        chain_entities, chain_rels = _chain_layer(rec_list)
        if chain_entities:
            chain_bundle: dict[str, Any] = {"entity": chain_entities}
            if chain_rels:
                chain_bundle["wasDerivedFrom"] = chain_rels
            doc.setdefault("bundle", {})["vaara:auditChain"] = chain_bundle

    return doc


def write_prov_json(
    records: Iterable[AuditRecord],
    out_path: Path,
    action_id: Optional[str] = None,
    include_chain: bool = True,
) -> int:
    """Write PROV-JSON to ``out_path``. Returns record count emitted."""
    rec_list = list(records)
    if action_id is not None:
        rec_list = [r for r in rec_list if r.action_id == action_id]
    doc = audit_to_prov_json(rec_list, action_id=action_id, include_chain=include_chain)
    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(doc, indent=2, default=str), encoding="utf-8")
    return len(rec_list)
