"""
Data lineage tracking system.

Tracks transformations and maintains lineage graph for audit and reproducibility.
"""

from datetime import datetime
from typing import Any
from uuid import UUID, uuid4

from pydantic import BaseModel, Field

from app.core.logging import get_logger

logger = get_logger(__name__)


class LineageNode(BaseModel):
    """A node in the lineage graph (dataset)."""

    dataset_id: UUID
    dataset_name: str
    created_at: datetime
    row_count: int = 0
    column_count: int = 0


class LineageEdge(BaseModel):
    """An edge in the lineage graph (transformation)."""

    id: UUID = Field(default_factory=uuid4)
    source_dataset_id: UUID
    target_dataset_id: UUID | None
    pipeline_id: UUID
    pipeline_name: str
    transformation_summary: dict[str, Any]
    created_at: datetime = Field(default_factory=datetime.utcnow)


class LineageGraph(BaseModel):
    """Complete lineage graph."""

    nodes: list[LineageNode] = Field(default_factory=list)
    edges: list[LineageEdge] = Field(default_factory=list)

    def add_node(self, node: LineageNode) -> None:
        """Add a node to the graph."""
        # Check if node already exists
        for existing in self.nodes:
            if existing.dataset_id == node.dataset_id:
                return
        self.nodes.append(node)

    def add_edge(self, edge: LineageEdge) -> None:
        """Add an edge to the graph."""
        self.edges.append(edge)

    def get_lineage_for_dataset(self, dataset_id: UUID) -> "LineageGraph":
        """Get lineage subgraph for a specific dataset."""
        # Find all edges involving this dataset
        related_edges = [
            edge
            for edge in self.edges
            if edge.source_dataset_id == dataset_id or edge.target_dataset_id == dataset_id
        ]

        # Find all related nodes
        related_dataset_ids = set()
        related_dataset_ids.add(dataset_id)
        for edge in related_edges:
            related_dataset_ids.add(edge.source_dataset_id)
            if edge.target_dataset_id:
                related_dataset_ids.add(edge.target_dataset_id)

        related_nodes = [node for node in self.nodes if node.dataset_id in related_dataset_ids]

        return LineageGraph(nodes=related_nodes, edges=related_edges)

    def get_ancestors(self, dataset_id: UUID) -> list[LineageNode]:
        """Get all ancestor datasets (upstream in lineage)."""
        ancestors = []
        visited = set()

        def dfs(current_id: UUID) -> None:
            if current_id in visited:
                return
            visited.add(current_id)

            # Find edges where current dataset is the target
            parent_edges = [
                edge for edge in self.edges if edge.target_dataset_id == current_id
            ]

            for edge in parent_edges:
                # Add parent node
                parent_node = next(
                    (n for n in self.nodes if n.dataset_id == edge.source_dataset_id), None
                )
                if parent_node and parent_node not in ancestors:
                    ancestors.append(parent_node)
                # Recursively find ancestors
                dfs(edge.source_dataset_id)

        dfs(dataset_id)
        return ancestors

    def get_descendants(self, dataset_id: UUID) -> list[LineageNode]:
        """Get all descendant datasets (downstream in lineage)."""
        descendants = []
        visited = set()

        def dfs(current_id: UUID) -> None:
            if current_id in visited:
                return
            visited.add(current_id)

            # Find edges where current dataset is the source
            child_edges = [
                edge for edge in self.edges if edge.source_dataset_id == current_id
            ]

            for edge in child_edges:
                if edge.target_dataset_id:
                    # Add child node
                    child_node = next(
                        (n for n in self.nodes if n.dataset_id == edge.target_dataset_id), None
                    )
                    if child_node and child_node not in descendants:
                        descendants.append(child_node)
                    # Recursively find descendants
                    dfs(edge.target_dataset_id)

        dfs(dataset_id)
        return descendants

    def to_mermaid(self) -> str:
        """Convert lineage graph to Mermaid diagram format."""
        lines = ["graph LR"]
        
        # Add nodes
        for node in self.nodes:
            node_id = str(node.dataset_id)[:8]
            lines.append(f'    {node_id}["{node.dataset_name}"]')
        
        # Add edges
        for edge in self.edges:
            source_id = str(edge.source_dataset_id)[:8]
            target_id = str(edge.target_dataset_id)[:8] if edge.target_dataset_id else "unknown"
            pipeline_name = edge.pipeline_name[:20]
            lines.append(f'    {source_id} -->|"{pipeline_name}"| {target_id}')
        
        return "\n".join(lines)


class LineageTracker:
    """Service for tracking data lineage."""

    def __init__(self) -> None:
        """Initialize lineage tracker."""
        self.graph = LineageGraph()

    def track_dataset(
        self, dataset_id: UUID, dataset_name: str, row_count: int = 0, column_count: int = 0
    ) -> None:
        """Track a new dataset."""
        node = LineageNode(
            dataset_id=dataset_id,
            dataset_name=dataset_name,
            created_at=datetime.utcnow(),
            row_count=row_count,
            column_count=column_count,
        )
        self.graph.add_node(node)
        logger.info("Dataset tracked in lineage", dataset_id=str(dataset_id))

    def track_transformation(
        self,
        source_dataset_id: UUID,
        target_dataset_id: UUID | None,
        pipeline_id: UUID,
        pipeline_name: str,
        transformation_summary: dict[str, Any],
    ) -> None:
        """Track a transformation (pipeline execution)."""
        edge = LineageEdge(
            source_dataset_id=source_dataset_id,
            target_dataset_id=target_dataset_id,
            pipeline_id=pipeline_id,
            pipeline_name=pipeline_name,
            transformation_summary=transformation_summary,
        )
        self.graph.add_edge(edge)
        logger.info(
            "Transformation tracked in lineage",
            pipeline_id=str(pipeline_id),
            source=str(source_dataset_id),
            target=str(target_dataset_id) if target_dataset_id else None,
        )

    def get_lineage(self, dataset_id: UUID) -> LineageGraph:
        """Get lineage for a specific dataset."""
        return self.graph.get_lineage_for_dataset(dataset_id)

    def get_lineage_summary(self, dataset_id: UUID) -> dict[str, Any]:
        """Get lineage summary for a dataset."""
        ancestors = self.graph.get_ancestors(dataset_id)
        descendants = self.graph.get_descendants(dataset_id)

        return {
            "dataset_id": str(dataset_id),
            "ancestor_count": len(ancestors),
            "descendant_count": len(descendants),
            "ancestors": [
                {"id": str(a.dataset_id), "name": a.dataset_name} for a in ancestors
            ],
            "descendants": [
                {"id": str(d.dataset_id), "name": d.dataset_name} for d in descendants
            ],
        }
