"""
Unified MCP tools for Data Catalog (Phase 3A consolidation).

Consolidates 19 individual catalog tools into 3 action-dispatched tools:
- catalog_asset(action, ...) — 14 actions for asset management
- catalog_term(action, ...) — 4 actions for glossary terms
- catalog_lineage(action, ...) — 4 actions for lineage extraction

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

import logging
import re
from pathlib import Path
from typing import Any, Dict, List, Optional

from .types import (
    AssetType,
    CatalogStats,
    DataAsset,
    DataClassification,
    DataQualityTier,
    GlossaryTerm,
    Owner,
    OwnershipRole,
    QualityMetrics,
    ScanConfig,
    SearchQuery,
    Tag,
    TermStatus,
)
from .catalog_store import CatalogStore
from .scanner import CatalogScanner
from src.utils.tool_wrapper import safe_tool

logger = logging.getLogger(__name__)

# ============================================================================
# Shared singletons (used by both unified.py and mcp_tools.py wrappers)
# ============================================================================

_catalog_store: Optional[CatalogStore] = None
_catalog_scanner: Optional[CatalogScanner] = None


def _get_query_func(settings):
    """Get the query function from connections API."""
    try:
        try:
            from src.connections_api import get_client
        except ImportError:
            from connections_api import get_client

        client = get_client(settings)

        def query_func(connection_id: str, query: str) -> List[Dict]:
            return client.execute_query(connection_id, query)

        return query_func
    except Exception as e:
        logger.warning(f"Failed to get connections API client: {e}")
        return None


def _ensure_catalog(settings) -> CatalogStore:
    """Ensure catalog store is initialized."""
    global _catalog_store
    if _catalog_store is None:
        data_dir = Path(settings.data_dir) / "data_catalog"
        _catalog_store = CatalogStore(data_dir=str(data_dir))
    return _catalog_store


def _ensure_scanner(settings) -> CatalogScanner:
    """Ensure scanner is initialized."""
    global _catalog_scanner
    if _catalog_scanner is None:
        catalog = _ensure_catalog(settings)
        query_func = _get_query_func(settings)
        _catalog_scanner = CatalogScanner(catalog, query_func)
    return _catalog_scanner


# ============================================================================
# catalog_asset action handlers
# ============================================================================

def _asset_create(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    name = kwargs.get("name")
    asset_type = kwargs.get("asset_type")
    if not name or not asset_type:
        return {"error": "name and asset_type are required for 'create' action"}

    try:
        parsed_type = AssetType(asset_type.lower())
    except ValueError:
        return {"error": f"Invalid asset_type: {asset_type}", "valid_types": [t.value for t in AssetType]}

    asset = DataAsset(
        name=name,
        asset_type=parsed_type,
        description=kwargs.get("description"),
        database=kwargs.get("database"),
        schema_name=kwargs.get("schema_name"),
        classification=DataClassification(kwargs.get("classification", "internal").lower()),
        quality_tier=DataQualityTier(kwargs.get("quality_tier", "unknown").lower()),
    )

    if kwargs.get("database") and kwargs.get("schema_name"):
        asset.fully_qualified_name = f"{kwargs['database']}.{kwargs['schema_name']}.{name}"
    elif kwargs.get("database"):
        asset.fully_qualified_name = f"{kwargs['database']}.{name}"

    tags = kwargs.get("tags")
    if tags:
        for tag_name in tags.split(","):
            tag_name = tag_name.strip()
            if tag_name:
                asset.tags.append(Tag(name=tag_name))

    owner_name = kwargs.get("owner_name")
    owner_email = kwargs.get("owner_email")
    if owner_name:
        asset.owners.append(Owner(
            user_id=owner_email or owner_name.lower().replace(" ", "_"),
            name=owner_name,
            email=owner_email,
            role=OwnershipRole.OWNER,
        ))

    catalog.create_asset(asset)
    return {
        "status": "created",
        "asset": {
            "id": asset.id,
            "name": asset.name,
            "type": asset.asset_type.value,
            "fully_qualified_name": asset.fully_qualified_name,
            "classification": asset.classification.value,
            "tags": [t.name for t in asset.tags],
        },
    }


def _asset_get(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    asset_id = kwargs.get("asset_id")
    name = kwargs.get("name")

    if asset_id:
        asset = catalog.get_asset(asset_id)
    elif name:
        asset = catalog.get_asset_by_name(name, database=kwargs.get("database"), schema_name=kwargs.get("schema_name"))
    else:
        return {"error": "Either asset_id or name is required"}

    if not asset:
        return {"error": "Asset not found"}

    return {
        "asset": {
            "id": asset.id,
            "name": asset.name,
            "type": asset.asset_type.value,
            "fully_qualified_name": asset.fully_qualified_name,
            "description": asset.description,
            "database": asset.database,
            "schema": asset.schema_name,
            "classification": asset.classification.value,
            "quality_tier": asset.quality_tier.value,
            "tags": [t.name for t in asset.tags],
            "owners": [{"name": o.name, "email": o.email, "role": o.role.value} for o in asset.owners],
            "columns": [
                {
                    "name": c.column_name, "type": c.data_type, "nullable": c.nullable,
                    "is_pk": c.is_primary_key, "is_fk": c.is_foreign_key,
                    "description": c.description,
                    "classification": c.classification.value if c.classification else None,
                }
                for c in asset.columns
            ],
            "row_count": asset.row_count,
            "size_bytes": asset.size_bytes,
            "quality_metrics": asset.quality_metrics.model_dump() if asset.quality_metrics else None,
            "upstream_count": len(asset.upstream_assets),
            "downstream_count": len(asset.downstream_assets),
            "created_at": asset.created_at.isoformat() if asset.created_at else None,
            "last_scanned_at": asset.last_scanned_at.isoformat() if asset.last_scanned_at else None,
        },
    }


def _asset_update(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    asset_id = kwargs.get("asset_id")
    if not asset_id:
        return {"error": "asset_id is required for 'update' action"}

    updates = {}
    description = kwargs.get("description")
    if description is not None:
        updates["description"] = description

    classification = kwargs.get("classification")
    if classification:
        try:
            updates["classification"] = DataClassification(classification.lower())
        except ValueError:
            return {"error": f"Invalid classification: {classification}"}

    quality_tier = kwargs.get("quality_tier")
    if quality_tier:
        try:
            updates["quality_tier"] = DataQualityTier(quality_tier.lower())
        except ValueError:
            return {"error": f"Invalid quality_tier: {quality_tier}"}

    if updates:
        catalog.update_asset(asset_id, updates)

    add_tags = kwargs.get("add_tags")
    if add_tags:
        for tag_name in add_tags.split(","):
            tag_name = tag_name.strip()
            if tag_name:
                catalog.add_tag_to_asset(asset_id, tag_name)

    remove_tags = kwargs.get("remove_tags")
    if remove_tags:
        for tag_name in remove_tags.split(","):
            tag_name = tag_name.strip()
            if tag_name:
                catalog.remove_tag_from_asset(asset_id, tag_name)

    add_owner = kwargs.get("add_owner")
    if add_owner:
        match = re.match(r"(.+?)\s*<(.+?)>", add_owner)
        if match:
            name, email = match.groups()
        else:
            name = add_owner
            email = None
        asset = catalog.get_asset(asset_id)
        if asset:
            asset.add_owner(Owner(
                user_id=email or name.lower().replace(" ", "_"),
                name=name.strip(),
                email=email,
                role=OwnershipRole.OWNER,
            ))
            catalog.update_asset(asset_id, {"owners": asset.owners})

    asset = catalog.get_asset(asset_id)
    if not asset:
        return {"error": "Asset not found"}

    return {
        "status": "updated",
        "asset": {
            "id": asset.id,
            "name": asset.name,
            "classification": asset.classification.value,
            "quality_tier": asset.quality_tier.value,
            "tags": [t.name for t in asset.tags],
            "owners": [o.name for o in asset.owners],
        },
    }


def _asset_list(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    asset_type = kwargs.get("asset_type")
    tags = kwargs.get("tags")
    classification = kwargs.get("classification")
    quality_tier = kwargs.get("quality_tier")
    limit = kwargs.get("limit", 50)
    offset = kwargs.get("offset", 0)

    parsed_type = AssetType(asset_type.lower()) if asset_type else None
    parsed_tags = [t.strip() for t in tags.split(",")] if tags else None
    parsed_classification = DataClassification(classification.lower()) if classification else None
    parsed_tier = DataQualityTier(quality_tier.lower()) if quality_tier else None

    assets = catalog.list_assets(
        asset_type=parsed_type,
        database=kwargs.get("database"),
        schema_name=kwargs.get("schema_name"),
        tags=parsed_tags,
        classification=parsed_classification,
        quality_tier=parsed_tier,
        owner_id=kwargs.get("owner_id"),
        limit=limit,
        offset=offset,
    )

    return {
        "count": len(assets),
        "offset": offset,
        "limit": limit,
        "assets": [
            {
                "id": a.id, "name": a.name, "type": a.asset_type.value,
                "fully_qualified_name": a.fully_qualified_name,
                "classification": a.classification.value,
                "quality_tier": a.quality_tier.value,
                "tags": [t.name for t in a.tags],
                "row_count": a.row_count,
            }
            for a in assets
        ],
    }


def _asset_delete(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    asset_id = kwargs.get("asset_id")
    if not asset_id:
        return {"error": "asset_id is required for 'delete' action"}

    asset = catalog.get_asset(asset_id)
    if not asset:
        return {"error": "Asset not found"}

    name = asset.name
    if catalog.delete_asset(asset_id):
        return {"status": "deleted", "asset_name": name}
    return {"error": "Failed to delete asset"}


def _asset_search(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    query = kwargs.get("query")
    if not query:
        return {"error": "query is required for 'search' action"}

    asset_types = kwargs.get("asset_types")
    tags = kwargs.get("tags")
    databases = kwargs.get("databases")
    classification = kwargs.get("classification")

    parsed_types = [AssetType(t.strip().lower()) for t in asset_types.split(",")] if asset_types else None
    parsed_tags = [t.strip() for t in tags.split(",")] if tags else None
    parsed_dbs = [d.strip() for d in databases.split(",")] if databases else None
    parsed_class = [DataClassification(classification.lower())] if classification else None

    search_query = SearchQuery(
        query=query,
        asset_types=parsed_types,
        tags=parsed_tags,
        databases=parsed_dbs,
        classifications=parsed_class,
        include_glossary=kwargs.get("include_glossary", True),
        limit=kwargs.get("limit", 20),
    )

    results = catalog.search(search_query)
    return {
        "query": results.query,
        "total_count": results.total_count,
        "took_ms": results.took_ms,
        "results": [
            {
                "id": r.asset_id, "type": r.asset_type.value, "name": r.name,
                "fully_qualified_name": r.fully_qualified_name,
                "description": r.description, "score": round(r.match_score, 3),
                "highlights": r.match_highlights, "tags": r.tags,
            }
            for r in results.results
        ],
    }


def _asset_stats(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    stats = catalog.get_stats()
    return {
        "assets": {
            "total": stats.total_assets,
            "by_type": stats.assets_by_type,
            "by_classification": stats.assets_by_classification,
            "by_quality_tier": stats.assets_by_quality_tier,
            "without_owners": stats.assets_without_owners,
            "without_descriptions": stats.assets_without_descriptions,
        },
        "glossary": {
            "total_terms": stats.total_glossary_terms,
            "by_status": stats.terms_by_status,
            "by_domain": stats.terms_by_domain,
        },
        "tags": {"total": stats.total_tags, "most_used": stats.most_used_tags[:5]},
        "owners": {"total": stats.total_owners},
        "catalog_updated_at": stats.catalog_updated_at.isoformat(),
    }


def _asset_scan_connection(settings, **kwargs) -> Dict[str, Any]:
    scanner = _ensure_scanner(settings)
    connection_id = kwargs.get("connection_id")
    if not connection_id:
        return {"error": "connection_id is required for 'scan_connection' action"}

    config = ScanConfig(
        connection_id=connection_id,
        database=kwargs.get("database"),
        schema_pattern=kwargs.get("schema_pattern"),
        table_pattern=kwargs.get("table_pattern"),
        include_views=kwargs.get("include_views", True),
        include_columns=True,
        profile_columns=kwargs.get("profile_columns", False),
        detect_pii=kwargs.get("detect_pii", True),
    )

    result = scanner.scan_connection(config)
    return {
        "scan_id": result.scan_id,
        "status": result.status,
        "started_at": result.started_at.isoformat(),
        "completed_at": result.completed_at.isoformat() if result.completed_at else None,
        "statistics": {
            "databases_scanned": result.databases_scanned,
            "schemas_scanned": result.schemas_scanned,
            "tables_scanned": result.tables_scanned,
            "columns_scanned": result.columns_scanned,
            "assets_created": result.assets_created,
            "assets_updated": result.assets_updated,
            "pii_columns_detected": result.pii_columns_detected,
        },
        "errors": result.errors,
    }


def _asset_scan_table(settings, **kwargs) -> Dict[str, Any]:
    scanner = _ensure_scanner(settings)
    connection_id = kwargs.get("connection_id")
    database = kwargs.get("database")
    schema_name = kwargs.get("schema_name")
    table_name = kwargs.get("table_name")
    if not all([connection_id, database, schema_name, table_name]):
        return {"error": "connection_id, database, schema_name, and table_name are required for 'scan_table'"}

    asset = scanner.scan_table(
        connection_id=connection_id,
        database=database,
        schema_name=schema_name,
        table_name=table_name,
        profile=kwargs.get("profile", True),
        detect_pii=kwargs.get("detect_pii", True),
    )

    if not asset:
        return {"error": "Table not found or scan failed"}

    return {
        "status": "scanned",
        "asset": {
            "id": asset.id,
            "name": asset.name,
            "fully_qualified_name": asset.fully_qualified_name,
            "column_count": len(asset.columns),
            "row_count": asset.row_count,
            "pii_columns": [
                c.column_name for c in asset.columns
                if c.classification in (DataClassification.PII, DataClassification.PHI, DataClassification.PCI)
            ],
        },
    }


def _asset_refresh(settings, **kwargs) -> Dict[str, Any]:
    scanner = _ensure_scanner(settings)
    asset_id = kwargs.get("asset_id")
    if not asset_id:
        return {"error": "asset_id is required for 'refresh' action"}

    asset = scanner.refresh_asset(asset_id)
    if not asset:
        return {"error": "Asset not found or refresh failed"}

    return {
        "status": "refreshed",
        "asset": {
            "id": asset.id,
            "name": asset.name,
            "column_count": len(asset.columns),
            "row_count": asset.row_count,
            "last_scanned_at": asset.last_scanned_at.isoformat() if asset.last_scanned_at else None,
        },
    }


def _asset_add_tag(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    asset_id = kwargs.get("asset_id")
    tag_name = kwargs.get("tag_name")
    if not asset_id or not tag_name:
        return {"error": "asset_id and tag_name are required for 'add_tag' action"}
    if catalog.add_tag_to_asset(asset_id, tag_name):
        return {"status": "added", "asset_id": asset_id, "tag": tag_name}
    return {"error": "Failed to add tag"}


def _asset_remove_tag(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    asset_id = kwargs.get("asset_id")
    tag_name = kwargs.get("tag_name")
    if not asset_id or not tag_name:
        return {"error": "asset_id and tag_name are required for 'remove_tag' action"}
    if catalog.remove_tag_from_asset(asset_id, tag_name):
        return {"status": "removed", "asset_id": asset_id, "tag": tag_name}
    return {"error": "Failed to remove tag"}


def _asset_list_tags(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    tags = catalog.list_tags(category=kwargs.get("tag_category"))
    return {
        "count": len(tags),
        "tags": [{"name": t.name, "category": t.category, "color": t.color} for t in tags],
    }


def _asset_create_tag(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    tag_name = kwargs.get("tag_name")
    if not tag_name:
        return {"error": "tag_name is required for 'create_tag' action"}
    tag = Tag(name=tag_name, category=kwargs.get("tag_category"))
    catalog.create_tag(tag)
    return {"status": "created", "tag": tag_name}


_ASSET_ACTIONS = {
    "create": _asset_create,
    "get": _asset_get,
    "update": _asset_update,
    "list": _asset_list,
    "delete": _asset_delete,
    "search": _asset_search,
    "stats": _asset_stats,
    "scan_connection": _asset_scan_connection,
    "scan_table": _asset_scan_table,
    "refresh": _asset_refresh,
    "add_tag": _asset_add_tag,
    "remove_tag": _asset_remove_tag,
    "list_tags": _asset_list_tags,
    "create_tag": _asset_create_tag,
}


# ============================================================================
# catalog_term action handlers
# ============================================================================

def _term_create(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    name = kwargs.get("name")
    definition = kwargs.get("definition")
    if not name or not definition:
        return {"error": "name and definition are required for 'create' action"}

    term = GlossaryTerm(
        name=name,
        definition=definition,
        domain=kwargs.get("domain"),
        category=kwargs.get("category"),
        status=TermStatus.DRAFT,
    )

    synonyms = kwargs.get("synonyms")
    if synonyms:
        term.synonyms = [s.strip() for s in synonyms.split(",")]

    examples = kwargs.get("examples")
    if examples:
        term.examples = [e.strip() for e in examples.split(",")]

    owner_name = kwargs.get("owner_name")
    if owner_name:
        term.owner = Owner(
            user_id=owner_name.lower().replace(" ", "_"),
            name=owner_name,
            role=OwnershipRole.OWNER,
        )

    catalog.create_term(term)
    return {
        "status": "created",
        "term": {"id": term.id, "name": term.name, "domain": term.domain, "status": term.status.value},
    }


def _term_get(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    term_id = kwargs.get("term_id")
    name = kwargs.get("name")

    if term_id:
        term = catalog.get_term(term_id)
    elif name:
        term = catalog.get_term_by_name(name)
    else:
        return {"error": "Either term_id or name is required"}

    if not term:
        return {"error": "Term not found"}

    return {
        "term": {
            "id": term.id, "name": term.name, "definition": term.definition,
            "domain": term.domain, "category": term.category, "status": term.status.value,
            "synonyms": term.synonyms, "examples": term.examples,
            "linked_assets": term.linked_asset_ids, "linked_columns": term.linked_column_refs,
            "owner": term.owner.name if term.owner else None,
            "created_at": term.created_at.isoformat() if term.created_at else None,
        },
    }


def _term_list(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    status = kwargs.get("status")
    limit = kwargs.get("limit", 50)
    offset = kwargs.get("offset", 0)

    parsed_status = TermStatus(status.lower()) if status else None
    terms = catalog.list_terms(
        domain=kwargs.get("domain"),
        status=parsed_status,
        limit=limit,
        offset=offset,
    )

    return {
        "count": len(terms),
        "offset": offset,
        "limit": limit,
        "terms": [
            {"id": t.id, "name": t.name, "domain": t.domain, "status": t.status.value, "linked_assets": len(t.linked_asset_ids)}
            for t in terms
        ],
    }


def _term_link(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    term_id = kwargs.get("term_id")
    if not term_id:
        return {"error": "term_id is required for 'link' action"}

    asset_id = kwargs.get("asset_id")
    column_ref = kwargs.get("column_ref")
    if not asset_id and not column_ref:
        return {"error": "Either asset_id or column_ref is required"}

    if asset_id:
        if catalog.link_term_to_asset(term_id, asset_id):
            return {"status": "linked", "term_id": term_id, "asset_id": asset_id}
        return {"error": "Failed to link term to asset"}

    if column_ref:
        if catalog.link_term_to_column(term_id, column_ref):
            return {"status": "linked", "term_id": term_id, "column_ref": column_ref}
        return {"error": "Failed to link term to column"}


_TERM_ACTIONS = {
    "create": _term_create,
    "get": _term_get,
    "list": _term_list,
    "link": _term_link,
}


# ============================================================================
# catalog_lineage action handlers
# ============================================================================

def _lineage_from_sql(settings, **kwargs) -> Dict[str, Any]:
    from .lineage_extractor import SQLLineageExtractor, LineageGraphBuilder

    sql = kwargs.get("sql")
    if not sql:
        return {"error": "sql is required for 'from_sql' action"}

    target_name = kwargs.get("target_name")
    target_type = kwargs.get("target_type", "VIEW")
    graph_name = kwargs.get("graph_name", "catalog_lineage")
    add_to_lineage_graph = kwargs.get("add_to_lineage_graph", True)

    extractor = SQLLineageExtractor()
    extraction = extractor.extract_from_sql(sql, target_name, target_type)

    builder = LineageGraphBuilder()
    graph = builder.build_from_extraction(extraction)
    mermaid = builder.generate_mermaid_diagram(graph)

    lineage_added = False
    if add_to_lineage_graph:
        try:
            from src.lineage import LineageTracker, NodeType, TransformationType

            tracker = LineageTracker()
            lin_graph = tracker.get_or_create_graph(graph_name)

            for node in graph.get("nodes", []):
                try:
                    node_type = NodeType(node.get("node_type", "TABLE"))
                except ValueError:
                    node_type = NodeType.TABLE
                tracker.add_node(graph_name=graph_name, name=node.get("name"), node_type=node_type)

            for edge in graph.get("edges", []):
                try:
                    trans_type = TransformationType(edge.get("transformation_type", "DIRECT"))
                except ValueError:
                    trans_type = TransformationType.DIRECT
                tracker.add_edge(
                    graph_name=graph_name,
                    source_node=edge.get("source"),
                    target_node=edge.get("target"),
                    transformation_type=trans_type,
                )

            for col_lin in graph.get("column_lineage", []):
                if col_lin.get("source_node") and col_lin.get("target_node"):
                    try:
                        trans_type = TransformationType(col_lin.get("transformation_type", "DIRECT"))
                    except ValueError:
                        trans_type = TransformationType.DIRECT
                    tracker.add_column_lineage(
                        graph_name=graph_name,
                        source_node=col_lin.get("source_node"),
                        source_columns=col_lin.get("source_columns", []),
                        target_node=col_lin.get("target_node"),
                        target_column=col_lin.get("target_column"),
                        transformation_type=trans_type,
                    )
            lineage_added = True
        except ImportError:
            logger.warning("Lineage module not available")
        except Exception as e:
            logger.warning(f"Failed to add to lineage graph: {e}")

    return {
        "status": "extracted",
        "target": extraction.get("target"),
        "source_count": len(extraction.get("sources", [])),
        "sources": extraction.get("sources", []),
        "transformation_types": [t.get("type") for t in extraction.get("transformations", [])],
        "column_lineage_count": len(extraction.get("column_lineage", [])),
        "column_lineage": extraction.get("column_lineage", [])[:10],
        "lineage_graph_name": graph_name if lineage_added else None,
        "mermaid_diagram": mermaid,
    }


def _lineage_from_dbt(settings, **kwargs) -> Dict[str, Any]:
    from .lineage_extractor import DbtLineageExtractor, LineageGraphBuilder

    manifest_path = kwargs.get("manifest_path")
    if not manifest_path:
        return {"error": "manifest_path is required for 'from_dbt' action"}

    model_name = kwargs.get("model_name")
    graph_name = kwargs.get("graph_name", "dbt_lineage")

    extractor = DbtLineageExtractor()

    if model_name:
        extraction = extractor.extract_column_lineage_from_model(manifest_path, model_name)
        if "error" in extraction:
            return extraction

        builder = LineageGraphBuilder()
        graph = builder.build_from_extraction(extraction)
        mermaid = builder.generate_mermaid_diagram(graph)

        return {
            "status": "extracted",
            "model_name": model_name,
            "target": extraction.get("target"),
            "source_count": len(extraction.get("sources", [])),
            "sources": extraction.get("sources", []),
            "column_lineage_count": len(extraction.get("column_lineage", [])),
            "column_lineage": extraction.get("column_lineage", [])[:10],
            "mermaid_diagram": mermaid,
        }
    else:
        manifest_data = extractor.extract_from_manifest(manifest_path)
        return {
            "status": "extracted",
            "project_name": manifest_data.get("project_name"),
            "dbt_version": manifest_data.get("dbt_version"),
            "model_count": len(manifest_data.get("nodes", [])),
            "source_count": len(manifest_data.get("sources", [])),
            "edge_count": len(manifest_data.get("edges", [])),
            "models": [
                {"name": n.get("name"), "materialized": n.get("materialized"), "depends_on_count": len(n.get("depends_on", []))}
                for n in manifest_data.get("nodes", [])
            ],
            "sources": [
                {"name": s.get("name"), "schema": s.get("schema"), "database": s.get("database")}
                for s in manifest_data.get("sources", [])
            ],
        }


def _lineage_visualize(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    asset_id = kwargs.get("asset_id")
    asset_name = kwargs.get("asset_name")

    if asset_id:
        asset = catalog.get_asset(asset_id)
    elif asset_name:
        asset = catalog.get_asset_by_name(asset_name)
    else:
        return {"error": "Either asset_id or asset_name is required"}

    if not asset:
        return {"error": "Asset not found"}

    from .lineage_extractor import LineageGraphBuilder

    direction = kwargs.get("direction", "both")
    fmt = kwargs.get("format", "mermaid")

    builder = LineageGraphBuilder()
    lineage = catalog.get_asset_lineage(asset.id, direction)

    nodes = [{"name": asset.name, "node_type": asset.asset_type.value.upper()}]
    edges = []

    for up_asset in lineage.get("upstream", []):
        nodes.append({"name": up_asset.name, "node_type": up_asset.asset_type.value.upper()})
        edges.append({"source": up_asset.name, "target": asset.name, "transformation_type": "DIRECT"})

    for down_asset in lineage.get("downstream", []):
        nodes.append({"name": down_asset.name, "node_type": down_asset.asset_type.value.upper()})
        edges.append({"source": asset.name, "target": down_asset.name, "transformation_type": "DIRECT"})

    graph = {"nodes": nodes, "edges": edges}

    if fmt == "mermaid":
        diagram = builder.generate_mermaid_diagram(graph, "TD")
    else:
        lines = ['digraph G {', '    rankdir=TB;', '    node [shape=box];']
        for node in nodes:
            lines.append(f'    "{node["name"]}" [label="{node["name"]}"];')
        for edge in edges:
            lines.append(f'    "{edge["source"]}" -> "{edge["target"]}";')
        lines.append("}")
        diagram = "\n".join(lines)

    return {
        "asset": {"id": asset.id, "name": asset.name, "type": asset.asset_type.value},
        "direction": direction,
        "upstream_count": len(lineage.get("upstream", [])),
        "downstream_count": len(lineage.get("downstream", [])),
        "upstream": [{"id": a.id, "name": a.name, "type": a.asset_type.value} for a in lineage.get("upstream", [])],
        "downstream": [{"id": a.id, "name": a.name, "type": a.asset_type.value} for a in lineage.get("downstream", [])],
        "diagram_format": fmt,
        "diagram": diagram,
    }


def _lineage_impact(settings, **kwargs) -> Dict[str, Any]:
    catalog = _ensure_catalog(settings)
    asset_id = kwargs.get("asset_id")
    asset_name = kwargs.get("asset_name")

    if asset_id:
        asset = catalog.get_asset(asset_id)
    elif asset_name:
        asset = catalog.get_asset_by_name(asset_name)
    else:
        return {"error": "Either asset_id or asset_name is required"}

    if not asset:
        return {"error": "Asset not found"}

    change_type = kwargs.get("change_type", "MODIFY")
    column_name = kwargs.get("column_name")

    lineage = catalog.get_asset_lineage(asset.id, "downstream")
    downstream = lineage.get("downstream", [])

    impacted = []
    severity_counts = {"CRITICAL": 0, "HIGH": 0, "MEDIUM": 0, "LOW": 0}

    for idx, down_asset in enumerate(downstream):
        distance = idx + 1
        if down_asset.asset_type.value in ("data_mart", "report", "dashboard"):
            severity = "CRITICAL"
        elif down_asset.asset_type.value in ("view", "dbt_model"):
            severity = "HIGH" if distance <= 2 else "MEDIUM"
        else:
            severity = "MEDIUM" if distance <= 2 else "LOW"

        if change_type.upper() == "REMOVE" and column_name:
            if severity == "MEDIUM":
                severity = "HIGH"
            elif severity == "LOW":
                severity = "MEDIUM"

        severity_counts[severity] += 1
        impacted.append({
            "asset_id": down_asset.id,
            "asset_name": down_asset.name,
            "asset_type": down_asset.asset_type.value,
            "severity": severity,
            "distance": distance,
            "description": f"{'Column ' + column_name + ' in ' if column_name else ''}{asset.name} impacts {down_asset.name}",
        })

    if severity_counts["CRITICAL"] > 0:
        overall_severity = "CRITICAL"
    elif severity_counts["HIGH"] > 0:
        overall_severity = "HIGH"
    elif severity_counts["MEDIUM"] > 0:
        overall_severity = "MEDIUM"
    else:
        overall_severity = "LOW"

    return {
        "source_asset": {"id": asset.id, "name": asset.name, "type": asset.asset_type.value},
        "change_type": change_type,
        "column": column_name,
        "overall_severity": overall_severity,
        "total_impacted": len(impacted),
        "severity_breakdown": severity_counts,
        "impacted_assets": impacted,
        "recommendation": (
            "Review all CRITICAL and HIGH severity assets before making changes. "
            "Consider updating dependent views/models to handle the change."
            if impacted else "No downstream assets found. Change is safe to proceed."
        ),
    }


_LINEAGE_ACTIONS = {
    "from_sql": _lineage_from_sql,
    "from_dbt": _lineage_from_dbt,
    "visualize": _lineage_visualize,
    "impact": _lineage_impact,
}


# ============================================================================
# Unified tool dispatch functions (called by deprecation wrappers too)
# ============================================================================

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


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


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


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

def register_unified_catalog_tools(mcp, settings):
    """Register the 3 unified catalog MCP tools."""

    @mcp.tool()
    @safe_tool("catalog_asset")
    def catalog_asset(
        action: str,
        name: Optional[str] = None,
        asset_id: Optional[str] = None,
        asset_type: Optional[str] = None,
        description: Optional[str] = None,
        database: Optional[str] = None,
        schema_name: Optional[str] = None,
        classification: Optional[str] = None,
        quality_tier: Optional[str] = None,
        tags: Optional[str] = None,
        owner_name: Optional[str] = None,
        owner_email: Optional[str] = None,
        add_tags: Optional[str] = None,
        remove_tags: Optional[str] = None,
        add_owner: Optional[str] = None,
        owner_id: Optional[str] = None,
        limit: int = 50,
        offset: int = 0,
        query: Optional[str] = None,
        asset_types: Optional[str] = None,
        databases: Optional[str] = None,
        include_glossary: bool = True,
        connection_id: Optional[str] = None,
        schema_pattern: Optional[str] = None,
        table_pattern: Optional[str] = None,
        include_views: bool = True,
        profile_columns: bool = False,
        profile: bool = True,
        detect_pii: bool = True,
        table_name: Optional[str] = None,
        tag_name: Optional[str] = None,
        tag_category: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Unified catalog asset management tool. Replaces 14 individual tools.

        Actions:
        - create: Create a new data asset (requires name, asset_type)
        - get: Get asset details (requires asset_id or name)
        - update: Update asset metadata (requires asset_id)
        - list: List assets with filters
        - delete: Delete an asset (requires asset_id)
        - search: Search the catalog (requires query)
        - stats: Get catalog statistics
        - scan_connection: Scan a data connection (requires connection_id)
        - scan_table: Scan a single table (requires connection_id, database, schema_name, table_name)
        - refresh: Refresh asset metadata (requires asset_id)
        - add_tag: Add tag to asset (requires asset_id, tag_name)
        - remove_tag: Remove tag from asset (requires asset_id, tag_name)
        - list_tags: List all tags
        - create_tag: Create a tag definition (requires tag_name)

        Args:
            action: The action to perform (see list above)
            name: Asset name (for create/get)
            asset_id: Asset ID (for get/update/delete/refresh/add_tag/remove_tag)
            asset_type: Asset type (for create/list)
            description: Description (for create/update)
            database: Database name
            schema_name: Schema name
            classification: Data classification
            quality_tier: Quality tier
            tags: Comma-separated tags (for create/list filter)
            owner_name: Owner name (for create)
            owner_email: Owner email (for create)
            add_tags: Tags to add (for update)
            remove_tags: Tags to remove (for update)
            add_owner: Owner to add (for update, format: "Name <email>")
            owner_id: Filter by owner (for list)
            limit: Max results (for list/search)
            offset: Pagination offset (for list)
            query: Search text (for search)
            asset_types: Comma-separated types (for search)
            databases: Comma-separated databases (for search)
            include_glossary: Include glossary in search
            connection_id: Connection ID (for scan_connection/scan_table)
            schema_pattern: Schema pattern (for scan_connection)
            table_pattern: Table pattern (for scan_connection)
            include_views: Include views in scan
            profile_columns: Profile columns during scan
            profile: Profile during table scan
            detect_pii: Detect PII columns
            table_name: Table name (for scan_table)
            tag_name: Tag name (for add_tag/remove_tag/create_tag)
            tag_category: Tag category (for list_tags/create_tag)

        Returns:
            Action-specific result dict
        """
        return dispatch_catalog_asset(settings, action, **{
            k: v for k, v in {
                "name": name, "asset_id": asset_id, "asset_type": asset_type,
                "description": description, "database": database, "schema_name": schema_name,
                "classification": classification, "quality_tier": quality_tier, "tags": tags,
                "owner_name": owner_name, "owner_email": owner_email, "add_tags": add_tags,
                "remove_tags": remove_tags, "add_owner": add_owner, "owner_id": owner_id,
                "limit": limit, "offset": offset, "query": query, "asset_types": asset_types,
                "databases": databases, "include_glossary": include_glossary,
                "connection_id": connection_id, "schema_pattern": schema_pattern,
                "table_pattern": table_pattern, "include_views": include_views,
                "profile_columns": profile_columns, "profile": profile, "detect_pii": detect_pii,
                "table_name": table_name, "tag_name": tag_name, "tag_category": tag_category,
            }.items() if v is not None or k in ("description",)
        })

    @mcp.tool()
    @safe_tool("catalog_term")
    def catalog_term(
        action: str,
        name: Optional[str] = None,
        term_id: Optional[str] = None,
        definition: Optional[str] = None,
        domain: Optional[str] = None,
        category: Optional[str] = None,
        synonyms: Optional[str] = None,
        examples: Optional[str] = None,
        owner_name: Optional[str] = None,
        status: Optional[str] = None,
        asset_id: Optional[str] = None,
        column_ref: Optional[str] = None,
        limit: int = 50,
        offset: int = 0,
    ) -> Dict[str, Any]:
        """
        Unified glossary term management tool. Replaces 4 individual tools.

        Actions:
        - create: Create a glossary term (requires name, definition)
        - get: Get term details (requires term_id or name)
        - list: List terms with filters
        - link: Link term to asset/column (requires term_id + asset_id or column_ref)

        Args:
            action: The action to perform
            name: Term name (for create/get)
            term_id: Term ID (for get/link)
            definition: Business definition (for create)
            domain: Business domain (for create/list filter)
            category: Term category (for create)
            synonyms: Comma-separated synonyms (for create)
            examples: Comma-separated examples (for create)
            owner_name: Term owner (for create)
            status: Filter by status (for list: draft, pending_review, approved, deprecated)
            asset_id: Asset ID to link (for link)
            column_ref: Column reference to link (for link)
            limit: Max results (for list)
            offset: Pagination offset (for list)

        Returns:
            Action-specific result dict
        """
        return dispatch_catalog_term(settings, action, **{
            k: v for k, v in {
                "name": name, "term_id": term_id, "definition": definition,
                "domain": domain, "category": category, "synonyms": synonyms,
                "examples": examples, "owner_name": owner_name, "status": status,
                "asset_id": asset_id, "column_ref": column_ref, "limit": limit, "offset": offset,
            }.items() if v is not None
        })

    @mcp.tool()
    @safe_tool("catalog_lineage")
    def catalog_lineage(
        action: str,
        sql: Optional[str] = None,
        target_name: Optional[str] = None,
        target_type: str = "VIEW",
        graph_name: Optional[str] = None,
        add_to_lineage_graph: bool = True,
        manifest_path: Optional[str] = None,
        model_name: Optional[str] = None,
        asset_id: Optional[str] = None,
        asset_name: Optional[str] = None,
        direction: str = "both",
        format: str = "mermaid",
        max_depth: int = 5,
        change_type: str = "MODIFY",
        column_name: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Unified lineage extraction and visualization tool. Replaces 4 individual tools.

        Actions:
        - from_sql: Extract lineage from SQL (requires sql)
        - from_dbt: Extract lineage from dbt manifest (requires manifest_path)
        - visualize: Visualize asset lineage (requires asset_id or asset_name)
        - impact: Analyze change impact (requires asset_id or asset_name)

        Args:
            action: The action to perform
            sql: SQL to parse (for from_sql)
            target_name: Target object name (for from_sql)
            target_type: Target type - VIEW, TABLE, MODEL (for from_sql)
            graph_name: Lineage graph name (for from_sql/from_dbt)
            add_to_lineage_graph: Add to lineage module (for from_sql)
            manifest_path: Path to dbt manifest.json (for from_dbt)
            model_name: Specific dbt model (for from_dbt)
            asset_id: Catalog asset ID (for visualize/impact)
            asset_name: Asset name (for visualize/impact)
            direction: Traversal direction (for visualize: upstream, downstream, both)
            format: Diagram format (for visualize: mermaid, dot)
            max_depth: Max traversal depth (for visualize)
            change_type: Type of change (for impact: REMOVE, MODIFY, RENAME)
            column_name: Column being changed (for impact)

        Returns:
            Action-specific result dict
        """
        return dispatch_catalog_lineage(settings, action, **{
            k: v for k, v in {
                "sql": sql, "target_name": target_name, "target_type": target_type,
                "graph_name": graph_name, "add_to_lineage_graph": add_to_lineage_graph,
                "manifest_path": manifest_path, "model_name": model_name,
                "asset_id": asset_id, "asset_name": asset_name, "direction": direction,
                "format": format, "max_depth": max_depth, "change_type": change_type,
                "column_name": column_name,
            }.items() if v is not None
        })

    logger.info("Registered 3 unified catalog tools: catalog_asset, catalog_term, catalog_lineage")
    return ["catalog_asset", "catalog_term", "catalog_lineage"]
