"""
Unified MCP tools for dbt Integration (Phase 3B consolidation).

Consolidates 8 individual dbt tools into 2 action-dispatched tools:
- dbt_project(action, ...) — 3 actions for project management
- dbt_generate(type, ...) — 5 types for code generation

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

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

from .types import (
    CiCdConfig,
    CiCdPlatform,
    DbtColumn,
    DbtModelConfig,
    DbtModelType,
    DbtProject,
    DbtProjectConfig,
    DbtSource,
    DbtSourceTable,
    ValidationResult,
)

logger = logging.getLogger(__name__)

# ============================================================================
# Lazy singletons
# ============================================================================

_project_gen = None
_model_gen = None
_source_gen = None
_metrics_gen = None
_cicd_gen = None


def _ensure_project_gen():
    global _project_gen
    if _project_gen is None:
        from .project_generator import DbtProjectGenerator
        _project_gen = DbtProjectGenerator()
    return _project_gen


def _ensure_model_gen():
    global _model_gen
    if _model_gen is None:
        from .model_generator import DbtModelGenerator
        _model_gen = DbtModelGenerator()
    return _model_gen


def _ensure_source_gen():
    global _source_gen
    if _source_gen is None:
        from .source_generator import DbtSourceGenerator
        _source_gen = DbtSourceGenerator()
    return _source_gen


def _ensure_metrics_gen():
    global _metrics_gen
    if _metrics_gen is None:
        from .source_generator import DbtMetricsGenerator
        _metrics_gen = DbtMetricsGenerator()
    return _metrics_gen


def _ensure_cicd_gen():
    global _cicd_gen
    if _cicd_gen is None:
        from .cicd_generator import CiCdGenerator
        _cicd_gen = CiCdGenerator()
    return _cicd_gen


# ============================================================================
# dbt_project action handlers
# ============================================================================

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

    project_gen = _ensure_project_gen()
    project = project_gen.create_project(
        name=name,
        profile=profile,
        target_database=kwargs.get("target_database"),
        target_schema=kwargs.get("target_schema"),
        hierarchy_project_id=kwargs.get("hierarchy_project_id"),
        include_cicd=kwargs.get("include_cicd", False),
    )
    output_dir = kwargs.get("output_dir")
    files = project_gen.scaffold_project(project, output_dir)

    return {
        "success": True,
        "project_id": project.id,
        "project_name": project.config.name,
        "profile": project.config.profile,
        "files_generated": list(files.keys()),
        "file_count": len(files),
        "output_dir": output_dir,
        "message": f"Created dbt project '{project.config.name}' with {len(files)} files",
    }


def _project_validate(settings, **kwargs) -> Dict[str, Any]:
    project_name = kwargs.get("project_name")
    if not project_name:
        return {"error": "project_name is required for 'validate' action"}

    project_gen = _ensure_project_gen()
    project = project_gen.get_project(project_name)
    if not project:
        return {"success": False, "error": f"Project '{project_name}' not found"}

    errors = []
    warnings = []
    missing_files = []
    files = project.generated_files

    required_files = ["dbt_project.yml", "profiles.yml.template"]
    for req_file in required_files:
        if req_file not in files:
            missing_files.append(req_file)
            errors.append(f"Missing required file: {req_file}")

    model_files = [f for f in files.keys() if f.endswith(".sql") and "models/" in f]
    if not model_files:
        warnings.append("No SQL models found in project")

    if "models/sources.yml" not in files:
        warnings.append("No sources.yml found - staging models may not work")

    for model in project.models:
        for ref in model.refs:
            ref_model = next((m for m in project.models if m.name == ref), None)
            if not ref_model:
                warnings.append(f"Model '{model.name}' references unknown model: {ref}")

    result = ValidationResult(
        valid=len(errors) == 0,
        errors=errors,
        warnings=warnings,
        missing_files=missing_files,
    )

    return {
        "success": True,
        "valid": result.valid,
        "errors": result.errors,
        "warnings": result.warnings,
        "missing_files": result.missing_files,
        "files_count": len(files),
        "models_count": len(project.models),
        "sources_count": len(project.sources),
    }


def _project_export(settings, **kwargs) -> Dict[str, Any]:
    project_name = kwargs.get("project_name")
    if not project_name:
        return {"error": "project_name is required for 'export' action"}

    project_gen = _ensure_project_gen()
    project = project_gen.get_project(project_name)
    if not project:
        return {"success": False, "error": f"Project '{project_name}' not found"}

    output_dir = kwargs.get("output_dir")
    if not output_dir:
        output_dir = f"./dbt_export/{project.config.name}"

    output_path = Path(output_dir)
    project_gen._write_files(str(output_path), project.generated_files)

    result = {
        "success": True,
        "project_name": project.config.name,
        "output_dir": str(output_path.absolute()),
        "files_exported": len(project.generated_files),
        "file_list": list(project.generated_files.keys()),
    }

    as_zip = kwargs.get("as_zip", False)
    if as_zip:
        zip_path = output_path.with_suffix(".zip")
        with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
            for filepath, content in project.generated_files.items():
                zf.writestr(filepath, content)
        result["zip_file"] = str(zip_path.absolute())
        result["message"] = f"Exported project to {output_path} and {zip_path}"
    else:
        result["message"] = f"Exported project to {output_path}"

    return result


# ============================================================================
# dbt_generate type handlers
# ============================================================================

def _generate_model(settings, **kwargs) -> Dict[str, Any]:
    project_name = kwargs.get("project_name")
    model_name = kwargs.get("model_name")
    if not project_name or not model_name:
        return {"error": "project_name and model_name are required for 'model' type"}

    project_gen = _ensure_project_gen()
    model_gen = _ensure_model_gen()

    project = project_gen.get_project(project_name)
    if not project:
        return {"success": False, "error": f"Project '{project_name}' not found"}

    model_type = kwargs.get("model_type", "staging")
    source_name = kwargs.get("source_name")
    source_table = kwargs.get("source_table")
    ref_models = kwargs.get("ref_models")
    columns = kwargs.get("columns")
    case_mappings = kwargs.get("case_mappings")
    description = kwargs.get("description")

    cols = json.loads(columns) if columns else None
    mappings = json.loads(case_mappings) if case_mappings else None
    refs = ref_models.split(",") if ref_models else []

    type_map = {
        "staging": DbtModelType.STAGING,
        "intermediate": DbtModelType.INTERMEDIATE,
        "dimension": DbtModelType.DIM,
        "dim": DbtModelType.DIM,
        "fact": DbtModelType.FACT,
        "mart": DbtModelType.MART,
    }
    model_type_enum = type_map.get(model_type.lower(), DbtModelType.STAGING)

    if model_type_enum == DbtModelType.STAGING:
        if not source_name or not source_table:
            return {"success": False, "error": "source_name and source_table required for staging models"}
        sql = model_gen.generate_staging_model(
            model_name=model_name, source_name=source_name, source_table=source_table,
            columns=cols, case_mappings=mappings, description=description,
        )
    elif model_type_enum == DbtModelType.INTERMEDIATE:
        if not refs:
            return {"success": False, "error": "ref_models required for intermediate models"}
        sql = model_gen.generate_intermediate_model(
            model_name=model_name, refs=refs, select_columns=cols, description=description,
        )
    elif model_type_enum == DbtModelType.DIM:
        if not refs:
            return {"success": False, "error": "ref_models required for dimension models"}
        hierarchy_cols = cols if cols else ["hierarchy_id", "hierarchy_name", "parent_id"]
        sql = model_gen.generate_dimension_model(
            model_name=model_name, ref_model=refs[0], hierarchy_columns=hierarchy_cols,
            description=description,
        )
    elif model_type_enum == DbtModelType.FACT:
        if not refs:
            return {"success": False, "error": "ref_models required for fact models"}
        sql = model_gen.generate_fact_model(
            model_name=model_name, ref_model=refs[0], dimension_refs=[],
            measure_columns=cols or ["amount"], description=description,
        )
    else:
        sql = model_gen.generate_staging_model(
            model_name=model_name, source_name=source_name or "raw",
            source_table=source_table or model_name, description=description,
        )

    model_config = DbtModelConfig(
        name=model_name, description=description, model_type=model_type_enum,
        source_name=source_name, source_table=source_table, refs=refs,
    )
    model_gen.add_model_to_project(project, model_config, sql)
    project_gen._save()

    prefix_map = {
        DbtModelType.STAGING: ("staging", "stg_"),
        DbtModelType.INTERMEDIATE: ("intermediate", "int_"),
        DbtModelType.DIM: ("marts", "dim_"),
        DbtModelType.FACT: ("marts", "fct_"),
        DbtModelType.MART: ("marts", ""),
    }
    folder, prefix = prefix_map.get(model_type_enum, ("marts", ""))
    file_path = f"models/{folder}/{prefix}{model_name}.sql"

    return {
        "success": True,
        "model_name": f"{prefix}{model_name}",
        "model_type": model_type,
        "file_path": file_path,
        "sql_preview": sql[:500] + "..." if len(sql) > 500 else sql,
        "message": f"Generated {model_type} model: {prefix}{model_name}",
    }


def _generate_sources(settings, **kwargs) -> Dict[str, Any]:
    project_name = kwargs.get("project_name")
    source_name = kwargs.get("source_name")
    if not project_name or not source_name:
        return {"error": "project_name and source_name are required for 'sources' type"}

    project_gen = _ensure_project_gen()
    source_gen = _ensure_source_gen()

    project = project_gen.get_project(project_name)
    if not project:
        return {"success": False, "error": f"Project '{project_name}' not found"}

    mappings = kwargs.get("mappings")
    tables = kwargs.get("tables")
    database = kwargs.get("database")
    schema_name = kwargs.get("schema_name")

    if mappings:
        mapping_data = json.loads(mappings)
        source = source_gen.generate_from_hierarchy_mappings(mapping_data, source_name)
    elif tables:
        table_data = json.loads(tables)
        source_tables = []
        for tbl in table_data:
            columns = [DbtColumn(name=col) for col in tbl.get("columns", [])]
            source_tables.append(DbtSourceTable(
                name=tbl.get("name"), description=tbl.get("description"), columns=columns,
            ))
        source = DbtSource(
            name=source_name, database=database, schema_name=schema_name, tables=source_tables,
        )
    else:
        return {"success": False, "error": "Either mappings or tables required"}

    if database:
        source.database = database
    if schema_name:
        source.schema_name = schema_name

    source_gen.add_source_to_project(project, source)
    project_gen._save()

    return {
        "success": True,
        "source_name": source_name,
        "tables_count": len(source.tables),
        "tables": [t.name for t in source.tables],
        "file_path": "models/sources.yml",
        "message": f"Generated sources.yml with {len(source.tables)} tables",
    }


def _generate_schema(settings, **kwargs) -> Dict[str, Any]:
    project_name = kwargs.get("project_name")
    if not project_name:
        return {"error": "project_name is required for 'schema' type"}

    project_gen = _ensure_project_gen()
    model_gen = _ensure_model_gen()
    source_gen = _ensure_source_gen()

    project = project_gen.get_project(project_name)
    if not project:
        return {"success": False, "error": f"Project '{project_name}' not found"}

    models_filter = kwargs.get("models")
    if models_filter:
        model_names = json.loads(models_filter)
        project_models = [m for m in project.models if m.name in model_names]
    else:
        project_models = project.models

    schema_models = []
    for model in project_models:
        schema_def = model_gen.generate_model_schema(model)
        schema_models.append(schema_def)

    schema_yml = source_gen.generate_schema_yml(schema_models)
    project.generated_files["models/schema.yml"] = schema_yml
    project_gen._save()

    return {
        "success": True,
        "models_count": len(schema_models),
        "models": [m["name"] for m in schema_models],
        "file_path": "models/schema.yml",
        "message": f"Generated schema.yml with {len(schema_models)} models",
    }


def _generate_metrics(settings, **kwargs) -> Dict[str, Any]:
    project_name = kwargs.get("project_name")
    if not project_name:
        return {"error": "project_name is required for 'metrics' type"}

    project_gen = _ensure_project_gen()
    metrics_gen = _ensure_metrics_gen()

    project = project_gen.get_project(project_name)
    if not project:
        return {"success": False, "error": f"Project '{project_name}' not found"}

    formula_groups = kwargs.get("formula_groups")
    metrics = kwargs.get("metrics")

    if formula_groups:
        groups_data = json.loads(formula_groups)
        metrics_list = metrics_gen.generate_from_formula_groups(groups_data)
    elif metrics:
        metrics_list = json.loads(metrics)
    else:
        return {"success": False, "error": "Either formula_groups or metrics required"}

    metrics_yml = metrics_gen.generate_metrics_yml(metrics_list)
    project.generated_files["models/metrics.yml"] = metrics_yml
    project_gen._save()

    return {
        "success": True,
        "metrics_count": len(metrics_list),
        "metrics": [m.get("name") for m in metrics_list],
        "file_path": "models/metrics.yml",
        "message": f"Generated metrics.yml with {len(metrics_list)} metrics",
    }


def _generate_cicd(settings, **kwargs) -> Dict[str, Any]:
    project_name = kwargs.get("project_name")
    if not project_name:
        return {"error": "project_name is required for 'cicd' type"}

    project_gen = _ensure_project_gen()
    cicd_gen = _ensure_cicd_gen()

    project = project_gen.get_project(project_name)
    if not project:
        return {"success": False, "error": f"Project '{project_name}' not found"}

    platform = kwargs.get("platform", "github_actions")
    platform_map = {
        "github_actions": CiCdPlatform.GITHUB_ACTIONS,
        "github": CiCdPlatform.GITHUB_ACTIONS,
        "gitlab_ci": CiCdPlatform.GITLAB_CI,
        "gitlab": CiCdPlatform.GITLAB_CI,
        "azure_devops": CiCdPlatform.AZURE_DEVOPS,
        "azure": CiCdPlatform.AZURE_DEVOPS,
    }
    platform_enum = platform_map.get(platform.lower(), CiCdPlatform.GITHUB_ACTIONS)

    trigger_branches = kwargs.get("trigger_branches")
    branches = trigger_branches.split(",") if trigger_branches else ["main", "develop"]

    config = CiCdConfig(
        platform=platform_enum,
        trigger_branches=branches,
        dbt_version=kwargs.get("dbt_version", "1.7.0"),
        run_tests=kwargs.get("run_tests", True),
        run_docs=kwargs.get("run_docs", True),
    )

    pipeline_content = cicd_gen.generate_pipeline(config, project.config.name)
    pipeline_path = cicd_gen.get_pipeline_path(platform_enum)

    project.cicd_config = config
    project.generated_files[pipeline_path] = pipeline_content
    project_gen._save()

    return {
        "success": True,
        "platform": platform,
        "file_path": pipeline_path,
        "trigger_branches": branches,
        "features": {"run_tests": config.run_tests, "run_docs": config.run_docs},
        "message": f"Generated {platform} pipeline at {pipeline_path}",
    }


# ============================================================================
# Action/type dispatch dictionaries
# ============================================================================

_PROJECT_ACTIONS = {
    "create": _project_create,
    "validate": _project_validate,
    "export": _project_export,
}

_GENERATE_TYPES = {
    "model": _generate_model,
    "sources": _generate_sources,
    "schema": _generate_schema,
    "metrics": _generate_metrics,
    "cicd": _generate_cicd,
}


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

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


def dispatch_dbt_generate(settings, type: str, **kwargs) -> Dict[str, Any]:
    """Dispatch a dbt_generate type."""
    handler = _GENERATE_TYPES.get(type)
    if not handler:
        return {
            "error": f"Unknown type: '{type}'",
            "valid_types": sorted(_GENERATE_TYPES.keys()),
        }
    try:
        return handler(settings, **kwargs)
    except Exception as e:
        logger.error(f"dbt_generate({type}) failed: {e}")
        return {"error": f"dbt_generate({type}) failed: {e}"}


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

def register_unified_dbt_tools(mcp, settings):
    """Register the 2 unified dbt MCP tools."""

    @mcp.tool()
    def dbt_project(
        action: str,
        name: Optional[str] = None,
        profile: Optional[str] = None,
        project_name: Optional[str] = None,
        target_database: Optional[str] = None,
        target_schema: Optional[str] = None,
        hierarchy_project_id: Optional[str] = None,
        include_cicd: bool = False,
        output_dir: Optional[str] = None,
        as_zip: bool = False,
    ) -> Dict[str, Any]:
        """
        Unified dbt project management tool. Replaces 3 individual tools.

        Actions:
        - create: Create a new dbt project scaffold (requires name, profile)
        - validate: Validate project structure and config (requires project_name)
        - export: Export project to directory or ZIP (requires project_name)

        Args:
            action: The action to perform (create, validate, export)
            name: Project name for create
            profile: dbt profile name for create
            project_name: Project name for validate/export
            target_database: Target database for create
            target_schema: Target schema for create
            hierarchy_project_id: Link to hierarchy project for create
            include_cicd: Include CI/CD workflow for create
            output_dir: Output directory for create/export
            as_zip: Create ZIP archive for export

        Returns:
            Action-specific result dict
        """
        return dispatch_dbt_project(settings, action, **{
            k: v for k, v in {
                "name": name, "profile": profile, "project_name": project_name,
                "target_database": target_database, "target_schema": target_schema,
                "hierarchy_project_id": hierarchy_project_id, "include_cicd": include_cicd,
                "output_dir": output_dir, "as_zip": as_zip,
            }.items() if v is not None and v is not False
        })

    @mcp.tool()
    def dbt_generate(
        type: str,
        project_name: Optional[str] = None,
        model_name: Optional[str] = None,
        model_type: str = "staging",
        source_name: Optional[str] = None,
        source_table: Optional[str] = None,
        ref_models: Optional[str] = None,
        columns: Optional[str] = None,
        case_mappings: Optional[str] = None,
        description: Optional[str] = None,
        mappings: Optional[str] = None,
        tables: Optional[str] = None,
        database: Optional[str] = None,
        schema_name: Optional[str] = None,
        models: Optional[str] = None,
        formula_groups: Optional[str] = None,
        metrics: Optional[str] = None,
        platform: str = "github_actions",
        trigger_branches: Optional[str] = None,
        dbt_version: str = "1.7.0",
        run_tests: bool = True,
        run_docs: bool = True,
    ) -> Dict[str, Any]:
        """
        Unified dbt code generation tool. Replaces 5 individual tools.

        Types:
        - model: Generate a dbt model SQL file (requires project_name, model_name)
        - sources: Generate sources.yml (requires project_name, source_name)
        - schema: Generate schema.yml (requires project_name)
        - metrics: Generate metrics.yml (requires project_name)
        - cicd: Generate CI/CD pipeline (requires project_name)

        Args:
            type: The generation type (model, sources, schema, metrics, cicd)
            project_name: dbt project name (all types)
            model_name: Model name (for model type)
            model_type: staging, intermediate, dimension, fact (for model type)
            source_name: Source name (for model/sources type)
            source_table: Source table (for model type)
            ref_models: Comma-separated model refs (for model type)
            columns: JSON list of columns (for model type)
            case_mappings: JSON list of CASE mappings (for model type)
            description: Description (for model type)
            mappings: JSON hierarchy mappings (for sources type)
            tables: JSON table definitions (for sources type)
            database: Database name (for sources type)
            schema_name: Schema name (for sources type)
            models: JSON list of model names (for schema type)
            formula_groups: JSON formula groups (for metrics type)
            metrics: JSON metric definitions (for metrics type)
            platform: CI/CD platform (for cicd type)
            trigger_branches: Comma-separated branches (for cicd type)
            dbt_version: dbt version (for cicd type)
            run_tests: Run tests in CI/CD (for cicd type)
            run_docs: Generate docs in CI/CD (for cicd type)

        Returns:
            Type-specific result dict
        """
        return dispatch_dbt_generate(settings, type, **{
            k: v for k, v in {
                "project_name": project_name, "model_name": model_name,
                "model_type": model_type, "source_name": source_name,
                "source_table": source_table, "ref_models": ref_models,
                "columns": columns, "case_mappings": case_mappings,
                "description": description, "mappings": mappings,
                "tables": tables, "database": database, "schema_name": schema_name,
                "models": models, "formula_groups": formula_groups, "metrics": metrics,
                "platform": platform, "trigger_branches": trigger_branches,
                "dbt_version": dbt_version, "run_tests": run_tests, "run_docs": run_docs,
            }.items() if v is not None
        })

    logger.info("Registered 2 unified dbt tools: dbt_project, dbt_generate")
    return ["dbt_project", "dbt_generate"]
