"""Causal compliance forecasting — Gap 6 from the Wolfram Completion Plan.

Uses temporal causal cones (from propagation analysis) and entanglement
(from the entanglement module) to answer:

1. **Confidence decay**: For each compliance control, how confident are we
   that the last-evaluated status is still accurate?  Recent code/graph
   changes whose causal cones intersect a control's evidence set erode
   confidence.

2. **Forecast from planned changes**: Given a set of planned feature changes,
   predict which controls will need re-evaluation and which are at risk of
   degradation.

3. **Entanglement-aware co-degradation**: Group controls by shared evidence
   features; when one control degrades, predict which others are likely to
   follow based on entanglement scores.
"""
from __future__ import annotations

import json
import logging
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

log = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# 1. Confidence decay
# ---------------------------------------------------------------------------

def compute_confidence_decay(
    scan_report: dict,
    recent_changes: Optional[List[str]] = None,
) -> Dict[str, Any]:
    """For each compliance control, compute confidence that its status is current.

    Uses causal cones of recent changes to determine which controls'
    evidence sets may have been affected.

    Parameters
    ----------
    scan_report : dict
        A combined-posture or scan report with ``articles[].controls[]``.
    recent_changes : list of str or None
        Feature IDs that changed since the last scan.  If *None*, attempts
        to infer from the graph (features modified after the scan timestamp).

    Returns
    -------
    dict
        {controls, summary: {fully_confident, uncertain, likely_stale}}
    """
    from graph.compliance_mapping import (
        _EVIDENCE_CLASSIFICATION,
        FRAMEWORK_MAPPINGS,
    )

    # Build control list from the scan report
    controls_out = []
    control_fns = _extract_control_evidence_fns(scan_report)

    # Compute causal cones for each changed feature
    change_cones = {}  # type: Dict[str, set]
    if recent_changes:
        change_cones = _compute_change_cones(recent_changes)

    # Get the scan timestamp for staleness baseline
    scan_ts = scan_report.get("generatedAt", "")

    # Map evidence functions → node types they depend on (from classification)
    fn_node_types = _build_fn_node_type_map()

    fully_confident = 0
    uncertain = 0
    likely_stale = 0

    for ctrl_id, ctrl_info in control_fns.items():
        fn_name = ctrl_info.get("evidence_fn_name", "")
        current_status = ctrl_info.get("status", "gap")

        classification = _EVIDENCE_CLASSIFICATION.get(fn_name, {})
        requires_graph = classification.get("requires_graph", False)

        # Determine which change cones intersect this control's evidence
        affected_by = []
        if requires_graph and change_cones:
            # Graph-required fns depend on features in the graph.
            # Check if any change cone reaches node types this fn uses.
            fn_deps = fn_node_types.get(fn_name, set())
            for change_id, cone_nodes in change_cones.items():
                # Cone nodes are feature IDs — graph-required fns implicitly
                # depend on the feature subgraph
                if cone_nodes & fn_deps or not fn_deps:
                    # No explicit deps or overlap → potentially affected
                    affected_by.append(change_id)

        # Confidence model:
        #   1.0 = no intersecting changes → fully confident
        #   0.5 = some intersecting changes → uncertain
        #   0.0 = many intersecting changes → likely stale
        if not recent_changes or not requires_graph:
            # Non-graph fns are unaffected by feature changes;
            # no recent changes means nothing to decay
            confidence = 1.0
        elif len(affected_by) == 0:
            confidence = 1.0
        else:
            # Decay proportional to fraction of changes that affect this control
            ratio = len(affected_by) / max(len(recent_changes), 1)
            confidence = round(max(0.0, 1.0 - ratio), 4)

        # Classify
        if confidence >= 0.8:
            fully_confident += 1
        elif confidence >= 0.4:
            uncertain += 1
        else:
            likely_stale += 1

        controls_out.append({
            "controlId": ctrl_id,
            "currentStatus": current_status,
            "evidenceFn": fn_name,
            "lastEvaluated": scan_ts,
            "confidence": confidence,
            "staleSince": None if confidence >= 0.8 else scan_ts,
            "affectedBy": affected_by,
        })

    return {
        "controls": controls_out,
        "summary": {
            "fully_confident": fully_confident,
            "uncertain": uncertain,
            "likely_stale": likely_stale,
            "total": len(controls_out),
        },
    }


# ---------------------------------------------------------------------------
# 2. Forecast from planned changes
# ---------------------------------------------------------------------------

def forecast_from_planned_changes(
    planned_feature_ids: List[str],
    current_report: dict,
) -> Dict[str, Any]:
    """Predict which controls will need re-evaluation after planned changes.

    Parameters
    ----------
    planned_feature_ids : list of str
        Feature node IDs for planned/upcoming changes.
    current_report : dict
        Current scan/posture report with ``articles[].controls[]``.

    Returns
    -------
    dict
        {planned_features, at_risk_controls, combined_cone_size,
         entanglement_clusters_hit}
    """
    from graph.compliance_mapping import _EVIDENCE_CLASSIFICATION

    # Compute combined wavefront from all planned features
    combined_cone = set()  # type: set
    per_feature_cones = {}  # type: Dict[str, set]
    timescale_arrivals = {}  # type: Dict[str, str]

    for fid in planned_feature_ids:
        try:
            from graph.propagation_analysis import compute_wavefront
            wf = compute_wavefront(fid)
            cone = set()
            for ts_key in ("machine", "human", "business"):
                for node in wf.get(ts_key, []):
                    nid = node["nodeId"]
                    cone.add(nid)
                    # Track the fastest arrival timescale per node
                    if nid not in timescale_arrivals:
                        timescale_arrivals[nid] = node["timescale"]
            per_feature_cones[fid] = cone
            combined_cone |= cone
        except Exception as exc:
            log.warning("Wavefront failed for %s: %s", fid, exc)
            per_feature_cones[fid] = set()

    # Build control → evidence fn map from the report
    control_fns = _extract_control_evidence_fns(current_report)
    fn_node_types = _build_fn_node_type_map()

    # Check entanglement clusters
    entanglement_clusters_hit = 0
    try:
        from graph.entanglement import compute_entanglement_matrix
        ent = compute_entanglement_matrix(feature_ids=planned_feature_ids)
        entanglement_clusters_hit = len(ent.get("clusters", []))
    except Exception as e:
        log.debug("Entanglement matrix unavailable: %s", e)

    # For each control, check if the combined cone intersects its evidence deps
    at_risk = []
    _STATUS_SEVERITY = {"met": 0, "partial": 1, "gap": 2}

    for ctrl_id, ctrl_info in control_fns.items():
        fn_name = ctrl_info.get("evidence_fn_name", "")
        current_status = ctrl_info.get("status", "gap")
        classification = _EVIDENCE_CLASSIFICATION.get(fn_name, {})

        if not classification.get("requires_graph", False):
            continue  # Non-graph fns unaffected by feature changes

        fn_deps = fn_node_types.get(fn_name, set())

        # Check intersection: any planned feature's cone reaches this fn's deps
        intersecting_features = []
        for fid in planned_feature_ids:
            cone = per_feature_cones.get(fid, set())
            if cone & fn_deps or (not fn_deps and cone):
                intersecting_features.append(fid)

        if not intersecting_features:
            continue

        # Risk level based on current status margin
        sev = _STATUS_SEVERITY.get(current_status, 2)
        if sev >= 2:
            risk = "high"      # Already a gap — change can't help
            margin = 0.0
        elif sev == 1:
            risk = "high"      # Partial — one more issue tips to gap
            margin = 0.3
        else:
            risk = "medium"    # Met — has some buffer
            margin = 0.7

        # Arrival timescale: fastest path from any planned feature to this control
        arrival_ts = "business"  # default to slowest
        for fid in intersecting_features:
            cone = per_feature_cones.get(fid, set())
            for nid in cone & fn_deps:
                ts = timescale_arrivals.get(nid, "business")
                if ts == "machine":
                    arrival_ts = "machine"
                    break
                elif ts == "human" and arrival_ts == "business":
                    arrival_ts = "human"

        at_risk.append({
            "controlId": ctrl_id,
            "currentStatus": current_status,
            "riskLevel": risk,
            "arrivalTimescale": arrival_ts,
            "evidenceMargin": round(margin, 2),
            "intersectingFeatures": intersecting_features,
        })

    # Sort by risk level (high first)
    risk_order = {"high": 0, "medium": 1, "low": 2}
    at_risk.sort(key=lambda c: risk_order.get(c["riskLevel"], 3))

    return {
        "planned_features": planned_feature_ids,
        "at_risk_controls": at_risk,
        "combined_cone_size": len(combined_cone),
        "entanglement_clusters_hit": entanglement_clusters_hit,
    }


# ---------------------------------------------------------------------------
# 3. Entanglement-aware co-degradation
# ---------------------------------------------------------------------------

def forecast_with_entanglement(
    current_report: dict,
    entanglement_matrix: Optional[dict] = None,
) -> Dict[str, Any]:
    """Predict which control clusters will degrade together.

    Groups controls by shared evidence function overlap.  For controls
    sharing evidence from entangled features, compute
    P(B degrades | A degrades) from the entanglement score.

    Parameters
    ----------
    current_report : dict
        Current scan/posture report.
    entanglement_matrix : dict or None
        Pre-computed entanglement matrix.  If *None*, computes from graph.

    Returns
    -------
    dict
        {clusters, independent_controls}
    """
    from graph.compliance_mapping import _EVIDENCE_CLASSIFICATION

    # Get entanglement matrix if not provided
    if entanglement_matrix is None:
        try:
            from graph.entanglement import compute_entanglement_matrix
            entanglement_matrix = compute_entanglement_matrix()
        except Exception as e:
            log.debug("Failed to compute entanglement matrix: %s", e)
            entanglement_matrix = {"features": [], "matrix": [], "clusters": []}

    # Build control → evidence fn map
    control_fns = _extract_control_evidence_fns(current_report)

    # Group controls by evidence function (controls sharing the same fn
    # will always co-degrade by definition)
    fn_to_controls = {}  # type: Dict[str, List[str]]
    for ctrl_id, info in control_fns.items():
        fn_name = info.get("evidence_fn_name", "")
        if fn_name:
            fn_to_controls.setdefault(fn_name, []).append(ctrl_id)

    # Build clusters: controls linked through shared evidence fns AND
    # entangled feature dependencies
    ent_features = entanglement_matrix.get("features", [])
    ent_mat = entanglement_matrix.get("matrix", [])
    ent_clusters = entanglement_matrix.get("clusters", [])

    # Map feature index for matrix lookups
    feat_idx = {fid: i for i, fid in enumerate(ent_features)}

    # For each pair of evidence fns that are graph-required,
    # check if they depend on entangled features
    graph_fns = [
        fn for fn, cls in _EVIDENCE_CLASSIFICATION.items()
        if cls.get("requires_graph", False)
    ]

    # Build fn → feature dependencies (features whose changes affect this fn)
    fn_features = {}  # type: Dict[str, set]
    fn_node_types = _build_fn_node_type_map()
    for fn_name in graph_fns:
        # All features are potential deps for graph-required fns
        fn_features[fn_name] = set(ent_features)

    # Find control clusters via shared entangled features
    clusters_out = []
    used_controls = set()  # type: set

    # First: controls sharing the same evidence fn are trivially clustered
    for fn_name, ctrls in fn_to_controls.items():
        if len(ctrls) > 1:
            cluster_ids = sorted(ctrls)
            clusters_out.append({
                "controls": cluster_ids,
                "shared_entanglement": 1.0,  # same fn → perfect entanglement
                "coupling_type": "shared_evidence_fn",
                "if_any_degrades_probability": 1.0,
            })
            used_controls.update(cluster_ids)

    # Second: controls with different evidence fns but entangled feature deps
    # Use entanglement clusters to find feature groups, then map back to controls
    for ent_cluster in ent_clusters:
        cluster_features = set(ent_cluster.get("features", []))
        avg_ent = ent_cluster.get("avg_entanglement", 0.0)

        if avg_ent < 0.3:
            continue  # Only significant entanglement

        # Find controls whose evidence fns are graph-required and touch
        # features in this entanglement cluster
        cluster_controls = set()
        for ctrl_id, info in control_fns.items():
            fn_name = info.get("evidence_fn_name", "")
            if fn_name in graph_fns:
                # Graph-required fn may be affected by features in this cluster
                cluster_controls.add(ctrl_id)

        if len(cluster_controls) > 1:
            sorted_ctrls = sorted(cluster_controls)
            # P(B degrades | A degrades) ≈ avg entanglement of the cluster
            clusters_out.append({
                "controls": sorted_ctrls,
                "shared_entanglement": round(avg_ent, 4),
                "coupling_type": "entangled_features",
                "if_any_degrades_probability": round(min(1.0, avg_ent * 1.5), 4),
            })
            used_controls.update(sorted_ctrls)

    # Independent controls: not in any cluster
    all_ctrl_ids = set(control_fns.keys())
    independent = sorted(all_ctrl_ids - used_controls)

    return {
        "clusters": clusters_out,
        "independent_controls": independent,
    }


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _extract_control_evidence_fns(report: dict) -> Dict[str, Dict[str, str]]:
    """Extract {controlId: {evidence_fn_name, status}} from a report.

    Works with both combined-posture reports and raw scan reports.
    """
    controls = {}
    report_data = report
    if isinstance(report_data.get("report"), str):
        try:
            report_data = json.loads(report_data["report"])
        except (json.JSONDecodeError, TypeError):
            pass
    elif isinstance(report_data.get("report"), dict):
        report_data = report_data["report"]

    for article in report_data.get("articles", []):
        for ctrl in article.get("controls", []):
            ctrl_id = ctrl.get("controlId", ctrl.get("id", ""))
            if not ctrl_id:
                continue
            fn_name = ctrl.get("evidence_fn", "")
            # evidence_fn may be a callable or string — normalize to string
            if callable(fn_name):
                fn_name = ""
            if not fn_name:
                fn_name = ctrl.get("evidence_fn_name", "")
            status = ctrl.get("combinedStatus", ctrl.get("status", "gap"))
            controls[ctrl_id] = {
                "evidence_fn_name": fn_name,
                "status": status,
            }
    return controls


def _compute_change_cones(feature_ids: List[str]) -> Dict[str, set]:
    """Compute causal cones (wavefronts) for a list of changed features.

    Returns {feature_id: set_of_reachable_node_ids}.
    """
    cones = {}
    try:
        from graph.propagation_analysis import compute_wavefront
    except ImportError:
        return cones

    for fid in feature_ids:
        try:
            wf = compute_wavefront(fid)
            nodes = set()
            for ts_key in ("machine", "human", "business"):
                for node in wf.get(ts_key, []):
                    nodes.add(node["nodeId"])
            # Include the origin feature itself
            nodes.add(fid)
            cones[fid] = nodes
        except Exception as exc:
            log.debug("Wavefront failed for %s: %s", fid, exc)
            cones[fid] = {fid}

    return cones


def _build_fn_node_type_map() -> Dict[str, set]:
    """Build a map of evidence_fn_name → set of feature IDs the fn depends on.

    For graph-required fns, we approximate by collecting all features in the
    graph.  For non-graph fns, returns empty set (no feature deps).
    """
    from graph.compliance_mapping import _EVIDENCE_CLASSIFICATION

    # Cache: all feature IDs from the graph
    all_features = _get_all_feature_ids()

    result = {}
    for fn_name, cls in _EVIDENCE_CLASSIFICATION.items():
        if cls.get("requires_graph", False):
            result[fn_name] = all_features
        else:
            result[fn_name] = set()
    return result


def _get_all_feature_ids() -> set:
    """Return the set of all FeatureNode IDs in the graph."""
    try:
        from graph.schema import get_connection
        conn = get_connection()
        result = conn.execute("MATCH (f:FeatureNode) RETURN f.id")
        ids = set()
        while result.has_next():
            ids.add(str(result.get_next()[0]))
        return ids
    except Exception as e:
        log.debug("Failed to fetch feature IDs from graph: %s", e)
        return set()
