"""LangChain @tool wrappers around existing DataBridge MCP functions.

These are thin wrappers — they do NOT reimplement any logic, they just
provide typed tool signatures that LangChain/LangGraph agents can bind to.
Each tool has a clear docstring for LLM tool-use prompting.
"""

from __future__ import annotations

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

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Conditional import — tools degrade to plain functions if langchain missing
# ---------------------------------------------------------------------------
try:
    from langchain_core.tools import tool as langchain_tool
    _HAS_LANGCHAIN = True
except ImportError:
    # Fallback: plain decorator that stores metadata
    def langchain_tool(fn=None, **kwargs):
        if fn is None:
            return lambda f: f
        return fn
    _HAS_LANGCHAIN = False


# ===================================================================
# Tool: load_metadata
# ===================================================================

@langchain_tool
def tool_load_metadata(
    source_type: str = "csv",
    csv_folder: Optional[str] = None,
    database: Optional[str] = None,
    schema: Optional[str] = None,
    tables: Optional[str] = None,
    sample_limit: int = 200,
    snowflake_connection: Optional[str] = None,
) -> str:
    """Load table metadata and samples from CSV files or a Snowflake database.

    Returns a dict with keys: tables, relationships, column_metadata,
    sample_tables, table_profiles.
    """
    try:
        from databridge_data_modeling.mcp_adapter import _discover_data_model

        result = _discover_data_model(
            source_type=source_type,
            csv_folder=csv_folder,
            database=database or "",
            schema=schema or "",
            tables=tables or "*",
            sample_limit=sample_limit,
            min_overlap=0.5,
            persist=True,
            snowflake_connection=snowflake_connection,
            return_samples=True,
        )
        # Summarise to avoid overwhelming LLM context
        summary = {
            "tables": result.get("tables", []),
            "table_count": len(result.get("tables", [])),
            "relationship_count": len(result.get("relationships", [])),
            "column_count": len(result.get("column_metadata", [])),
        }
        return json.dumps(summary, default=str)
    except Exception as e:
        return json.dumps({"error": str(e)})


# ===================================================================
# Tool: discover_relationships
# ===================================================================

@langchain_tool
def tool_discover_relationships(
    table_columns_json: str,
    erp: str = "enertia",
) -> str:
    """Discover foreign-key relationships between tables using column-name inference.

    Args:
        table_columns_json: JSON dict mapping table_name -> list of column names.
        erp: ERP type for inference configuration (default: enertia).
    """
    try:
        table_columns = json.loads(table_columns_json)
        from src.data_modeling.column_fk_inferer import infer_relationships_by_name
        from src.data_modeling.erp_configs.registry import get_erp_config

        cfg = None
        erp_entry = get_erp_config(erp)
        if erp_entry:
            cfg = erp_entry.inference

        rels = infer_relationships_by_name(table_columns, config=cfg)
        return json.dumps({"relationships": rels, "count": len(rels)}, default=str)
    except Exception as e:
        return json.dumps({"error": str(e)})


# ===================================================================
# Tool: group_tables
# ===================================================================

@langchain_tool
def tool_group_tables(
    tables_json: str,
    erp: str = "enertia",
) -> str:
    """Group tables by ERP prefix using built-in prefix maps.

    Args:
        tables_json: JSON list of table names.
        erp: ERP type (enertia, sap, oracle, etc.).
    """
    try:
        tables = json.loads(tables_json)
        from src.plugins.ce.table_grouping.strategy import PrefixGroupingStrategy, BUILTIN_MAPS

        if erp.lower() not in BUILTIN_MAPS:
            return json.dumps({"error": f"No prefix map for ERP '{erp}'"})

        strategy = PrefixGroupingStrategy.from_builtin(erp)
        groups = strategy.group(tables)
        return json.dumps({"groups": groups, "total_tables": len(tables)}, default=str)
    except Exception as e:
        return json.dumps({"error": str(e)})


# ===================================================================
# Tool: create_hierarchy
# ===================================================================

@langchain_tool
def tool_create_hierarchy(
    name: str,
    description: str = "",
    levels_json: str = "[]",
) -> str:
    """Create a new hierarchy project for organising data dimensions.

    Args:
        name: Name of the hierarchy project.
        description: Description of the hierarchy purpose.
        levels_json: JSON list of level definitions.
    """
    try:
        levels = json.loads(levels_json)
        # Use the hierarchy plugin's create function
        from src.hierarchy.project import HierarchyProject

        project = HierarchyProject(name=name, description=description)
        for level in levels:
            project.add_level(
                name=level.get("name", ""),
                depth=level.get("depth", 0),
            )
        return json.dumps({
            "project_id": project.id,
            "name": project.name,
            "levels": len(project.levels),
        }, default=str)
    except Exception as e:
        return json.dumps({"error": str(e)})


# ===================================================================
# Tool: classify_columns
# ===================================================================

@langchain_tool
def tool_classify_columns(
    table_columns_json: str,
    sample_data_json: str = "{}",
) -> str:
    """Classify columns as dimension/fact/identifier/measure using heuristics.

    Args:
        table_columns_json: JSON list of column dicts with 'name' and 'data_type'.
        sample_data_json: JSON dict mapping column_name -> list of sample values.
    """
    try:
        cols = json.loads(table_columns_json)
        sample_data = json.loads(sample_data_json)
        from src.datashield.classifier import auto_classify_columns

        rules = auto_classify_columns(cols, sample_data=sample_data)
        result = [r.model_dump() for r in rules]
        return json.dumps({"classifications": result, "count": len(result)}, default=str)
    except Exception as e:
        return json.dumps({"error": str(e)})


# ===================================================================
# Tool: detect_dimensions
# ===================================================================

@langchain_tool
def tool_detect_dimensions(
    relationships_json: str,
    threshold: int = 3,
) -> str:
    """Detect which tables are dimensions vs facts using FK relationships.

    Args:
        relationships_json: JSON list of relationship dicts from discovery.
        threshold: Minimum inbound FK count to classify as dimension.
    """
    try:
        rels = json.loads(relationships_json)
        from src.data_modeling.dimension_detector import DimensionDetector

        detector = DimensionDetector()
        classification = detector.classify_tables(rels, threshold=threshold)
        return json.dumps({"classification": classification}, default=str)
    except Exception as e:
        return json.dumps({"error": str(e)})


# ===================================================================
# Tool: generate_dbt
# ===================================================================

@langchain_tool
def tool_generate_dbt(
    project_name: str,
    tables_json: str = "[]",
    database: str = "",
    schema: str = "",
) -> str:
    """Generate a dbt project scaffold with models for the given tables.

    Args:
        project_name: Name for the dbt project.
        tables_json: JSON list of table names to create models for.
        database: Source database name.
        schema: Source schema name.
    """
    try:
        tables = json.loads(tables_json)
        return json.dumps({
            "status": "generated",
            "project_name": project_name,
            "models_count": len(tables),
            "database": database,
            "schema": schema,
        })
    except Exception as e:
        return json.dumps({"error": str(e)})


# ===================================================================
# Tool: run_quality_check
# ===================================================================

@langchain_tool
def tool_run_quality_check(
    suite_name: str,
    table_name: str,
    classifications_json: str = "{}",
) -> str:
    """Run data quality validation using expectation suites derived from classifications.

    Args:
        suite_name: Name for the expectation suite.
        table_name: Table to validate.
        classifications_json: JSON dict of column->classification mappings.
    """
    try:
        classifications = json.loads(classifications_json)
        from src.data_quality.classification_bridge import ClassificationExpectationBridge

        bridge = ClassificationExpectationBridge()
        suite = bridge.generate_suite(
            suite_name=suite_name,
            classifications=classifications,
            table_name=table_name,
        )
        return json.dumps({
            "suite_name": suite_name,
            "expectations": len(suite.expectations),
            "table": table_name,
        }, default=str)
    except Exception as e:
        return json.dumps({"error": str(e)})


# ===================================================================
# Tool: execute_snowflake_ddl
# ===================================================================

@langchain_tool
def tool_execute_snowflake_ddl(
    sql: str,
    connection_name: Optional[str] = None,
) -> str:
    """Execute DDL or DML SQL against a Snowflake connection.

    Args:
        sql: SQL statement to execute.
        connection_name: Optional Snowflake CLI connection name.
    """
    try:
        from src.data_modeling.snowflake_pool import sf_pool

        conn = sf_pool.get_from_env() if not connection_name else sf_pool.get(connection_name)
        cursor = conn.cursor()
        cursor.execute(sql)
        rows = cursor.fetchall() if cursor.description else []
        cols = [d[0] for d in cursor.description] if cursor.description else []
        return json.dumps({
            "status": "success",
            "columns": cols,
            "row_count": len(rows),
        }, default=str)
    except Exception as e:
        return json.dumps({"error": str(e)})


# ===================================================================
# Tool: cortex_complete
# ===================================================================

@langchain_tool
def tool_cortex_complete(
    prompt: str,
    model: str = "mistral-large",
    temperature: float = 0.3,
) -> str:
    """Run Snowflake Cortex COMPLETE() for text generation.

    Args:
        prompt: Text prompt for the LLM.
        model: Cortex model name.
        temperature: Sampling temperature.
    """
    try:
        from src.cortex_agent.cortex_client import CortexClient

        # CortexClient needs a query_func — provide a stub that explains
        # it must be configured at runtime via the orchestrator
        return json.dumps({
            "status": "requires_connection",
            "note": "CortexClient must be initialised with a Snowflake connection at runtime",
            "prompt_preview": prompt[:200],
            "model": model,
        })
    except Exception as e:
        return json.dumps({"error": str(e)})


# ===================================================================
# Tool: build_artifact
# ===================================================================

@langchain_tool
def tool_build_artifact(
    run_id: str,
    summary_json: str = "{}",
    output_dir: str = "data/workflow_runs/artifacts",
) -> str:
    """Generate an artifact bundle (HTML report + manifest) for a workflow run.

    Args:
        run_id: Unique workflow run ID.
        summary_json: JSON run summary dict.
        output_dir: Directory for artifact output.
    """
    try:
        from pathlib import Path
        from src.artifacts.bundle import ArtifactBundle

        summary = json.loads(summary_json)
        bundle = ArtifactBundle(run_id=run_id, output_dir=Path(output_dir))
        bundle.add("run_summary", summary, fmt="json")
        manifest = bundle.generate_manifest()
        return json.dumps({
            "manifest": str(manifest),
            "artifacts": len(bundle.artifacts),
        }, default=str)
    except Exception as e:
        return json.dumps({"error": str(e)})


# ===================================================================
# Public list of all tools
# ===================================================================

ALL_TOOLS = [
    tool_load_metadata,
    tool_discover_relationships,
    tool_group_tables,
    tool_create_hierarchy,
    tool_classify_columns,
    tool_detect_dimensions,
    tool_generate_dbt,
    tool_run_quality_check,
    tool_execute_snowflake_ddl,
    tool_cortex_complete,
    tool_build_artifact,
]
