"""MCP tools for data modeling discovery and dictionary generation.

Provides 3 MCP tools:
- 2 unified tools: model_discover, model_relationships (Phase 3A)
- 1 standalone: model_load_dimensional_model (unchanged -- heavy cross-account ETL)
"""
from __future__ import annotations

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

from src.config import settings
from src.data_catalog.catalog_store import CatalogStore
from src.data_catalog.types import AssetType, ColumnProfile, DataAsset
from src.knowledge_base import KBIngestionPipeline
from src.knowledge_base.types import GraphNode, GraphEdge, VectorMetadata
from src.dbt_integration import DbtProjectGenerator, run_dbt_command, load_run_results, ensure_profiles_yml

from .sample_loader import (
    load_csv_samples,
    load_csv_folder,
    load_snowflake_samples,
    load_sqlalchemy_samples,
)
from .relationship_inferer import infer_relationships, set_default_config, get_default_config
from .dictionary_builder import classify_columns, generate_dictionary, to_catalog_classification
from .dbt_bridge import generate_dbt_assets, ingest_dbt_artifacts
from .snowflake_utils import ensure_snowflake_env_from_cli
from .dimension_detector import DimensionDetector
from .focus_config import InferenceConfig, FocusConfig
from .erp_configs.registry import detect_erp_type, get_erp_config

logger = logging.getLogger(__name__)


# ============================================================================
# Shared helpers (used by both unified.py handlers and legacy wrappers)
# ============================================================================

def _ensure_catalog() -> CatalogStore:
    data_dir = Path(settings.data_dir) / "data_catalog"
    return CatalogStore(data_dir=str(data_dir))


def _save_run(payload: Dict[str, Any]) -> Path:
    run_dir = Path(settings.data_dir) / "data_modeling"
    run_dir.mkdir(parents=True, exist_ok=True)
    path = run_dir / "latest_run.json"
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    return path


def _persist_to_kb(
    pipeline: KBIngestionPipeline,
    dictionary: Dict[str, Dict[str, Dict[str, str]]],
    relationships: List[dict],
) -> None:
    nodes: List[GraphNode] = []
    edges: List[GraphEdge] = []
    vectors: List[VectorMetadata] = []

    for table, cols in dictionary.items():
        table_id = f"TABLE:{table}"
        nodes.append(GraphNode(id=table_id, type="Table", properties={"name": table}))
        for col_name, meta in cols.items():
            col_id = f"COLUMN:{table}.{col_name}"
            nodes.append(GraphNode(id=col_id, type="Column", properties={"name": col_name}))
            edges.append(GraphEdge(source=table_id, target=col_id, type="HAS_COLUMN"))
            vectors.append(VectorMetadata(
                id=col_id,
                source_type="data_dictionary",
                source_id=f"{table}.{col_name}",
                tags=[meta.get("classification", "internal")],
                extra={"description": meta.get("description", "")},
            ))

    for rel in relationships:
        src_table = rel["source_table"]
        tgt_table = rel["target_table"]
        edges.append(GraphEdge(
            source=f"TABLE:{src_table}",
            target=f"TABLE:{tgt_table}",
            type="FOREIGN_KEY",
            properties=rel,
        ))

    if nodes:
        pipeline.ingest_graph(nodes, edges)
    if vectors:
        pipeline.ingest_vector_metadata(vectors)


def _persist_to_catalog(
    catalog: CatalogStore,
    dictionary: Dict[str, Dict[str, Dict[str, str]]],
    relationships: List[dict],
    database: Optional[str],
    schema: Optional[str],
) -> None:
    asset_map: Dict[str, DataAsset] = {}

    for table, cols in dictionary.items():
        asset = catalog.get_asset_by_name(table, AssetType.TABLE, database=database, schema_name=schema)
        if not asset:
            asset = DataAsset(
                name=table,
                asset_type=AssetType.TABLE,
                database=database,
                schema_name=schema,
                fully_qualified_name=".".join([p for p in [database, schema, table] if p]),
            )
            catalog.create_asset(asset)

        profiles: List[ColumnProfile] = []
        for col_name, meta in cols.items():
            profiles.append(ColumnProfile(
                column_name=col_name,
                data_type="",
                description=meta.get("description"),
                classification=to_catalog_classification(meta.get("classification", "internal")),
            ))
        catalog.update_asset(asset.id, {"columns": profiles})
        asset_map[table] = asset

    for rel in relationships:
        src = asset_map.get(rel["source_table"]) or catalog.get_asset_by_name(rel["source_table"], AssetType.TABLE)
        tgt = asset_map.get(rel["target_table"]) or catalog.get_asset_by_name(rel["target_table"], AssetType.TABLE)
        if src and tgt:
            src.downstream_assets.append(tgt.id)
            tgt.upstream_assets.append(src.id)
            catalog.update_asset(src.id, {"downstream_assets": src.downstream_assets})
            catalog.update_asset(tgt.id, {"upstream_assets": tgt.upstream_assets})


def _discover_data_model(
    source_type: str,
    csv_paths: Optional[str],
    csv_folder: Optional[str],
    connection_string: Optional[str],
    database: Optional[str],
    schema: Optional[str],
    tables: Optional[str],
    sample_limit: int,
    min_overlap: float,
    persist: bool,
    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,
    return_samples: bool = False,
    sf_connection=None,
) -> Dict[str, Any]:
    _reuse_conn = sf_connection

    sample_tables: Dict[str, Dict[str, List[dict]]] = {}
    if source_type == "csv":
        if csv_folder:
            sample_tables = load_csv_folder(csv_folder, limit=sample_limit)
        elif csv_paths:
            paths = [p.strip() for p in csv_paths.split(",") if p.strip()]
            sample_tables = load_csv_samples(paths, limit=sample_limit)
    elif source_type == "snowflake":
        if not database or not schema or not tables:
            return {"error": "database, schema, and tables are required for snowflake source"}
        table_list = [t.strip() for t in tables.split(",") if t.strip()]
        ensure_snowflake_env_from_cli(connection_name=snowflake_connection)

        if _reuse_conn is None:
            from .snowflake_pool import sf_pool
            _reuse_conn = sf_pool.get_from_env(database=database, schema=schema)

        if table_list == ["*"]:
            _cur = _reuse_conn.cursor()
            _cur.execute(
                f"SELECT TABLE_NAME FROM {database}.INFORMATION_SCHEMA.TABLES "
                f"WHERE TABLE_SCHEMA = '{schema}' AND TABLE_TYPE = 'BASE TABLE' "
                f"ORDER BY TABLE_NAME"
            )
            table_list = [row[0] for row in _cur.fetchall()]
            _cur.close()

        sample_tables = load_snowflake_samples(
            database, schema, table_list, limit=sample_limit,
            sf_connection=_reuse_conn,
        )
    elif source_type == "sql":
        if not connection_string or not tables:
            return {"error": "connection_string and tables are required for sql source"}
        table_list = [t.strip() for t in tables.split(",") if t.strip()]
        sample_tables = load_sqlalchemy_samples(connection_string, database, schema, table_list, limit=sample_limit)
    else:
        return {"error": f"Unknown source_type: {source_type}"}

    if not sample_tables:
        return {"error": "No samples loaded"}

    column_metadata: List[dict] = []
    if return_samples and source_type == "snowflake" and database and schema:
        try:
            if _reuse_conn is None:
                from .snowflake_pool import sf_pool
                _reuse_conn = sf_pool.get_from_env(database=database, schema=schema)
            cur = _reuse_conn.cursor()
            table_in = ",".join(f"'{t}'" for t in sample_tables.keys())
            cur.execute(
                f"SELECT TABLE_NAME, COLUMN_NAME, DATA_TYPE, ORDINAL_POSITION "
                f"FROM {database}.INFORMATION_SCHEMA.COLUMNS "
                f"WHERE TABLE_SCHEMA = '{schema}' "
                f"AND TABLE_NAME IN ({table_in}) "
                f"ORDER BY TABLE_NAME, ORDINAL_POSITION"
            )
            for row in cur.fetchall():
                column_metadata.append({
                    "TABLE_NAME": row[0],
                    "COLUMN_NAME": row[1],
                    "DATA_TYPE": row[2],
                    "ORDINAL_POSITION": row[3],
                })
            cur.close()
        except Exception:
            column_metadata = []

    detected_erp = detect_erp_type(list(sample_tables.keys()))
    if detected_erp:
        entry = get_erp_config(detected_erp)
        if entry:
            current = get_default_config()
            if current.min_overlap >= 0.5:
                set_default_config(entry.inference)

    logger.info("Inferring relationships across %d tables...", len(sample_tables))
    relationships = infer_relationships(sample_tables, min_overlap=min_overlap)
    logger.info("Found %d relationships", len(relationships))

    dictionary: Dict[str, Dict[str, Dict[str, str]]] = {}
    total = len(sample_tables)
    for i, (table, data) in enumerate(sample_tables.items(), 1):
        columns = [{"name": c, "data_type": ""} for c in data.get("columns", [])]
        logger.info("Classifying %d/%d: %s (%d cols)", i, total, table, len(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,
            connection_name=snowflake_connection, sf_connection=_reuse_conn,
        )
        dictionary[table] = generate_dictionary(
            table, columns, samples, classifications,
            database=database, schema=schema,
            connection_name=snowflake_connection, sf_connection=_reuse_conn,
        )

    if persist:
        logger.info("Persisting to catalog and knowledge base...")
        catalog = _ensure_catalog()
        pipeline = KBIngestionPipeline()
        _persist_to_catalog(catalog, dictionary, relationships, database, schema)
        _persist_to_kb(pipeline, dictionary, relationships)

    dbt_results = None
    if dbt_project_name:
        generator = DbtProjectGenerator()
        project = generator.get_project(dbt_project_name)
        if not project:
            project = generator.create_project(name=dbt_project_name, profile=dbt_project_name)
        project_dir = Path("data/dbt_projects") / project.config.name
        generator.scaffold_project(project, output_path=str(project_dir))
        source_schema = schema.lower() if schema else None
        generate_dbt_assets(project_dir, dictionary, relationships, source_name="raw", source_schema=source_schema)

        ensure_profiles_yml(profile_name=project.config.profile)
        if dbt_run:
            run_dbt_command("run", project_dir)
        if dbt_docs:
            run_dbt_command("docs generate", project_dir)
        dbt_results = load_run_results(project_dir)
        if ingest_dbt:
            ingest_dbt_artifacts(project_dir)

    if dbt_manifest_path and Path(dbt_manifest_path).exists():
        try:
            from src.data_catalog.lineage_extractor import DbtLineageExtractor
            extractor = DbtLineageExtractor()
            extraction = extractor.extract_from_manifest(dbt_manifest_path)
            pipeline = KBIngestionPipeline()
            nodes = []
            edges = []
            vectors = []
            for node in extraction.get("nodes", []):
                unique_id = node.get("unique_id", node.get("name"))
                node_id = f"DBT:{unique_id}"
                nodes.append(GraphNode(id=node_id, type="dbt_model", properties=node))
                vectors.append(VectorMetadata(
                    id=node_id, source_type="dbt_model", source_id=unique_id,
                    tags=["dbt"],
                    extra={"name": node.get("name"), "materialized": node.get("materialized")},
                ))
            for edge in extraction.get("edges", []):
                edges.append(GraphEdge(
                    source=f"DBT:{edge.get('source')}",
                    target=f"DBT:{edge.get('target')}",
                    type="DEPENDS_ON",
                ))
            if nodes:
                pipeline.ingest_graph(nodes, edges)
                pipeline.ingest_vector_metadata(vectors)
        except Exception:
            pass

    run_path = _save_run({
        "source_type": source_type,
        "tables": list(sample_tables.keys()),
        "relationship_count": len(relationships),
        "dictionary_tables": len(dictionary),
    })

    result = {
        "status": "completed",
        "tables": list(sample_tables.keys()),
        "relationship_count": len(relationships),
        "relationships": relationships,
        "dictionary_tables": list(dictionary.keys()),
        "dbt_results": dbt_results,
        "run_path": str(run_path),
    }
    if return_samples:
        result["sample_tables"] = sample_tables
        result["column_metadata"] = column_metadata
    return result


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

def register_data_modeling_tools(mcp, mcp_settings):
    """Register data modeling MCP tools (unified + standalone)."""

    # Import and register the 2 unified tools
    from .unified import (
        dispatch_model_discover,
        dispatch_model_relationships,
        register_unified_modeling_tools,
    )
    register_unified_modeling_tools(mcp, mcp_settings)

    # =========================================================================
    # Standalone tool (unchanged -- heavy cross-account ETL)
    # =========================================================================

    @mcp.tool()
    def model_load_dimensional_model(
        spec_path: Optional[str] = None,
        spec_json: Optional[str] = None,
        source_db: str = "",
        source_schema: str = "",
        dest_db: str = "",
        dest_schema: str = "",
        dim_limit: int = 200000,
        txn_limit: int = 1000000,
        batch_size: int = 10000,
        truncate: bool = True,
        parallel: bool = False,
        max_workers: int = 4,
    ) -> Dict[str, Any]:
        """Load a dimensional model from source to destination using table specs.

        Provide either spec_path (JSON file) or spec_json (inline JSON).
        Supports parallel dimension loading for faster execution.

        spec_path: Path to JSON spec file defining table mappings.
        spec_json: Inline JSON string of table specs.
        source_db: Source Snowflake database.
        source_schema: Source schema.
        dest_db: Destination database.
        dest_schema: Destination schema.
        parallel: Load dimensions concurrently (default False).
        """
        from .dm_loader import DMSpecLoader, GenericDMLoader

        if spec_path:
            specs = DMSpecLoader.from_json(spec_path)
        elif spec_json:
            try:
                data = json.loads(spec_json)
                specs = [__import__("src.data_modeling.dm_loader", fromlist=["TableSpec"]).TableSpec.from_dict(d) for d in data]
            except (json.JSONDecodeError, KeyError) as exc:
                return {"error": f"Invalid spec JSON: {exc}"}
        else:
            return {"error": "Provide either spec_path or spec_json"}

        if not source_db or not source_schema or not dest_db or not dest_schema:
            return {"error": "source_db, source_schema, dest_db, and dest_schema are required"}

        loader = GenericDMLoader(
            source_db=source_db,
            source_schema=source_schema,
            dest_db=dest_db,
            dest_schema=dest_schema,
        )
        return loader.load(
            specs=specs,
            dim_limit=dim_limit,
            txn_limit=txn_limit,
            batch_size=batch_size,
            truncate=truncate,
            parallel=parallel,
            max_workers=max_workers,
        )

    logger.info("Registered 3 Data Modeling MCP tools (2 unified + 1 standalone)")
    return {
        "tools_registered": 3,
        "unified_tools": ["model_discover", "model_relationships"],
        "standalone_tools": ["model_load_dimensional_model"],
        "tools": [
            "model_discover", "model_relationships", "model_load_dimensional_model",
        ],
    }
