"""Graph diffing logic and snapshot resolution."""

import json
import warnings
from typing import Any

from metaxy.graph.diff.diff_models import (
    AddedNode,
    FieldChange,
    GraphDiff,
    NodeChange,
    RemovedNode,
)
from metaxy.metadata_store.base import FEATURE_VERSIONS_KEY, MetadataStore
from metaxy.models.feature import FeatureGraph
from metaxy.models.types import FeatureKey, FieldKey


class SnapshotResolver:
    """Resolves snapshot version literals to actual snapshot hashes."""

    def resolve_snapshot(
        self, literal: str, store: MetadataStore | None, graph: FeatureGraph | None
    ) -> str:
        """Resolve a snapshot literal to its actual version hash.

        Args:
            literal: Snapshot identifier ("latest", "current", or version hash)
            store: Metadata store to query for snapshots (required for "latest")
            graph: Optional active graph for "current" resolution

        Returns:
            Resolved snapshot version hash

        Raises:
            ValueError: If literal is invalid or cannot be resolved
        """
        if literal == "latest":
            if store is None:
                raise ValueError(
                    "Cannot resolve 'latest': no metadata store provided. "
                    "Provide a store to query for snapshots."
                )
            return self._resolve_latest(store)
        elif literal == "current":
            return self._resolve_current(graph)
        else:
            # Treat as explicit snapshot version
            return literal

    def _resolve_latest(self, store: MetadataStore) -> str:
        """Resolve 'latest' to most recent snapshot in store."""
        snapshots_df = store.read_graph_snapshots()

        if snapshots_df.height == 0:
            raise ValueError(
                "No snapshots found in store. Cannot resolve 'latest'. "
                "Run 'metaxy push' to record a snapshot."
            )

        # read_graph_snapshots() returns sorted by recorded_at descending
        latest_snapshot = snapshots_df["snapshot_version"][0]
        return latest_snapshot

    def _resolve_current(self, graph: FeatureGraph | None) -> str:
        """Resolve 'current' to active graph's snapshot version."""
        if graph is None:
            raise ValueError(
                "Cannot resolve 'current': no active graph provided. "
                "Ensure features are loaded before using 'current'."
            )

        if len(graph.features_by_key) == 0:
            raise ValueError(
                "Cannot resolve 'current': active graph is empty. "
                "Ensure features are loaded before using 'current'."
            )

        return graph.snapshot_version


class GraphDiffer:
    """Compares two graph snapshots and produces a diff."""

    def diff(
        self,
        snapshot1_data: dict[str, dict[str, Any]],
        snapshot2_data: dict[str, dict[str, Any]],
        from_snapshot_version: str = "unknown",
        to_snapshot_version: str = "unknown",
    ) -> GraphDiff:
        """Compute diff between two snapshots.

        Args:
            snapshot1_data: First snapshot (feature_key -> {feature_version, feature_spec, fields})
            snapshot2_data: Second snapshot (feature_key -> {feature_version, feature_spec, fields})
            from_snapshot_version: Source snapshot version
            to_snapshot_version: Target snapshot version

        Returns:
            GraphDiff with added, removed, and changed features
        """
        # Extract feature keys
        keys1 = set(snapshot1_data.keys())
        keys2 = set(snapshot2_data.keys())

        # Identify added and removed features
        added_keys = keys2 - keys1
        removed_keys = keys1 - keys2
        common_keys = keys1 & keys2

        # Build added nodes
        added_nodes = []
        for key_str in sorted(added_keys):
            feature_data = snapshot2_data[key_str]
            feature_spec = feature_data.get("feature_spec", {})

            # Extract fields
            fields_list = []
            for field_dict in feature_spec.get("fields", []):
                field_key_list = field_dict.get("key", [])
                field_key_str = (
                    "/".join(field_key_list)
                    if isinstance(field_key_list, list)
                    else field_key_list
                )
                fields_list.append(
                    {
                        "key": field_key_str,
                        "version": feature_data.get("fields", {}).get(
                            field_key_str, ""
                        ),
                        "code_version": field_dict.get("code_version"),
                    }
                )

            # Extract dependencies
            deps = []
            if feature_spec.get("deps"):
                for dep in feature_spec["deps"]:
                    dep_key = dep.get("key", [])
                    if isinstance(dep_key, list):
                        deps.append(FeatureKey(dep_key))
                    else:
                        deps.append(FeatureKey(dep_key.split("/")))

            added_nodes.append(
                AddedNode(
                    feature_key=FeatureKey(key_str.split("/")),
                    version=feature_data["feature_version"],
                    code_version=feature_spec.get("code_version"),
                    fields=fields_list,
                    dependencies=deps,
                )
            )

        # Build removed nodes
        removed_nodes = []
        for key_str in sorted(removed_keys):
            feature_data = snapshot1_data[key_str]
            feature_spec = feature_data.get("feature_spec", {})

            # Extract fields
            fields_list = []
            for field_dict in feature_spec.get("fields", []):
                field_key_list = field_dict.get("key", [])
                field_key_str = (
                    "/".join(field_key_list)
                    if isinstance(field_key_list, list)
                    else field_key_list
                )
                fields_list.append(
                    {
                        "key": field_key_str,
                        "version": feature_data.get("fields", {}).get(
                            field_key_str, ""
                        ),
                        "code_version": field_dict.get("code_version"),
                    }
                )

            # Extract dependencies
            deps = []
            if feature_spec.get("deps"):
                for dep in feature_spec["deps"]:
                    dep_key = dep.get("key", [])
                    if isinstance(dep_key, list):
                        deps.append(FeatureKey(dep_key))
                    else:
                        deps.append(FeatureKey(dep_key.split("/")))

            removed_nodes.append(
                RemovedNode(
                    feature_key=FeatureKey(key_str.split("/")),
                    version=feature_data["feature_version"],
                    code_version=feature_spec.get("code_version"),
                    fields=fields_list,
                    dependencies=deps,
                )
            )

        # Identify changed features
        changed_nodes = []
        for key_str in sorted(common_keys):
            feature1 = snapshot1_data[key_str]
            feature2 = snapshot2_data[key_str]

            version1 = feature1["feature_version"]
            version2 = feature2["feature_version"]

            spec1 = feature1.get("feature_spec", {})
            spec2 = feature2.get("feature_spec", {})

            fields1 = feature1.get("fields", {})
            fields2 = feature2.get("fields", {})

            # Check if feature version changed
            if version1 != version2:
                # Compute field changes
                field_changes = self._compute_field_changes(fields1, fields2)

                changed_nodes.append(
                    NodeChange(
                        feature_key=FeatureKey(key_str.split("/")),
                        old_version=version1,
                        new_version=version2,
                        old_code_version=spec1.get("code_version"),
                        new_code_version=spec2.get("code_version"),
                        added_fields=[fc for fc in field_changes if fc.is_added],
                        removed_fields=[fc for fc in field_changes if fc.is_removed],
                        changed_fields=[fc for fc in field_changes if fc.is_changed],
                    )
                )

        return GraphDiff(
            from_snapshot_version=from_snapshot_version,
            to_snapshot_version=to_snapshot_version,
            added_nodes=added_nodes,
            removed_nodes=removed_nodes,
            changed_nodes=changed_nodes,
        )

    def _compute_field_changes(
        self, fields1: dict[str, str], fields2: dict[str, str]
    ) -> list[FieldChange]:
        """Compute changes between two field version mappings.

        Args:
            fields1: Field key (string) -> field version (hash) from snapshot1
            fields2: Field key (string) -> field version (hash) from snapshot2

        Returns:
            List of FieldChange objects
        """
        field_keys1 = set(fields1.keys())
        field_keys2 = set(fields2.keys())

        added_fields = field_keys2 - field_keys1
        removed_fields = field_keys1 - field_keys2
        common_fields = field_keys1 & field_keys2

        changes = []

        # Added fields
        for field_key_str in sorted(added_fields):
            changes.append(
                FieldChange(
                    field_key=FieldKey(field_key_str.split("/")),
                    old_version=None,
                    new_version=fields2[field_key_str],
                    old_code_version=None,
                    new_code_version=None,  # TODO: Extract from spec if available
                )
            )

        # Removed fields
        for field_key_str in sorted(removed_fields):
            changes.append(
                FieldChange(
                    field_key=FieldKey(field_key_str.split("/")),
                    old_version=fields1[field_key_str],
                    new_version=None,
                    old_code_version=None,  # TODO: Extract from spec if available
                    new_code_version=None,
                )
            )

        # Changed fields
        for field_key_str in sorted(common_fields):
            version1 = fields1[field_key_str]
            version2 = fields2[field_key_str]

            if version1 != version2:
                changes.append(
                    FieldChange(
                        field_key=FieldKey(field_key_str.split("/")),
                        old_version=version1,
                        new_version=version2,
                        old_code_version=None,  # TODO: Extract from spec if available
                        new_code_version=None,  # TODO: Extract from spec if available
                    )
                )

        return changes

    def create_merged_graph_data(
        self,
        snapshot1_data: dict[str, dict[str, Any]],
        snapshot2_data: dict[str, dict[str, Any]],
        diff: GraphDiff,
    ) -> dict[str, Any]:
        """Create merged graph data structure with status annotations.

        This combines features from both snapshots into a single unified view,
        annotating each feature with its status (added/removed/changed/unchanged).

        Args:
            snapshot1_data: First snapshot data (feature_key -> {feature_version, fields})
            snapshot2_data: Second snapshot data (feature_key -> {feature_version, fields})
            diff: Computed diff between snapshots

        Returns:
            Dict with structure:
            {
                'nodes': {
                    feature_key_str: {
                        'status': 'added' | 'removed' | 'changed' | 'unchanged',
                        'old_version': str | None,
                        'new_version': str | None,
                        'fields': {...},  # fields from relevant snapshot
                        'field_changes': [...],  # FieldChange objects for changed nodes
                        'dependencies': [feature_key_str, ...],  # deps from relevant snapshot
                    }
                },
                'edges': [
                    {'from': feature_key_str, 'to': feature_key_str}
                ]
            }
        """
        # Create status mapping for efficient lookup
        added_keys = {node.feature_key.to_string() for node in diff.added_nodes}
        removed_keys = {node.feature_key.to_string() for node in diff.removed_nodes}
        changed_keys = {
            node.feature_key.to_string(): node for node in diff.changed_nodes
        }

        # Get all feature keys from both snapshots
        all_keys = set(snapshot1_data.keys()) | set(snapshot2_data.keys())

        nodes = {}
        edges = []

        for feature_key_str in all_keys:
            # Determine status
            if feature_key_str in added_keys:
                status = "added"
                old_version = None
                new_version = snapshot2_data[feature_key_str]["feature_version"]
                fields = snapshot2_data[feature_key_str].get("fields", {})
                field_changes = []
                # Dependencies from snapshot2
                deps = self._extract_dependencies(
                    snapshot2_data[feature_key_str].get("feature_spec", {})
                )
            elif feature_key_str in removed_keys:
                status = "removed"
                old_version = snapshot1_data[feature_key_str]["feature_version"]
                new_version = None
                fields = snapshot1_data[feature_key_str].get("fields", {})
                field_changes = []
                # Dependencies from snapshot1
                deps = self._extract_dependencies(
                    snapshot1_data[feature_key_str].get("feature_spec", {})
                )
            elif feature_key_str in changed_keys:
                status = "changed"
                node_change = changed_keys[feature_key_str]
                old_version = node_change.old_version
                new_version = node_change.new_version
                fields = snapshot2_data[feature_key_str].get("fields", {})
                # Combine all field changes from the NodeChange
                field_changes = (
                    node_change.added_fields
                    + node_change.removed_fields
                    + node_change.changed_fields
                )
                # Dependencies from snapshot2 (current version)
                deps = self._extract_dependencies(
                    snapshot2_data[feature_key_str].get("feature_spec", {})
                )
            else:
                # Unchanged
                status = "unchanged"
                old_version = snapshot1_data[feature_key_str]["feature_version"]
                new_version = snapshot2_data[feature_key_str]["feature_version"]
                fields = snapshot2_data[feature_key_str].get("fields", {})
                field_changes = []
                # Dependencies from snapshot2
                deps = self._extract_dependencies(
                    snapshot2_data[feature_key_str].get("feature_spec", {})
                )

            nodes[feature_key_str] = {
                "status": status,
                "old_version": old_version,
                "new_version": new_version,
                "fields": fields,
                "field_changes": field_changes,
                "dependencies": deps,
            }

            # Create edges for dependencies (arrow points from dependency to feature)
            for dep_key in deps:
                edges.append({"from": dep_key, "to": feature_key_str})

        return {
            "nodes": nodes,
            "edges": edges,
        }

    def _extract_dependencies(self, feature_spec: dict[str, Any]) -> list[str]:
        """Extract dependency feature keys from a feature spec.

        Args:
            feature_spec: Parsed feature spec dict

        Returns:
            List of dependency feature keys as strings
        """
        deps = feature_spec.get("deps", [])
        if deps is None:
            return []

        dep_keys = []
        for dep in deps:
            # dep is a dict with 'key' field
            dep_key = dep.get("key", [])
            if isinstance(dep_key, list):
                dep_keys.append("/".join(dep_key))
            else:
                dep_keys.append(dep_key)

        return dep_keys

    def filter_merged_graph(
        self,
        merged_data: dict[str, Any],
        focus_feature: str | None = None,
        up: int | None = None,
        down: int | None = None,
    ) -> dict[str, Any]:
        """Filter merged graph to show only relevant features.

        Args:
            merged_data: Merged graph data with nodes and edges
            focus_feature: Feature key to focus on (string format with / or __)
            up: Number of upstream levels (None = all if focus_feature is set, 0 otherwise)
            down: Number of downstream levels (None = all if focus_feature is set, 0 otherwise)

        Returns:
            Filtered merged graph data with same structure

        Raises:
            ValueError: If focus_feature is specified but not found in graph
        """
        if focus_feature is None:
            # No filtering
            return merged_data

        # Parse feature key (support both / and __ formats)
        if "/" in focus_feature:
            focus_key = focus_feature
        else:
            focus_key = focus_feature.replace("__", "/")

        # Check if focus feature exists
        if focus_key not in merged_data["nodes"]:
            raise ValueError(f"Feature '{focus_feature}' not found in graph")

        # Build dependency graph for traversal
        # Build forward edges (feature -> dependents) and backward edges (feature -> dependencies)
        forward_edges: dict[str, list[str]] = {}  # feature -> list of dependents
        backward_edges: dict[str, list[str]] = {}  # feature -> list of dependencies

        for edge in merged_data["edges"]:
            dep = edge["from"]  # dependency
            feat = edge["to"]  # dependent feature

            if feat not in backward_edges:
                backward_edges[feat] = []
            backward_edges[feat].append(dep)

            if dep not in forward_edges:
                forward_edges[dep] = []
            forward_edges[dep].append(feat)

        # Find features to include
        features_to_include = {focus_key}

        # Add upstream (dependencies)
        # Default behavior: if focus_feature is set but up is not specified, include all upstream
        if up is None:
            # Include all upstream
            upstream = self._get_upstream_features(
                focus_key, backward_edges, max_levels=None
            )
            features_to_include.update(upstream)
        elif up > 0:
            # Include specified number of levels
            upstream = self._get_upstream_features(
                focus_key, backward_edges, max_levels=up
            )
            features_to_include.update(upstream)
        # else: up == 0, don't include upstream

        # Add downstream (dependents)
        # Default behavior: if focus_feature is set but down is not specified, include all downstream
        if down is None:
            # Include all downstream
            downstream = self._get_downstream_features(
                focus_key, forward_edges, max_levels=None
            )
            features_to_include.update(downstream)
        elif down > 0:
            # Include specified number of levels
            downstream = self._get_downstream_features(
                focus_key, forward_edges, max_levels=down
            )
            features_to_include.update(downstream)
        # else: down == 0, don't include downstream

        # Filter nodes and edges
        filtered_nodes = {
            k: v for k, v in merged_data["nodes"].items() if k in features_to_include
        }
        filtered_edges = [
            e
            for e in merged_data["edges"]
            if e["from"] in features_to_include and e["to"] in features_to_include
        ]

        return {
            "nodes": filtered_nodes,
            "edges": filtered_edges,
        }

    def _get_upstream_features(
        self,
        start_key: str,
        backward_edges: dict[str, list[str]],
        max_levels: int | None = None,
        visited: set[str] | None = None,
        level: int = 0,
    ) -> set[str]:
        """Get upstream features (dependencies) recursively."""
        if visited is None:
            visited = set()

        if start_key in visited:
            return set()

        if max_levels is not None and level >= max_levels:
            return set()

        visited.add(start_key)
        upstream: set[str] = set()

        deps = backward_edges.get(start_key, [])
        for dep in deps:
            if dep not in visited:
                upstream.add(dep)
                # Recurse
                upstream.update(
                    self._get_upstream_features(
                        dep, backward_edges, max_levels, visited, level + 1
                    )
                )

        return upstream

    def _get_downstream_features(
        self,
        start_key: str,
        forward_edges: dict[str, list[str]],
        max_levels: int | None = None,
        visited: set[str] | None = None,
        level: int = 0,
    ) -> set[str]:
        """Get downstream features (dependents) recursively."""
        if visited is None:
            visited = set()

        if start_key in visited:
            return set()

        if max_levels is not None and level >= max_levels:
            return set()

        visited.add(start_key)
        downstream: set[str] = set()

        dependents = forward_edges.get(start_key, [])
        for dependent in dependents:
            if dependent not in visited:
                downstream.add(dependent)
                # Recurse
                downstream.update(
                    self._get_downstream_features(
                        dependent, forward_edges, max_levels, visited, level + 1
                    )
                )

        return downstream

    def load_snapshot_data(
        self, store: MetadataStore, snapshot_version: str
    ) -> dict[str, dict[str, Any]]:
        """Load snapshot data from store.

        Args:
            store: Metadata store to query
            snapshot_version: Snapshot version to load

        Returns:
            Dict mapping feature_key (string) -> {feature_version, feature_spec, fields}
            where fields is dict mapping field_key (string) -> field_version (hash)

        Raises:
            ValueError: If snapshot not found in store
        """
        # Query feature_versions table for this snapshot
        try:
            features_lazy = store._read_metadata_native(FEATURE_VERSIONS_KEY)
            if features_lazy is None:
                raise ValueError(
                    f"No feature_versions table found in store. Cannot load snapshot {snapshot_version}."
                )

            # Filter by snapshot_version
            import narwhals as nw

            features_df = (
                features_lazy.filter(nw.col("snapshot_version") == snapshot_version)
                .collect()
                .to_polars()
            )

            if features_df.height == 0:
                raise ValueError(
                    f"Snapshot {snapshot_version} not found in store. "
                    "Run 'metaxy push' to record snapshots or check the version hash."
                )

        except Exception as e:
            raise ValueError(f"Failed to load snapshot {snapshot_version}: {e}") from e

        # Build snapshot data structure
        snapshot_dict = {}
        for row in features_df.iter_rows(named=True):
            feature_key_str = row["feature_key"]
            feature_version = row["feature_version"]
            feature_spec_json = row["feature_spec"]
            feature_class_path = row.get("feature_class_path", "")

            feature_spec_dict = json.loads(feature_spec_json)

            snapshot_dict[feature_key_str] = {
                "feature_version": feature_version,
                "feature_spec": feature_spec_dict,
                "feature_class_path": feature_class_path,
            }

        # Try to reconstruct FeatureGraph from snapshot to compute field versions
        # This may fail if features have been removed/moved, so we handle that gracefully
        graph: FeatureGraph | None = None
        try:
            graph = FeatureGraph.from_snapshot(snapshot_dict)
            graph_available = True
        except ImportError:
            # Some features can't be imported (likely removed) - proceed without graph
            # For diff purposes, we can still show feature-level changes
            # We'll use feature_version as a fallback for all field versions
            graph_available = False
            warnings.warn(
                "Using feature_version as field_version fallback for features that cannot be imported. "
                "This may occur when features have been removed or moved.",
                UserWarning,
                stacklevel=2,
            )

        # Compute field versions using the reconstructed graph (if available)
        from metaxy.models.plan import FQFieldKey

        snapshot_data = {}
        for feature_key_str in snapshot_dict.keys():
            feature_version = snapshot_dict[feature_key_str]["feature_version"]
            feature_spec = snapshot_dict[feature_key_str]["feature_spec"]
            feature_key_obj = FeatureKey(feature_key_str.split("/"))

            # Compute field versions using graph (if available)
            fields_data = {}
            if (
                graph_available
                and graph is not None
                and feature_key_obj in graph.features_by_key
            ):
                # Feature exists in reconstructed graph - compute precise field versions
                for field_dict in feature_spec.get("fields", []):
                    field_key_list = field_dict.get("key")
                    if isinstance(field_key_list, list):
                        field_key = FieldKey(field_key_list)
                        field_key_str_normalized = "/".join(field_key_list)
                    else:
                        field_key = FieldKey([field_key_list])
                        field_key_str_normalized = field_key_list

                    # Compute field version using the graph
                    fq_key = FQFieldKey(feature=feature_key_obj, field=field_key)
                    field_version = graph.get_field_version(fq_key)
                    fields_data[field_key_str_normalized] = field_version
            else:
                # Feature doesn't exist in graph (removed/moved) - use feature_version as fallback
                # All fields get the same version (the feature version)
                for field_dict in feature_spec.get("fields", []):
                    field_key_list = field_dict.get("key")
                    if isinstance(field_key_list, list):
                        field_key_str_normalized = "/".join(field_key_list)
                    else:
                        field_key_str_normalized = field_key_list

                    # Use feature_version directly as fallback
                    fields_data[field_key_str_normalized] = feature_version

            snapshot_data[feature_key_str] = {
                "feature_version": feature_version,
                "fields": fields_data,
                "feature_spec": feature_spec,
            }

        return snapshot_data
