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

Consolidates 9 individual model tools into 2 action-dispatched tools:
- model_discover(mode, ...) — 6 modes for discovery and classification
- model_relationships(action, ...) — 3 actions for relationship management

model_load_dimensional_model stays standalone (heavy cross-account ETL).

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

from __future__ import annotations

import json
import logging
import os
from pathlib import Path
from typing import Any, Dict, List, Optional

from src.utils.tool_wrapper import safe_tool

logger = logging.getLogger(__name__)


# ============================================================================
# Lazy AI Advisor singleton
# ============================================================================

_ai_advisor = None


def _get_ai_advisor():
    """Return a shared AiAdvisor instance (lazy-initialized)."""
    global _ai_advisor
    if _ai_advisor is None:
        from .ai_advisor import AiAdvisor

        # Try Cortex first
        cortex = None
        try:
            from src.cortex_agent.cortex_client import CortexClient
            from src.data_modeling.snowflake_pool import sf_pool
            if sf_pool and sf_pool.default_connection:
                cortex = CortexClient(
                    connection_id=sf_pool.default_connection,
                    query_func=sf_pool.query,
                )
        except Exception:
            logger.debug("Cortex client unavailable for AI advisor")

        # Try Anthropic key
        api_key = os.environ.get("ANTHROPIC_API_KEY")

        _ai_advisor = AiAdvisor(cortex_client=cortex, anthropic_key=api_key)
    return _ai_advisor


# ============================================================================
# model_discover mode handlers
# ============================================================================

def _discover_full(mcp_settings, **kwargs) -> Dict[str, Any]:
    """Full discovery: relationships + classify + dictionary."""
    from .mcp_tools import _discover_data_model
    return _discover_data_model(
        source_type=kwargs.get("source_type", "csv"),
        csv_paths=kwargs.get("csv_paths"),
        csv_folder=kwargs.get("csv_folder"),
        connection_string=kwargs.get("connection_string"),
        database=kwargs.get("database"),
        schema=kwargs.get("schema"),
        tables=kwargs.get("tables"),
        sample_limit=kwargs.get("sample_limit", 200),
        min_overlap=kwargs.get("min_overlap", 0.5),
        persist=kwargs.get("persist", True),
        dbt_project_name=kwargs.get("dbt_project_name"),
        dbt_manifest_path=kwargs.get("dbt_manifest_path"),
        dbt_run=kwargs.get("dbt_run", False),
        dbt_docs=kwargs.get("dbt_docs", False),
        snowflake_connection=kwargs.get("snowflake_connection"),
        ingest_dbt=kwargs.get("ingest_dbt", True),
    )


def _discover_classify(mcp_settings, **kwargs) -> Dict[str, Any]:
    """Classify columns and build data dictionary."""
    from .sample_loader import load_csv_samples, load_csv_folder
    from .dictionary_builder import classify_columns, generate_dictionary
    from .mcp_tools import _ensure_catalog, _persist_to_catalog, _persist_to_kb

    csv_paths = kwargs.get("csv_paths")
    csv_folder = kwargs.get("csv_folder")
    database = kwargs.get("database")
    schema = kwargs.get("schema")
    persist = kwargs.get("persist", True)

    sample_tables: Dict[str, Dict[str, List[dict]]] = {}
    if csv_folder:
        sample_tables = load_csv_folder(csv_folder, limit=200)
    elif csv_paths:
        paths = [p.strip() for p in csv_paths.split(",") if p.strip()]
        sample_tables = load_csv_samples(paths, limit=200)
    if not sample_tables:
        return {"error": "No samples loaded"}

    dictionary: Dict[str, Dict[str, Dict[str, str]]] = {}
    for table, data in sample_tables.items():
        columns = [{"name": c, "data_type": ""} for c in data.get("columns", [])]
        samples = {c: [r.get(c) for r in data.get("rows", [])][:5] for c in data.get("columns", [])}
        classifications = classify_columns(columns, samples, database=database, schema=schema)
        dictionary[table] = generate_dictionary(table, columns, samples, classifications, database=database, schema=schema)

    if persist:
        from src.knowledge_base import KBIngestionPipeline
        catalog = _ensure_catalog()
        pipeline = KBIngestionPipeline()
        _persist_to_catalog(catalog, dictionary, [], database, schema)
        _persist_to_kb(pipeline, dictionary, [])

    return {"dictionary_tables": list(dictionary.keys())}


def _discover_detect_dimensions(mcp_settings, **kwargs) -> Dict[str, Any]:
    """Detect dimension vs fact tables using Kimball heuristics."""
    from .dimension_detector import DimensionDetector

    relationships_json = kwargs.get("relationships_json")
    if not relationships_json:
        return {"error": "relationships_json is required for 'detect_dimensions' mode"}

    # threshold param renamed to avoid conflict with bus_matrix
    threshold = kwargs.get("threshold", 3)
    ai_enhance = kwargs.get("ai_enhance", False)

    try:
        relationships = json.loads(relationships_json)
    except json.JSONDecodeError:
        return {"error": "Invalid JSON"}

    advisor = _get_ai_advisor() if ai_enhance else None

    detector = DimensionDetector()
    candidates = detector.detect(relationships, threshold=threshold)
    classification = detector.classify_tables(
        relationships, threshold=threshold, ai_advisor=advisor,
    )

    result = {
        "dimension_candidates": [
            {"table": c.table, "inbound_references": c.inbound_references, "referencing_tables": c.referencing_tables}
            for c in candidates
        ],
        "classification": classification,
        "summary": {
            "dimensions": sum(1 for v in classification.values() if v == "dimension"),
            "facts": sum(1 for v in classification.values() if v == "fact"),
            "bridges": sum(1 for v in classification.values() if v == "bridge"),
            "unknown": sum(1 for v in classification.values() if v == "unknown"),
        },
    }

    return result


def _discover_bus_matrix(mcp_settings, **kwargs) -> Dict[str, Any]:
    """Generate Kimball bus matrix."""
    from .bus_matrix import BusMatrixGenerator

    relationships_json = kwargs.get("relationships_json")
    classification_json = kwargs.get("classification_json")
    if not relationships_json or not classification_json:
        return {"error": "relationships_json and classification_json required for 'bus_matrix' mode"}

    # bus_matrix_threshold to avoid conflict with detect_dimensions threshold
    threshold = kwargs.get("bus_matrix_threshold", 0.0)
    output_format = kwargs.get("output_format", "json")

    try:
        relationships = json.loads(relationships_json)
        classification = json.loads(classification_json)
    except json.JSONDecodeError:
        return {"error": "Invalid JSON input"}

    gen = BusMatrixGenerator()
    result = gen.generate(relationships, classification, threshold=threshold)

    if output_format == "markdown":
        result["rendered"] = gen.to_markdown(result)
    elif output_format == "html":
        result["rendered"] = gen.to_html(result)
        result["css"] = gen.bus_matrix_css()

    return result


def _discover_auto_specs(mcp_settings, **kwargs) -> Dict[str, Any]:
    """Auto-generate DM table specs from discovery metadata."""
    from .spec_generator import DMSpecGenerator

    relationships_json = kwargs.get("relationships_json")
    classification_json = kwargs.get("classification_json")
    if not relationships_json or not classification_json:
        return {"error": "relationships_json and classification_json required for 'auto_specs' mode"}

    column_metadata_json = kwargs.get("column_metadata_json")
    erp_type = kwargs.get("erp_type", "UNKNOWN")
    run_id = kwargs.get("run_id", "GENERATED")
    output_path = kwargs.get("output_path")
    ai_enhance = kwargs.get("ai_enhance", False)

    try:
        relationships = json.loads(relationships_json)
        classification = json.loads(classification_json)
        column_metadata = json.loads(column_metadata_json) if column_metadata_json else None
    except json.JSONDecodeError:
        return {"error": "Invalid JSON input"}

    advisor = _get_ai_advisor() if ai_enhance else None

    generator = DMSpecGenerator()
    specs = generator.generate(
        relationships=relationships,
        classification=classification,
        column_metadata=column_metadata,
        erp_type=erp_type,
        run_id=run_id,
        ai_advisor=advisor,
    )

    result: Dict[str, Any] = {
        "specs_generated": len(specs),
        "dimensions": sum(1 for s in specs if not s.is_transaction),
        "facts": sum(1 for s in specs if s.is_transaction),
        "specs": [s.to_dict() for s in specs],
    }

    if generator.last_ai_enrichment:
        result["_ai_advisor"] = generator.last_ai_enrichment

    if output_path:
        generator.to_json(specs, output_path)
        result["output_path"] = output_path

    return result


def _discover_configure(mcp_settings, **kwargs) -> Dict[str, Any]:
    """Get or set inference config for relationship detection."""
    from .relationship_inferer import set_default_config, get_default_config
    from .focus_config import InferenceConfig
    from .erp_configs.registry import get_erp_config, list_erp_configs

    erp_preset = kwargs.get("erp_preset")
    key_suffixes = kwargs.get("key_suffixes")
    strip_prefixes = kwargs.get("strip_prefixes")
    exclude_columns = kwargs.get("exclude_columns")
    # config_min_overlap to avoid conflict with other modes
    min_overlap = kwargs.get("config_min_overlap") or kwargs.get("min_overlap")

    current = get_default_config()

    if all(v is None for v in [erp_preset, key_suffixes, strip_prefixes, exclude_columns, min_overlap]):
        return {
            "mode": "read",
            "config": {
                "key_suffixes": current.key_suffixes,
                "strip_prefixes": current.strip_prefixes,
                "exclude_columns": current.exclude_columns,
                "min_overlap": current.min_overlap,
            },
            "available_presets": [e["erp"] for e in list_erp_configs()],
        }

    if erp_preset:
        entry = get_erp_config(erp_preset)
        if not entry:
            return {"error": f"Unknown ERP preset: {erp_preset}. Available: {[e['erp'] for e in list_erp_configs()]}"}
        set_default_config(entry.inference)
        return {
            "mode": "preset_applied",
            "erp": erp_preset,
            "config": {
                "key_suffixes": entry.inference.key_suffixes,
                "strip_prefixes": entry.inference.strip_prefixes,
                "exclude_columns": entry.inference.exclude_columns,
                "min_overlap": entry.inference.min_overlap,
            },
            "focus": {
                "patterns": entry.focus.patterns,
                "sample_limit": entry.focus.sample_limit,
                "overlap_threshold": entry.focus.overlap_threshold,
            },
        }

    new_cfg = InferenceConfig(
        key_suffixes=[s.strip() for s in key_suffixes.split(",")] if key_suffixes else current.key_suffixes,
        strip_prefixes=[s.strip() for s in strip_prefixes.split(",")] if strip_prefixes else current.strip_prefixes,
        exclude_columns=[s.strip() for s in exclude_columns.split(",")] if exclude_columns else current.exclude_columns,
        min_overlap=float(min_overlap) if min_overlap is not None else current.min_overlap,
    )
    set_default_config(new_cfg)
    return {
        "mode": "updated",
        "config": {
            "key_suffixes": new_cfg.key_suffixes,
            "strip_prefixes": new_cfg.strip_prefixes,
            "exclude_columns": new_cfg.exclude_columns,
            "min_overlap": new_cfg.min_overlap,
        },
    }


_DISCOVER_MODES = {
    "full": _discover_full,
    "classify": _discover_classify,
    "detect_dimensions": _discover_detect_dimensions,
    "bus_matrix": _discover_bus_matrix,
    "auto_specs": _discover_auto_specs,
    "configure": _discover_configure,
}


# ============================================================================
# model_relationships action handlers
# ============================================================================

def _relationships_infer_data(mcp_settings, **kwargs) -> Dict[str, Any]:
    """Infer relationships from CSV sample data."""
    from .sample_loader import load_csv_samples, load_csv_folder
    from .relationship_inferer import infer_relationships

    csv_paths = kwargs.get("csv_paths")
    csv_folder = kwargs.get("csv_folder")
    min_overlap = kwargs.get("min_overlap", 0.5)

    sample_tables: Dict[str, Dict[str, List[dict]]] = {}
    if csv_folder:
        sample_tables = load_csv_folder(csv_folder, limit=200)
    elif csv_paths:
        paths = [p.strip() for p in csv_paths.split(",") if p.strip()]
        sample_tables = load_csv_samples(paths, limit=200)
    if not sample_tables:
        return {"error": "No samples loaded"}

    relationships = infer_relationships(sample_tables, min_overlap=min_overlap)
    return {"relationship_count": len(relationships), "relationships": relationships}


def _relationships_infer_name(mcp_settings, **kwargs) -> Dict[str, Any]:
    """Infer FK relationships from column/table names (no data needed)."""
    from .column_fk_inferer import infer_relationships_by_name as _infer_by_name
    from .erp_configs.registry import get_erp_config as _get_erp
    from .relationship_inferer import get_default_config

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

    try:
        table_columns = json.loads(table_columns_json)
    except json.JSONDecodeError:
        return {"error": "Invalid JSON"}

    erp_preset = kwargs.get("erp_preset")
    min_score = kwargs.get("min_score", 2)
    ai_enhance = kwargs.get("ai_enhance", False)

    cfg = None
    if erp_preset:
        entry = _get_erp(erp_preset)
        if entry:
            cfg = entry.inference
    if cfg is None:
        cfg = get_default_config()

    advisor = _get_ai_advisor() if ai_enhance else None

    relationships = _infer_by_name(
        table_columns, config=cfg, min_score=min_score, ai_advisor=advisor,
    )
    return {
        "relationship_count": len(relationships),
        "relationships": relationships,
        "tables_analyzed": len(table_columns),
        "method": "column_name",
    }


def _relationships_persist(mcp_settings, **kwargs) -> Dict[str, Any]:
    """Persist relationships into KB graph and catalog."""
    from .mcp_tools import _ensure_catalog, _persist_to_catalog, _persist_to_kb
    from src.knowledge_base import KBIngestionPipeline

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

    try:
        relationships = json.loads(relationships_json)
    except json.JSONDecodeError:
        return {"error": "Invalid JSON"}

    catalog = _ensure_catalog()
    pipeline = KBIngestionPipeline()
    _persist_to_catalog(catalog, {}, relationships, None, None)
    _persist_to_kb(pipeline, {}, relationships)
    return {"status": "persisted", "relationship_count": len(relationships)}


_RELATIONSHIP_ACTIONS = {
    "infer_data": _relationships_infer_data,
    "infer_name": _relationships_infer_name,
    "persist": _relationships_persist,
}


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

def dispatch_model_discover(mcp_settings, mode: str, **kwargs) -> Dict[str, Any]:
    """Dispatch a model_discover mode."""
    handler = _DISCOVER_MODES.get(mode)
    if not handler:
        return {
            "error": f"Unknown mode: '{mode}'",
            "valid_modes": sorted(_DISCOVER_MODES.keys()),
        }
    try:
        return handler(mcp_settings, **kwargs)
    except ValueError as e:
        return {"error": str(e)}
    except Exception as e:
        logger.error(f"model_discover({mode}) failed: {e}")
        return {"error": f"model_discover({mode}) failed: {e}"}


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


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

def register_unified_modeling_tools(mcp, mcp_settings):
    """Register the 2 unified modeling MCP tools."""

    @mcp.tool()
    @safe_tool("model_discover")
    def model_discover(
        mode: str,
        source_type: Optional[str] = None,
        csv_paths: Optional[str] = None,
        csv_folder: Optional[str] = None,
        connection_string: Optional[str] = None,
        database: Optional[str] = None,
        schema: Optional[str] = None,
        tables: Optional[str] = None,
        sample_limit: int = 200,
        min_overlap: float = 0.5,
        persist: bool = True,
        dbt_project_name: Optional[str] = None,
        dbt_manifest_path: Optional[str] = None,
        dbt_run: bool = False,
        dbt_docs: bool = False,
        snowflake_connection: Optional[str] = None,
        ingest_dbt: bool = True,
        relationships_json: Optional[str] = None,
        classification_json: Optional[str] = None,
        column_metadata_json: Optional[str] = None,
        threshold: int = 3,
        bus_matrix_threshold: float = 0.0,
        output_format: str = "json",
        erp_type: str = "UNKNOWN",
        erp_preset: Optional[str] = None,
        run_id: str = "GENERATED",
        output_path: Optional[str] = None,
        key_suffixes: Optional[str] = None,
        strip_prefixes: Optional[str] = None,
        exclude_columns: Optional[str] = None,
        config_min_overlap: Optional[float] = None,
        ai_enhance: bool = False,
    ) -> Dict[str, Any]:
        """
        Unified data model discovery and classification tool. Replaces 6 individual tools.

        Modes:
        - full: Full discovery — relationships + classify + dictionary (requires source_type)
        - classify: Classify columns and build dictionary (requires csv_paths or csv_folder)
        - detect_dimensions: Detect dim/fact via Kimball heuristics (requires relationships_json)
        - bus_matrix: Generate Kimball bus matrix (requires relationships_json, classification_json)
        - auto_specs: Auto-generate DM table specs (requires relationships_json, classification_json)
        - configure: Get/set inference config for relationship detection

        Args:
            mode: The discovery mode (see list above)
            source_type: Data source type - csv, snowflake, sql (for full)
            csv_paths: Comma-separated CSV paths (for full/classify)
            csv_folder: Folder of CSVs (for full/classify)
            connection_string: SQL connection string (for full)
            database: Database name (for full/classify)
            schema: Schema name (for full/classify)
            tables: Comma-separated table names (for full)
            sample_limit: Sample row limit (for full)
            min_overlap: Minimum FK overlap threshold (for full)
            persist: Persist to catalog/KB (for full/classify)
            dbt_project_name: dbt project (for full)
            dbt_manifest_path: dbt manifest path (for full)
            dbt_run: Run dbt after generating (for full)
            dbt_docs: Generate dbt docs (for full)
            snowflake_connection: Snowflake connection name (for full)
            ingest_dbt: Ingest dbt artifacts (for full)
            relationships_json: JSON relationships (for detect_dimensions/bus_matrix/auto_specs)
            classification_json: JSON table classifications (for bus_matrix/auto_specs)
            column_metadata_json: JSON column metadata (for auto_specs)
            threshold: Min inbound refs for dimension classification (for detect_dimensions)
            bus_matrix_threshold: Min overlap for bus matrix (for bus_matrix)
            output_format: Output format - json, markdown, html (for bus_matrix)
            erp_type: ERP system name (for auto_specs)
            erp_preset: ERP config preset name (for configure)
            run_id: Run identifier (for auto_specs)
            output_path: Output file path (for auto_specs)
            key_suffixes: FK column suffixes (for configure)
            strip_prefixes: Prefixes to strip (for configure)
            exclude_columns: Columns to exclude (for configure)
            config_min_overlap: Default overlap threshold (for configure)
            ai_enhance: Enable optional AI advisor enrichment (for detect_dimensions/auto_specs)

        Returns:
            Mode-specific result dict
        """
        return dispatch_model_discover(mcp_settings, mode, **{
            k: v for k, v in {
                "source_type": source_type, "csv_paths": csv_paths, "csv_folder": csv_folder,
                "connection_string": connection_string, "database": database, "schema": schema,
                "tables": tables, "sample_limit": sample_limit, "min_overlap": min_overlap,
                "persist": persist, "dbt_project_name": dbt_project_name,
                "dbt_manifest_path": dbt_manifest_path, "dbt_run": dbt_run, "dbt_docs": dbt_docs,
                "snowflake_connection": snowflake_connection, "ingest_dbt": ingest_dbt,
                "relationships_json": relationships_json, "classification_json": classification_json,
                "column_metadata_json": column_metadata_json, "threshold": threshold,
                "bus_matrix_threshold": bus_matrix_threshold, "output_format": output_format,
                "erp_type": erp_type, "erp_preset": erp_preset, "run_id": run_id,
                "output_path": output_path, "key_suffixes": key_suffixes,
                "strip_prefixes": strip_prefixes, "exclude_columns": exclude_columns,
                "config_min_overlap": config_min_overlap, "ai_enhance": ai_enhance,
            }.items() if v is not None
        })

    @mcp.tool()
    @safe_tool("model_relationships")
    def model_relationships(
        action: str,
        csv_paths: Optional[str] = None,
        csv_folder: Optional[str] = None,
        min_overlap: float = 0.5,
        table_columns_json: Optional[str] = None,
        erp_preset: Optional[str] = None,
        min_score: int = 2,
        relationships_json: Optional[str] = None,
        ai_enhance: bool = False,
    ) -> Dict[str, Any]:
        """
        Unified relationship management tool. Replaces 3 individual tools.

        Actions:
        - infer_data: Infer relationships from CSV sample data (requires csv_paths or csv_folder)
        - infer_name: Infer FK relationships from column names only (requires table_columns_json)
        - persist: Persist relationships to KB and catalog (requires relationships_json)

        Args:
            action: The action to perform (see list above)
            csv_paths: Comma-separated CSV paths (for infer_data)
            csv_folder: CSV folder (for infer_data)
            min_overlap: Minimum overlap threshold (for infer_data)
            table_columns_json: JSON dict of {table: [columns]} (for infer_name)
            erp_preset: ERP preset name (for infer_name)
            min_score: Min match score (for infer_name)
            relationships_json: JSON relationships to persist (for persist)
            ai_enhance: Enable optional AI advisor enrichment (for infer_name)

        Returns:
            Action-specific result dict
        """
        return dispatch_model_relationships(mcp_settings, action, **{
            k: v for k, v in {
                "csv_paths": csv_paths, "csv_folder": csv_folder, "min_overlap": min_overlap,
                "table_columns_json": table_columns_json, "erp_preset": erp_preset,
                "min_score": min_score, "relationships_json": relationships_json,
                "ai_enhance": ai_enhance,
            }.items() if v is not None
        })

    logger.info("Registered 2 unified modeling tools: model_discover, model_relationships")
    return ["model_discover", "model_relationships"]
