"""
Unified MCP tools for Lineage & Impact Analysis (Phase 3D-5 consolidation).

Consolidates 11 individual lineage tools into 3 action-dispatched tools:
- lineage_graph(action, ...) -- 3 actions: add_node, list, validate
- lineage_track(action, ...) -- 5 actions: track_column, get_column, get_table, downstream_impact, upstream_deps
- lineage_visualize(action, ...) -- 3 actions: change_impact, build_graph, export_diagram

Original tool names are preserved as deprecation wrappers in mcp_tools.py.
"""

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

from .types import (
    NodeType,
    TransformationType,
    ChangeType,
)

logger = logging.getLogger(__name__)

# ============================================================================
# Lazy singletons
# ============================================================================

_tracker = None
_analyzer = None


def _ensure_tracker():
    global _tracker
    if _tracker is None:
        from .lineage_tracker import LineageTracker
        _tracker = LineageTracker()
    return _tracker


def _ensure_analyzer():
    global _analyzer
    if _analyzer is None:
        from .impact_analyzer import ImpactAnalyzer
        _analyzer = ImpactAnalyzer(_ensure_tracker())
    return _analyzer


# ============================================================================
# lineage_graph action handlers
# ============================================================================

def _graph_add_node(settings, **kwargs) -> Dict[str, Any]:
    graph_name = kwargs.get("graph_name")
    node_name = kwargs.get("node_name")
    node_type = kwargs.get("node_type")
    if not graph_name or not node_name or not node_type:
        return {"error": "graph_name, node_name, and node_type are required for 'add_node' action"}

    tracker = _ensure_tracker()

    try:
        nt = NodeType(node_type.upper())
    except ValueError:
        return {"success": False, "error": f"Invalid node type: {node_type}"}

    columns = kwargs.get("columns")
    col_list = None
    if columns:
        col_list = json.loads(columns)

    tags = kwargs.get("tags")
    tag_list = None
    if tags:
        tag_list = [t.strip() for t in tags.split(",")]

    node = tracker.add_node(
        graph_name=graph_name,
        name=node_name,
        node_type=nt,
        database=kwargs.get("database"),
        schema_name=kwargs.get("schema_name"),
        columns=col_list,
        description=kwargs.get("description"),
        tags=tag_list,
    )

    return {
        "success": True,
        "node_id": node.id,
        "node_name": node.name,
        "node_type": node.node_type.value,
        "fully_qualified_name": node.fully_qualified_name,
        "column_count": len(node.columns),
        "message": f"Added node '{node_name}' to graph '{graph_name}'",
    }


def _graph_list(settings, **kwargs) -> Dict[str, Any]:
    tracker = _ensure_tracker()
    graphs = tracker.list_graphs()

    return {
        "success": True,
        "graph_count": len(graphs),
        "graphs": graphs,
    }


def _graph_validate(settings, **kwargs) -> Dict[str, Any]:
    graph_name = kwargs.get("graph_name")
    if not graph_name:
        return {"error": "graph_name is required for 'validate' action"}

    analyzer = _ensure_analyzer()
    result = analyzer.validate_lineage(graph_name)

    return {
        "success": True,
        "is_valid": result.is_valid,
        "completeness_score": round(result.completeness_score * 100, 1),
        "node_count": result.node_count,
        "edge_count": result.edge_count,
        "column_lineage_count": result.column_lineage_count,
        "orphan_nodes": result.orphan_nodes,
        "missing_sources": result.missing_sources,
        "circular_dependencies": result.circular_dependencies,
        "warnings": result.warnings,
    }


# ============================================================================
# lineage_track action handlers
# ============================================================================

def _track_track_column(settings, **kwargs) -> Dict[str, Any]:
    graph_name = kwargs.get("graph_name")
    source_node = kwargs.get("source_node")
    source_columns = kwargs.get("source_columns")
    target_node = kwargs.get("target_node")
    target_column = kwargs.get("target_column")
    if not all([graph_name, source_node, source_columns, target_node, target_column]):
        return {"error": "graph_name, source_node, source_columns, target_node, and target_column are required for 'track_column' action"}

    tracker = _ensure_tracker()

    transformation_type = kwargs.get("transformation_type", "DIRECT")
    try:
        tt = TransformationType(transformation_type.upper())
    except ValueError:
        tt = TransformationType.DIRECT

    source_col_list = [c.strip() for c in source_columns.split(",")]

    lineage = tracker.add_column_lineage(
        graph_name=graph_name,
        source_node=source_node,
        source_columns=source_col_list,
        target_node=target_node,
        target_column=target_column,
        transformation_type=tt,
        transformation_expression=kwargs.get("transformation_expression"),
    )

    return {
        "success": True,
        "lineage_id": lineage.id,
        "source_node_id": lineage.source_node_id,
        "source_columns": lineage.source_columns,
        "target_node_id": lineage.target_node_id,
        "target_column": lineage.target_column,
        "transformation_type": lineage.transformation_type.value,
        "message": f"Added column lineage: {source_columns} -> {target_column}",
    }


def _track_get_column(settings, **kwargs) -> Dict[str, Any]:
    graph_name = kwargs.get("graph_name")
    node = kwargs.get("node")
    column = kwargs.get("column")
    if not graph_name or not node or not column:
        return {"error": "graph_name, node, and column are required for 'get_column' action"}

    tracker = _ensure_tracker()
    direction = kwargs.get("direction", "upstream")

    lineage = tracker.get_column_lineage(
        graph_name=graph_name,
        node=node,
        column=column,
        direction=direction,
    )

    return {
        "success": True,
        "node": node,
        "column": column,
        "direction": direction,
        "lineage_count": len(lineage),
        "lineage": lineage,
    }


def _track_get_table(settings, **kwargs) -> Dict[str, Any]:
    graph_name = kwargs.get("graph_name")
    node = kwargs.get("node")
    if not graph_name or not node:
        return {"error": "graph_name and node are required for 'get_table' action"}

    tracker = _ensure_tracker()
    direction = kwargs.get("direction", "both")
    max_depth = kwargs.get("max_depth", 10)

    result = tracker.get_table_lineage(
        graph_name=graph_name,
        node=node,
        direction=direction,
        max_depth=max_depth,
    )

    if "error" in result:
        return {"success": False, "error": result["error"]}

    return {
        "success": True,
        "node": result.get("node_name"),
        "node_id": result.get("node_id"),
        "upstream_count": len(result.get("upstream", [])),
        "downstream_count": len(result.get("downstream", [])),
        "upstream": result.get("upstream", []),
        "downstream": result.get("downstream", []),
    }


def _track_downstream_impact(settings, **kwargs) -> Dict[str, Any]:
    graph_name = kwargs.get("graph_name")
    node = kwargs.get("node")
    if not graph_name or not node:
        return {"error": "graph_name and node are required for 'downstream_impact' action"}

    analyzer = _ensure_analyzer()
    max_depth = kwargs.get("max_depth", 10)

    result = analyzer.get_downstream_impact(graph_name, node, max_depth)

    if "error" in result:
        return {"success": False, "error": result["error"]}

    return {
        "success": True,
        **result,
    }


def _track_upstream_deps(settings, **kwargs) -> Dict[str, Any]:
    graph_name = kwargs.get("graph_name")
    node = kwargs.get("node")
    if not graph_name or not node:
        return {"error": "graph_name and node are required for 'upstream_deps' action"}

    analyzer = _ensure_analyzer()
    max_depth = kwargs.get("max_depth", 10)

    result = analyzer.get_upstream_dependencies(graph_name, node, max_depth)

    if "error" in result:
        return {"success": False, "error": result["error"]}

    return {
        "success": True,
        **result,
    }


# ============================================================================
# lineage_visualize action handlers
# ============================================================================

def _visualize_change_impact(settings, **kwargs) -> Dict[str, Any]:
    graph_name = kwargs.get("graph_name")
    node = kwargs.get("node")
    change_type = kwargs.get("change_type")
    if not graph_name or not node or not change_type:
        return {"error": "graph_name, node, and change_type are required for 'change_impact' action"}

    analyzer = _ensure_analyzer()

    try:
        ct = ChangeType(change_type.upper())
    except ValueError:
        return {"success": False, "error": f"Invalid change type: {change_type}"}

    column = kwargs.get("column")
    new_column_name = kwargs.get("new_column_name")

    if ct == ChangeType.REMOVE_COLUMN and column:
        result = analyzer.analyze_column_removal(graph_name, node, column)
    elif ct == ChangeType.RENAME_COLUMN and column and new_column_name:
        result = analyzer.analyze_column_rename(
            graph_name, node, column, new_column_name
        )
    else:
        result = analyzer.analyze_hierarchy_change(graph_name, node, ct)

    return {
        "success": True,
        "change_type": result.change_type.value,
        "source_node": result.source_node_name,
        "change_description": result.change_description,
        "total_impacted": result.total_impacted,
        "max_severity": result.max_severity.value,
        "critical_count": result.critical_count,
        "high_count": result.high_count,
        "medium_count": result.medium_count,
        "low_count": result.low_count,
        "impacted_objects": [i.to_dict() for i in result.impacted_objects],
    }


def _visualize_build_graph(settings, **kwargs) -> Dict[str, Any]:
    graph_name = kwargs.get("graph_name")
    node = kwargs.get("node")
    if not graph_name or not node:
        return {"error": "graph_name and node are required for 'build_graph' action"}

    analyzer = _ensure_analyzer()
    direction = kwargs.get("direction", "downstream")
    max_depth = kwargs.get("max_depth", 10)

    dep_graph = analyzer.build_dependency_graph(
        graph_name, node, direction, max_depth
    )

    return {
        "success": True,
        "root_node_id": dep_graph.root_node_id,
        "direction": dep_graph.direction,
        "total_nodes": dep_graph.total_nodes,
        "max_depth": dep_graph.max_depth,
        "nodes": [n.to_dict() for n in dep_graph.nodes],
        "edges": dep_graph.edges,
    }


def _visualize_export_diagram(settings, **kwargs) -> Dict[str, Any]:
    graph_name = kwargs.get("graph_name")
    node = kwargs.get("node")
    if not graph_name or not node:
        return {"error": "graph_name and node are required for 'export_diagram' action"}

    analyzer = _ensure_analyzer()
    direction = kwargs.get("direction", "downstream")
    format = kwargs.get("format", "mermaid")
    max_depth = kwargs.get("max_depth", 10)

    dep_graph = analyzer.build_dependency_graph(
        graph_name, node, direction, max_depth
    )

    if format.lower() == "mermaid":
        diagram = dep_graph.to_mermaid()
    elif format.lower() == "dot":
        diagram = dep_graph.to_dot()
    else:
        return {"success": False, "error": f"Unknown format: {format}"}

    return {
        "success": True,
        "format": format,
        "node_count": dep_graph.total_nodes,
        "edge_count": len(dep_graph.edges),
        "diagram": diagram,
    }


# ============================================================================
# Action dispatch dictionaries
# ============================================================================

_GRAPH_ACTIONS = {
    "add_node": _graph_add_node,
    "list": _graph_list,
    "validate": _graph_validate,
}

_TRACK_ACTIONS = {
    "track_column": _track_track_column,
    "get_column": _track_get_column,
    "get_table": _track_get_table,
    "downstream_impact": _track_downstream_impact,
    "upstream_deps": _track_upstream_deps,
}

_VISUALIZE_ACTIONS = {
    "change_impact": _visualize_change_impact,
    "build_graph": _visualize_build_graph,
    "export_diagram": _visualize_export_diagram,
}


# ============================================================================
# Unified dispatch functions
# ============================================================================

def dispatch_lineage_graph(settings, action: str, **kwargs) -> Dict[str, Any]:
    """Dispatch a lineage_graph action."""
    handler = _GRAPH_ACTIONS.get(action)
    if not handler:
        return {
            "error": f"Unknown action: '{action}'",
            "valid_actions": sorted(_GRAPH_ACTIONS.keys()),
        }
    try:
        return handler(settings, **kwargs)
    except Exception as e:
        logger.error(f"lineage_graph({action}) failed: {e}")
        return {"error": f"lineage_graph({action}) failed: {e}"}


def dispatch_lineage_track(settings, action: str, **kwargs) -> Dict[str, Any]:
    """Dispatch a lineage_track action."""
    handler = _TRACK_ACTIONS.get(action)
    if not handler:
        return {
            "error": f"Unknown action: '{action}'",
            "valid_actions": sorted(_TRACK_ACTIONS.keys()),
        }
    try:
        return handler(settings, **kwargs)
    except Exception as e:
        logger.error(f"lineage_track({action}) failed: {e}")
        return {"error": f"lineage_track({action}) failed: {e}"}


def dispatch_lineage_visualize(settings, action: str, **kwargs) -> Dict[str, Any]:
    """Dispatch a lineage_visualize action."""
    handler = _VISUALIZE_ACTIONS.get(action)
    if not handler:
        return {
            "error": f"Unknown action: '{action}'",
            "valid_actions": sorted(_VISUALIZE_ACTIONS.keys()),
        }
    try:
        return handler(settings, **kwargs)
    except Exception as e:
        logger.error(f"lineage_visualize({action}) failed: {e}")
        return {"error": f"lineage_visualize({action}) failed: {e}"}


# ============================================================================
# Registration
# ============================================================================

def register_unified_lineage_tools(mcp, settings):
    """Register the 3 unified lineage MCP tools."""

    @mcp.tool()
    def lineage_graph(
        action: str,
        graph_name: Optional[str] = None,
        node_name: Optional[str] = None,
        node_type: Optional[str] = None,
        database: Optional[str] = None,
        schema_name: Optional[str] = None,
        columns: Optional[str] = None,
        description: Optional[str] = None,
        tags: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Unified lineage graph management tool. Replaces 3 individual tools.

        Actions:
        - add_node: Add a node to the lineage graph (requires graph_name, node_name, node_type)
        - list: List all lineage graphs (no additional params)
        - validate: Validate lineage graph completeness (requires graph_name)

        Args:
            action: The action to perform (add_node, list, validate)
            graph_name: Name of the lineage graph (for add_node, validate)
            node_name: Node name / object name (for add_node)
            node_type: Type of node: TABLE, VIEW, DYNAMIC_TABLE, HIERARCHY, DATA_MART, DBT_MODEL (for add_node)
            database: Database name (for add_node)
            schema_name: Schema name (for add_node)
            columns: JSON array of column definitions (for add_node)
            description: Node description (for add_node)
            tags: Comma-separated tags (for add_node)

        Returns:
            Action-specific result dict
        """
        return dispatch_lineage_graph(settings, action, **{
            k: v for k, v in {
                "graph_name": graph_name, "node_name": node_name,
                "node_type": node_type, "database": database,
                "schema_name": schema_name, "columns": columns,
                "description": description, "tags": tags,
            }.items() if v is not None
        })

    @mcp.tool()
    def lineage_track(
        action: str,
        graph_name: Optional[str] = None,
        source_node: Optional[str] = None,
        source_columns: Optional[str] = None,
        target_node: Optional[str] = None,
        target_column: Optional[str] = None,
        transformation_type: str = "DIRECT",
        transformation_expression: Optional[str] = None,
        node: Optional[str] = None,
        column: Optional[str] = None,
        direction: Optional[str] = None,
        max_depth: int = 10,
    ) -> Dict[str, Any]:
        """
        Unified lineage tracking and dependency tool. Replaces 5 individual tools.

        Actions:
        - track_column: Add column-level lineage (requires graph_name, source_node, source_columns, target_node, target_column)
        - get_column: Get lineage for a specific column (requires graph_name, node, column)
        - get_table: Get lineage for a table/object (requires graph_name, node)
        - downstream_impact: Get all downstream dependents (requires graph_name, node)
        - upstream_deps: Get all upstream dependencies (requires graph_name, node)

        Args:
            action: The action to perform (track_column, get_column, get_table, downstream_impact, upstream_deps)
            graph_name: Name of the lineage graph (all actions)
            source_node: Source node name or ID (for track_column)
            source_columns: Comma-separated source column names (for track_column)
            target_node: Target node name or ID (for track_column)
            target_column: Target column name (for track_column)
            transformation_type: DIRECT, AGGREGATION, CALCULATION, FILTER, JOIN, CASE (for track_column)
            transformation_expression: Optional expression (for track_column)
            node: Node name or ID (for get_column, get_table, downstream_impact, upstream_deps)
            column: Column name (for get_column)
            direction: "upstream", "downstream", or "both" (for get_column, get_table)
            max_depth: Maximum traversal depth (for get_table, downstream_impact, upstream_deps)

        Returns:
            Action-specific result dict
        """
        return dispatch_lineage_track(settings, action, **{
            k: v for k, v in {
                "graph_name": graph_name, "source_node": source_node,
                "source_columns": source_columns, "target_node": target_node,
                "target_column": target_column, "transformation_type": transformation_type,
                "transformation_expression": transformation_expression,
                "node": node, "column": column, "direction": direction,
                "max_depth": max_depth,
            }.items() if v is not None
        })

    @mcp.tool()
    def lineage_visualize(
        action: str,
        graph_name: Optional[str] = None,
        node: Optional[str] = None,
        change_type: Optional[str] = None,
        column: Optional[str] = None,
        new_column_name: Optional[str] = None,
        direction: str = "downstream",
        format: str = "mermaid",
        max_depth: int = 10,
    ) -> Dict[str, Any]:
        """
        Unified lineage visualization and impact analysis tool. Replaces 3 individual tools.

        Actions:
        - change_impact: Analyze impact of a proposed change (requires graph_name, node, change_type)
        - build_graph: Build a dependency graph for visualization (requires graph_name, node)
        - export_diagram: Export lineage as Mermaid or DOT diagram (requires graph_name, node)

        Args:
            action: The action to perform (change_impact, build_graph, export_diagram)
            graph_name: Name of the lineage graph (all actions)
            node: Node name or ID (all actions)
            change_type: REMOVE_COLUMN, RENAME_COLUMN, REMOVE_NODE, MODIFY_MAPPING, MODIFY_FORMULA (for change_impact)
            column: Column name for column changes (for change_impact)
            new_column_name: New column name for renames (for change_impact)
            direction: "upstream" or "downstream" (for build_graph, export_diagram)
            format: Output format "mermaid" or "dot" (for export_diagram)
            max_depth: Maximum traversal depth (for build_graph, export_diagram)

        Returns:
            Action-specific result dict
        """
        return dispatch_lineage_visualize(settings, action, **{
            k: v for k, v in {
                "graph_name": graph_name, "node": node,
                "change_type": change_type, "column": column,
                "new_column_name": new_column_name, "direction": direction,
                "format": format, "max_depth": max_depth,
            }.items() if v is not None
        })

    logger.info("Registered 3 unified lineage tools: lineage_graph, lineage_track, lineage_visualize")
    return ["lineage_graph", "lineage_track", "lineage_visualize"]
