"""GraphContextMixin — reusable graph/IKE enrichment for workflows.

Replaces copy-pasted graph integration blocks across workflows with a
single mixin that handles:
- Fallback policy resolution (params > config > "best_effort")
- IKE graph and formula graph collection
- Per-run caching via ``SDKContext.cache_get/set``
- Provenance timing for auditability
- ``strict`` / ``best_effort`` / ``off`` error handling

Usage::

    class MyWorkflow(GraphContextMixin, BaseWorkflow):
        def execute(self, params, ctx):
            gctx = self.collect_graph_context(params, ctx)
            # gctx.available, gctx.ike_graph, gctx.formula_graph, etc.
            result.metadata["graph_context"] = gctx.to_dict()
"""

from __future__ import annotations

import logging
import time
from typing import Any, Dict, Optional

from databridge_core.sdk.core.types import GraphContext

logger = logging.getLogger(__name__)


class GraphContextMixin:
    """Mixin providing standardized graph/IKE context collection.

    Opt-in: workflows inherit from ``(GraphContextMixin, BaseWorkflow)``
    and call ``self.collect_graph_context(params, ctx)`` during execute.
    """

    def collect_graph_context(
        self,
        params: Dict[str, Any],
        ctx: Optional[Any] = None,
    ) -> GraphContext:
        """Collect graph enrichment context for a workflow run.

        Reads fallback policy from ``params["graph_fallback_policy"]``,
        falling back to ``ctx.config.graph_fallback_policy`` then
        ``"best_effort"``.

        Returns a populated ``GraphContext`` dataclass.
        """
        gctx = GraphContext()

        if ctx is None:
            return gctx

        # Resolve policy
        policy = params.get(
            "graph_fallback_policy",
            getattr(ctx.config, "graph_fallback_policy", "best_effort"),
        )

        if policy == "off":
            return gctx

        # Set trace_id from context
        gctx.trace_id = getattr(ctx, "trace_id", "")

        t0 = time.monotonic()

        # IKE graph
        self._collect_ike_graph(params, ctx, gctx, policy)

        # Formula graph
        self._collect_formula_graph(params, ctx, gctx, policy)

        elapsed_ms = round((time.monotonic() - t0) * 1000, 1)
        gctx.provenance["elapsed_ms"] = elapsed_ms
        gctx.provenance["policy"] = policy
        gctx.available = (
            gctx.ike_graph is not None
            or gctx.formula_graph is not None
            or len(gctx.lineage_nodes) > 0
        )

        return gctx

    def _collect_ike_graph(
        self,
        params: Dict[str, Any],
        ctx: Any,
        gctx: GraphContext,
        policy: str,
    ) -> None:
        """Try to collect IKE graph data, respecting cache and policy."""
        if not hasattr(ctx, "graphs") or not ctx.graphs.ike_graph_available:
            return

        cache_key = "graph_mixin:ike_graph"
        cached = ctx.cache_get(cache_key) if hasattr(ctx, "cache_get") else None
        if cached is not None:
            gctx.ike_graph = cached
            gctx.provenance["ike_graph_source"] = "cache"
            return

        try:
            query = params.get("ike_query", "")
            result = ctx.graphs.query_ike_graph(
                query=query,
                trace_id=getattr(ctx, "trace_id", None),
            )
            if result is not None:
                gctx.ike_graph = result
                gctx.provenance["ike_graph_source"] = "live"
                if hasattr(ctx, "cache_set"):
                    ctx.cache_set(cache_key, result)
        except Exception as exc:
            msg = f"IKE graph collection failed: {exc}"
            gctx.errors["ike_graph"] = msg
            if policy == "strict":
                raise RuntimeError(msg) from exc
            logger.warning(msg)

    def _collect_formula_graph(
        self,
        params: Dict[str, Any],
        ctx: Any,
        gctx: GraphContext,
        policy: str,
    ) -> None:
        """Try to collect formula dependency graph, respecting cache and policy."""
        if not hasattr(ctx, "graphs") or not ctx.graphs.formula_graph_available:
            return

        workbook_id = params.get("workbook_id")
        if not workbook_id:
            return

        cache_key = f"graph_mixin:formula_graph:{workbook_id}"
        cached = ctx.cache_get(cache_key) if hasattr(ctx, "cache_get") else None
        if cached is not None:
            gctx.formula_graph = cached
            gctx.provenance["formula_graph_source"] = "cache"
            return

        try:
            result = ctx.graphs.get_formula_dependency_graph(workbook_id)
            if result is not None:
                gctx.formula_graph = result
                gctx.provenance["formula_graph_source"] = "live"
                if hasattr(ctx, "cache_set"):
                    ctx.cache_set(cache_key, result)
        except Exception as exc:
            msg = f"Formula graph collection failed: {exc}"
            gctx.errors["formula_graph"] = msg
            if policy == "strict":
                raise RuntimeError(msg) from exc
            logger.warning(msg)

    def _register_lineage_node(self, node_id: str, gctx: GraphContext) -> None:
        """Track a lineage node registration in the graph context."""
        gctx.lineage_nodes.append(node_id)

    def _register_evidence_id(self, evidence_id: str, gctx: GraphContext) -> None:
        """Track an evidence entry in the graph context."""
        gctx.evidence_ids.append(evidence_id)

    def _register_vector_doc(self, doc_id: str, gctx: GraphContext) -> None:
        """Track a vector-indexed document in the graph context."""
        gctx.vector_doc_ids.append(doc_id)
