"""Data lineage tracking via WorldDAG edges.

Records data ingestion, transformation, and derivation as edges in the
WorldDAG. Lineage IS causal: "this dataset exists BECAUSE of that source."
"""

from __future__ import annotations

from orxaq.kernel import (
    AuditLog,
    CausalEdge,
    CausalRole,
    Entity,
    VariableType,
    WorldDAG,
)


class LineageTracker:
    """Track data lineage as edges in the WorldDAG."""

    def __init__(self, dag: WorldDAG, audit_log: AuditLog) -> None:
        self._dag = dag
        self._audit = audit_log

    def record_ingest(
        self,
        source: str,
        dataset_id: str,
        columns: list[str] | None = None,
    ) -> None:
        """Record a data ingestion event.

        Creates Entity nodes for source and dataset, adds a lineage edge.
        """
        # Create source entity if not exists
        if self._dag.get_entity(source) is None:
            self._dag.add_entity(
                Entity(
                    entity_id=source,
                    name=f"Source: {source}",
                    causal_role=CausalRole.EXOGENOUS,
                    variable_type=VariableType.STATIC,
                    domain="lineage",
                )
            )

        # Create dataset entity if not exists
        if self._dag.get_entity(dataset_id) is None:
            desc = f"columns: {', '.join(columns)}" if columns else ""
            self._dag.add_entity(
                Entity(
                    entity_id=dataset_id,
                    name=f"Dataset: {dataset_id}",
                    description=desc,
                    causal_role=CausalRole.OUTCOME,
                    variable_type=VariableType.DERIVED,
                    domain="lineage",
                )
            )

        # Add lineage edge
        self._dag.add_edge(
            CausalEdge(
                source=source,
                target=dataset_id,
                mechanism="data_ingest",
                confidence=1.0,
                edge_type="lineage",
                provenance=f"Ingested from {source}",
            )
        )

        self._audit.append(
            action="data_ingest",
            actor="fabric.lineage",
            details=f"Ingested {dataset_id} from {source}",
        )

    def record_transform(
        self,
        input_id: str,
        output_id: str,
        operation: str,
    ) -> None:
        """Record a data transformation.

        Adds an edge from input dataset to output dataset.
        """
        # Create output entity if not exists
        if self._dag.get_entity(output_id) is None:
            self._dag.add_entity(
                Entity(
                    entity_id=output_id,
                    name=f"Dataset: {output_id}",
                    causal_role=CausalRole.OUTCOME,
                    variable_type=VariableType.DERIVED,
                    domain="lineage",
                )
            )

        self._dag.add_edge(
            CausalEdge(
                source=input_id,
                target=output_id,
                mechanism=operation,
                confidence=1.0,
                edge_type="lineage",
                provenance=f"Transform: {operation}",
            )
        )

        self._audit.append(
            action="data_transform",
            actor="fabric.lineage",
            details=f"{input_id} -> {output_id} via {operation}",
        )

    def trace_upstream(self, dataset_id: str) -> list[str]:
        """Return all ancestor datasets (where did this data come from?)."""
        ancestors = self._dag.ancestors(dataset_id)
        return sorted(ancestors)

    def trace_downstream(self, dataset_id: str) -> list[str]:
        """Return all descendant datasets (what depends on this?)."""
        descendants = self._dag.descendants(dataset_id)
        return sorted(descendants)

    def impact_analysis(self, source_id: str) -> dict[str, int]:
        """If this source changes, what downstream assets are affected?

        Returns {dataset_id: hop_distance_from_source}.
        """
        result: dict[str, int] = {}
        current_level = {source_id}
        distance = 0

        visited: set[str] = set()
        while current_level:
            next_level: set[str] = set()
            for node in current_level:
                if node in visited:  # pragma: no cover
                    continue
                visited.add(node)
                if node != source_id:
                    result[node] = distance
                for child in self._dag.children(node):
                    if child not in visited:
                        next_level.add(child)
            current_level = next_level
            distance += 1

        return result
