"""MCP tool registration for the Mart Factory plugin.

Registers 10 tools: 3 config, 3 pipeline, 2 discovery, 2 validation.
"""

from __future__ import annotations

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

logger = logging.getLogger(__name__)


def register_mart_factory_tools(mcp, settings) -> Dict[str, Any]:
    """Register all 10 Mart Factory MCP tools with the server.

    Args:
        mcp: FastMCP server instance.
        settings: Application settings (Pydantic Settings).

    Returns:
        Dict with registration summary.
    """
    try:
        from databridge.mart import (
            MartConfigGenerator,
            CortexDiscoveryAgent,
            FormulaPrecedenceEngine,
            MartPipelineGenerator,
        )
    except ImportError:
        from .config_generator import MartConfigGenerator
        from .cortex_discovery import CortexDiscoveryAgent
        from .formula_engine import FormulaPrecedenceEngine
        from .pipeline_generator import MartPipelineGenerator

    # Shared instances
    _config_gen = MartConfigGenerator()
    _formula_engine = FormulaPrecedenceEngine()
    _pipeline_gen = MartPipelineGenerator(formula_engine=_formula_engine)
    _discovery = CortexDiscoveryAgent()

    # ==================================================================
    # CONFIG TOOLS (1-3)
    # ==================================================================

    @mcp.tool()
    def create_mart_config(
        project_name: str,
        report_type: str = "P&L",
        hierarchy_table: str = "",
        mapping_table: str = "",
        account_segment: str = "GROSS",
        measure_prefix: str = "AMT_",
        schema_name: str = "ANALYTICS",
        description: str = "",
    ) -> str:
        """Create a new data mart configuration with the 7-variable contract.

        Define the foundation for a 4-object DDL pipeline (VW_1→DT_2→DT_3A→DT_3).
        Supports P&L, Balance Sheet, LOS, and Custom report types.

        Args:
            project_name: Unique identifier for this mart project.
            report_type: One of 'P&L', 'Balance Sheet', 'LOS', 'Custom'.
            hierarchy_table: Source hierarchy table/view name.
            mapping_table: ID_SOURCE mapping table/view name.
            account_segment: GROSS or NET filter value.
            measure_prefix: Column prefix for measures (e.g. AMT_).
            schema_name: Target Snowflake schema name.
            description: Human-readable description.
        """
        config = _config_gen.create_config(
            project_name=project_name,
            report_type=report_type,
            hierarchy_table=hierarchy_table,
            mapping_table=mapping_table,
            account_segment=account_segment,
            measure_prefix=measure_prefix,
            schema_name=schema_name,
            description=description,
        )
        return json.dumps({
            "status": "created",
            "project_name": config.project_name,
            "report_type": config.report_type.value,
            "hierarchy_table": config.hierarchy_table,
            "mapping_table": config.mapping_table,
            "configs_count": len(_config_gen.list_configs()),
        }, indent=2)

    @mcp.tool()
    def add_mart_join_pattern(
        project_name: str,
        name: str,
        join_keys: str,
        fact_keys: str,
        source_table: str = "",
        filter_expr: str = "",
    ) -> str:
        """Add a UNION ALL branch to a mart configuration.

        Each branch defines a data source that feeds into DT_3A
        pre-aggregation via UNION ALL.

        Args:
            project_name: Target mart config.
            name: Branch identifier (e.g. 'revenue', 'headcount').
            join_keys: Comma-separated join columns (e.g. 'HIERARCHY_KEY,PERIOD').
            fact_keys: Comma-separated fact/measure columns (e.g. 'AMT_USD,AMT_LOCAL').
            source_table: Optional source table override.
            filter_expr: Optional WHERE clause fragment.
        """
        try:
            config = _config_gen.add_join_pattern(
                project_name=project_name,
                name=name,
                join_keys=[k.strip() for k in join_keys.split(",") if k.strip()],
                fact_keys=[k.strip() for k in fact_keys.split(",") if k.strip()],
                source_table=source_table,
                filter_expr=filter_expr or None,
            )
            return json.dumps({
                "status": "added",
                "project_name": project_name,
                "pattern_name": name,
                "total_patterns": len(config.join_patterns),
            }, indent=2)
        except KeyError as e:
            return json.dumps({"error": str(e)}, indent=2)

    @mcp.tool()
    def export_mart_config(
        project_name: str,
        format: str = "yaml",
    ) -> str:
        """Export a mart configuration as YAML or JSON.

        Args:
            project_name: Config to export.
            format: Output format — 'yaml' for dbt YAML, 'json' for raw JSON.
        """
        try:
            config = _config_gen.get_config(project_name)
            if config is None:
                return json.dumps({"error": f"Config not found: {project_name}"})

            if format.lower() == "yaml":
                return _config_gen.export_yaml(project_name)
            else:
                return json.dumps(config.model_dump(), indent=2, default=str)
        except KeyError as e:
            return json.dumps({"error": str(e)}, indent=2)

    # ==================================================================
    # PIPELINE TOOLS (4-6)
    # ==================================================================

    @mcp.tool()
    def generate_mart_pipeline(
        project_name: str,
    ) -> str:
        """Generate the full 4-object DDL pipeline for a mart config.

        Creates VW_1 (Translation View), DT_2 (Granularity Table),
        DT_3A (Pre-Aggregation Fact), and DT_3 (Data Mart) in order.

        Args:
            project_name: Mart config to generate pipeline for.
        """
        try:
            config = _config_gen.get_config(project_name)
            if config is None:
                return json.dumps({"error": f"Config not found: {project_name}"})

            pipeline = _pipeline_gen.generate_full_pipeline(config)
            return json.dumps({
                "project_name": pipeline.project_name,
                "objects_count": pipeline.total_objects,
                "objects": [
                    {
                        "layer": o.layer.value,
                        "object_name": o.object_name,
                        "object_type": o.object_type.value,
                        "ddl_length": len(o.ddl),
                        "depends_on": o.depends_on,
                    }
                    for o in pipeline.objects
                ],
                "notes": pipeline.generation_notes,
            }, indent=2, default=str)
        except KeyError as e:
            return json.dumps({"error": str(e)}, indent=2)

    @mcp.tool()
    def generate_mart_object(
        project_name: str,
        layer: str,
    ) -> str:
        """Generate DDL for a specific pipeline layer.

        Args:
            project_name: Mart config identifier.
            layer: Pipeline layer — 'VW_1', 'DT_2', 'DT_3A', or 'DT_3'.
        """
        try:
            config = _config_gen.get_config(project_name)
            if config is None:
                return json.dumps({"error": f"Config not found: {project_name}"})

            layer_upper = layer.upper()
            generators = {
                "VW_1": _pipeline_gen.generate_vw1,
                "DT_2": _pipeline_gen.generate_dt2,
                "DT_3A": _pipeline_gen.generate_dt3a,
                "DT_3": lambda c: _pipeline_gen.generate_dt3(c),
            }

            gen_fn = generators.get(layer_upper)
            if gen_fn is None:
                return json.dumps({
                    "error": f"Unknown layer: {layer}. Use VW_1, DT_2, DT_3A, or DT_3."
                })

            obj = gen_fn(config)
            return json.dumps({
                "layer": obj.layer.value,
                "object_name": obj.object_name,
                "object_type": obj.object_type.value,
                "ddl": obj.ddl,
                "depends_on": obj.depends_on,
            }, indent=2, default=str)
        except KeyError as e:
            return json.dumps({"error": str(e)}, indent=2)

    @mcp.tool()
    def generate_mart_dbt_project(
        project_name: str,
    ) -> str:
        """Generate pipeline as dbt model SQL files.

        Outputs one .sql file per pipeline layer with dbt config blocks.

        Args:
            project_name: Mart config identifier.
        """
        try:
            config = _config_gen.get_config(project_name)
            if config is None:
                return json.dumps({"error": f"Config not found: {project_name}"})

            models = _pipeline_gen.generate_dbt_models(config)
            return json.dumps({
                "project_name": project_name,
                "model_count": len(models),
                "models": {
                    filename: {"lines": sql.count("\n") + 1, "preview": sql[:200]}
                    for filename, sql in models.items()
                },
            }, indent=2, default=str)
        except KeyError as e:
            return json.dumps({"error": str(e)}, indent=2)

    # ==================================================================
    # DISCOVERY TOOLS (7-8)
    # ==================================================================

    @mcp.tool()
    def discover_hierarchy_pattern(
        table_name: str,
        columns: str,
        sample_data: str = "{}",
    ) -> str:
        """Discover hierarchy patterns in table metadata using heuristic AI.

        Scans column names and sample values to detect:
        - Report type (P&L, Balance Sheet, LOS)
        - Hierarchy levels
        - ID_SOURCE distribution
        - Data quality issues

        Args:
            table_name: Name of the source table.
            columns: Comma-separated column names.
            sample_data: JSON string of column→values mapping for deeper analysis.
        """
        col_list = [c.strip() for c in columns.split(",") if c.strip()]
        try:
            sample_values = json.loads(sample_data) if sample_data else {}
        except json.JSONDecodeError:
            sample_values = {}

        table_data = {
            "columns": col_list,
            "sample_values": sample_values,
        }

        result = _discovery.discover_hierarchy(table_data, table_name)
        return json.dumps({
            "table_name": table_name,
            "hierarchy_type": result.hierarchy_type.value,
            "level_count": result.level_count,
            "detected_levels": result.detected_levels,
            "id_source_count": len(result.id_source_values),
            "id_source_values": result.id_source_values[:20],
            "join_patterns": len(result.recommended_join_patterns),
            "confidence": result.confidence,
            "quality_issues": [qi.model_dump() for qi in result.quality_issues],
            "suggestions": result.suggestions,
        }, indent=2, default=str)

    @mcp.tool()
    def suggest_mart_config(
        table_name: str,
        columns: str,
        sample_data: str = "{}",
    ) -> str:
        """Generate a recommended MartConfig from table discovery.

        Combines hierarchy pattern discovery with config generation
        to produce a ready-to-use mart configuration.

        Args:
            table_name: Source table name.
            columns: Comma-separated column names.
            sample_data: JSON string of column→values mapping.
        """
        col_list = [c.strip() for c in columns.split(",") if c.strip()]
        try:
            sample_values = json.loads(sample_data) if sample_data else {}
        except json.JSONDecodeError:
            sample_values = {}

        table_data = {
            "columns": col_list,
            "sample_values": sample_values,
        }

        recommendation = _discovery.generate_config_recommendation(table_name, table_data)
        return json.dumps(recommendation, indent=2, default=str)

    # ==================================================================
    # VALIDATION TOOLS (9-10)
    # ==================================================================

    @mcp.tool()
    def validate_mart_config(
        project_name: str,
    ) -> str:
        """Validate a mart configuration for completeness and consistency.

        Checks required fields, join pattern integrity, column mapping
        uniqueness, and flag consistency.

        Args:
            project_name: Config to validate.
        """
        try:
            result = _config_gen.validate_config(project_name)
            return json.dumps({
                "valid": result.valid,
                "checked": result.checked,
                "passed": result.passed,
                "failed": result.failed,
                "messages": [m.model_dump() for m in result.messages],
            }, indent=2, default=str)
        except KeyError as e:
            return json.dumps({"error": str(e)}, indent=2)

    @mcp.tool()
    def validate_mart_pipeline(
        project_name: str,
    ) -> str:
        """Validate a generated mart pipeline for correctness.

        Checks all 4 layers are present, DDL is non-empty,
        and dependency chain is valid.

        Args:
            project_name: Config whose pipeline to validate.
        """
        try:
            config = _config_gen.get_config(project_name)
            if config is None:
                return json.dumps({"error": f"Config not found: {project_name}"})

            pipeline = _pipeline_gen.generate_full_pipeline(config)
            result = _pipeline_gen.validate_pipeline(pipeline)
            return json.dumps({
                "valid": result.valid,
                "checked": result.checked,
                "passed": result.passed,
                "failed": result.failed,
                "messages": [m.model_dump() for m in result.messages],
            }, indent=2, default=str)
        except KeyError as e:
            return json.dumps({"error": str(e)}, indent=2)

    # ------------------------------------------------------------------
    logger.info("[Plugin] mart_factory registered 10 tools")
    return {"tools_registered": 10, "plugin": "mart_factory"}
