"""
Unified MCP tools for DataShield (Phase 3D-6 consolidation).

Consolidates 12 individual datashield tools into 3 action-dispatched tools:
- shield_project(action, ...) — 4 actions: create, list, get, delete
- shield_table(action, ...) — 4 actions: auto_classify, add, remove, preview
- shield_deploy(action, ...) — 4 actions: generate_ddl, deploy, shield_file, status

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

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

logger = logging.getLogger(__name__)

# ============================================================================
# Lazy singleton
# ============================================================================

_service = None


def _ensure_service(settings):
    global _service
    if _service is None:
        from .service import ShieldService
        data_dir = Path(settings.data_dir) / "datashield"
        _service = ShieldService(data_dir=str(data_dir))
    return _service


# ============================================================================
# shield_project action handlers
# ============================================================================

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

    service = _ensure_service(settings)
    try:
        project = service.create_project(
            name=name,
            passphrase=passphrase,
            description=kwargs.get("description") or None,
        )
        return {
            "status": "success",
            "project": {
                "id": project.id,
                "name": project.name,
                "key_alias": project.key_alias,
                "created_at": project.created_at,
            },
            "message": f"Shield project '{name}' created. Add tables with shield_table(action='add', ...).",
        }
    except Exception as e:
        logger.error("Failed to create shield project: %s", e)
        return {"error": str(e)}


def _project_list(settings, **kwargs) -> Dict[str, Any]:
    service = _ensure_service(settings)
    projects = service.list_projects()
    return {
        "status": "success",
        "count": len(projects),
        "projects": projects,
    }


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

    service = _ensure_service(settings)
    project = service.get_project(project_id)
    if not project:
        return {"error": f"Project not found: {project_id}"}

    return {
        "status": "success",
        "project": project.model_dump(),
    }


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

    service = _ensure_service(settings)
    try:
        deleted = service.delete_project(project_id)
        if deleted:
            return {"status": "success", "message": f"Project {project_id} deleted"}
        return {"error": f"Project not found: {project_id}"}
    except Exception as e:
        return {"error": str(e)}


# ============================================================================
# shield_table action handlers
# ============================================================================

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

    try:
        from .classifier import auto_classify_columns

        cols = json.loads(columns)
        sample_data = kwargs.get("sample_data")
        samples = json.loads(sample_data) if sample_data else None
        row_count = kwargs.get("row_count", 0)

        rules = auto_classify_columns(cols, samples, row_count)

        return {
            "status": "success",
            "column_count": len(rules),
            "rules": [r.model_dump() for r in rules],
            "message": "Review and adjust classifications before adding to a shield project.",
        }
    except json.JSONDecodeError as e:
        return {"error": f"Invalid JSON: {e}"}
    except Exception as e:
        return {"error": str(e)}


def _table_add(settings, **kwargs) -> Dict[str, Any]:
    project_id = kwargs.get("project_id")
    database = kwargs.get("database")
    schema_name = kwargs.get("schema_name")
    table_name = kwargs.get("table_name")
    column_rules = kwargs.get("column_rules")
    if not all([project_id, database, schema_name, table_name, column_rules]):
        return {"error": "project_id, database, schema_name, table_name, column_rules are required for 'add' action"}

    try:
        from .types import ColumnRule, TableShieldConfig

        rules = [ColumnRule(**r) for r in json.loads(column_rules)]
        key_columns = json.loads(kwargs.get("key_columns", "[]"))
        skip_columns = json.loads(kwargs.get("skip_columns", "[]"))

        config = TableShieldConfig(
            database=database,
            schema_name=schema_name,
            table_name=table_name,
            table_type=kwargs.get("table_type", "unknown"),
            column_rules=rules,
            key_columns=key_columns,
            skip_columns=skip_columns,
        )

        service = _ensure_service(settings)
        service.add_table_shield(project_id, config)

        return {
            "status": "success",
            "table": f"{database}.{schema_name}.{table_name}",
            "rules_count": len(rules),
            "message": f"Table shield added with {len(rules)} column rules.",
        }
    except json.JSONDecodeError as e:
        return {"error": f"Invalid JSON: {e}"}
    except Exception as e:
        return {"error": str(e)}


def _table_remove(settings, **kwargs) -> Dict[str, Any]:
    project_id = kwargs.get("project_id")
    database = kwargs.get("database")
    schema_name = kwargs.get("schema_name")
    table_name = kwargs.get("table_name")
    if not all([project_id, database, schema_name, table_name]):
        return {"error": "project_id, database, schema_name, table_name are required for 'remove' action"}

    try:
        service = _ensure_service(settings)
        removed = service.remove_table_shield(
            project_id, database, schema_name, table_name
        )
        if removed:
            return {
                "status": "success",
                "message": f"Removed {database}.{schema_name}.{table_name} from project",
            }
        return {"error": "Table not found in project"}
    except Exception as e:
        return {"error": str(e)}


def _table_preview(settings, **kwargs) -> Dict[str, Any]:
    project_id = kwargs.get("project_id")
    passphrase = kwargs.get("passphrase")
    database = kwargs.get("database")
    schema_name = kwargs.get("schema_name")
    table_name = kwargs.get("table_name")
    if not all([project_id, passphrase, database, schema_name, table_name]):
        return {"error": "project_id, passphrase, database, schema_name, table_name are required for 'preview' action"}

    try:
        service = _ensure_service(settings)
        project = service.get_project(project_id)
        if not project:
            return {"error": f"Project not found: {project_id}"}

        # Find table config
        fqn = f"{database}.{schema_name}.{table_name}"
        table_config = None
        for t in project.tables:
            if f"{t.database}.{t.schema_name}.{t.table_name}" == fqn:
                table_config = t
                break
        if not table_config:
            return {"error": f"Table {fqn} not found in project"}

        engine = service.get_engine(project_id, passphrase)

        sample_rows = kwargs.get("sample_rows", "[]")
        rows = json.loads(sample_rows) if isinstance(sample_rows, str) else sample_rows
        if not rows:
            return {"error": "No sample rows provided"}

        # Limit to 5 rows
        rows = rows[:5]
        previews = []

        rules_by_col = {r.column_name: r for r in table_config.column_rules}

        for row in rows:
            preview_row = {}
            for col_name, original in row.items():
                rule = rules_by_col.get(col_name)
                if rule and col_name not in table_config.skip_columns:
                    scrambled = engine.scramble(original, rule)
                    preview_row[col_name] = {
                        "original": original,
                        "scrambled": scrambled,
                        "strategy": rule.strategy.value,
                    }
                else:
                    preview_row[col_name] = {
                        "original": original,
                        "scrambled": original,
                        "strategy": "passthrough (no rule)",
                    }
            previews.append(preview_row)

        return {
            "status": "success",
            "table": fqn,
            "preview_rows": len(previews),
            "data": previews,
        }
    except Exception as e:
        return {"error": str(e)}


# ============================================================================
# shield_deploy action handlers
# ============================================================================

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

    try:
        from .snowflake_generator import generate_full_ddl

        service = _ensure_service(settings)
        project = service.get_project(project_id)
        if not project:
            return {"error": f"Project not found: {project_id}"}

        if not project.tables:
            return {"error": "Project has no tables configured. Add tables with shield_table(action='add', ...) first."}

        key_ref = kwargs.get("key_ref")
        ddl = generate_full_ddl(project, key_ref or None)

        return {
            "status": "success",
            "project": project.name,
            "tables": len(project.tables),
            "ddl": ddl,
            "message": "Review the DDL before deploying. Use shield_deploy(action='deploy', ...) to execute.",
        }
    except Exception as e:
        return {"error": str(e)}


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

    try:
        from .snowflake_generator import generate_full_ddl

        service = _ensure_service(settings)
        project = service.get_project(project_id)
        if not project:
            return {"error": f"Project not found: {project_id}"}

        key_ref = kwargs.get("key_ref")
        ddl = generate_full_ddl(project, key_ref or None)

        # Try to get the connection from server context
        try:
            try:
                from src.server import _connections
            except ImportError:
                from server import _connections

            if connection_id not in _connections:
                return {"error": f"Connection not found: {connection_id}. Use get_console_connections() to list available connections."}

            conn = _connections[connection_id]
            engine = conn.get("engine")
            if not engine:
                return {"error": "Connection has no active engine"}

            # Execute DDL statements
            from sqlalchemy import text
            statements = [s.strip() for s in ddl.split(";") if s.strip()]
            executed = 0
            with engine.connect() as connection:
                for stmt in statements:
                    if stmt.startswith("--"):
                        continue
                    connection.execute(text(stmt))
                    executed += 1
                connection.commit()

            return {
                "status": "success",
                "statements_executed": executed,
                "message": f"Deployed {executed} DDL statements to Snowflake via {connection_id}",
            }
        except ImportError:
            return {
                "error": "Database connections not available. Use shield_deploy(action='generate_ddl', ...) to get the DDL and execute manually.",
            }
    except Exception as e:
        return {"error": str(e)}


def _deploy_shield_file(settings, **kwargs) -> Dict[str, Any]:
    project_id = kwargs.get("project_id")
    passphrase = kwargs.get("passphrase")
    input_path = kwargs.get("input_path")
    output_path = kwargs.get("output_path")
    if not all([project_id, passphrase, input_path, output_path]):
        return {"error": "project_id, passphrase, input_path, output_path are required for 'shield_file' action"}

    try:
        from .interceptor import DataShieldInterceptor

        service = _ensure_service(settings)
        project = service.get_project(project_id)
        if not project:
            return {"error": f"Project not found: {project_id}"}

        database = kwargs.get("database", "LOCAL")
        schema_name = kwargs.get("schema_name", "FILES")
        table_name = kwargs.get("table_name", "")

        if not table_name:
            table_name = Path(input_path).stem.upper()

        fqn = f"{database}.{schema_name}.{table_name}"
        table_config = None
        for t in project.tables:
            if f"{t.database}.{t.schema_name}.{t.table_name}" == fqn:
                table_config = t
                break
        if not table_config:
            return {"error": f"Table {fqn} not found in project. Add it with shield_table(action='add', ...) first."}

        engine = service.get_engine(project_id, passphrase)
        interceptor = DataShieldInterceptor(engine)

        ext = Path(input_path).suffix.lower()
        if ext == ".csv":
            result = interceptor.shield_csv(input_path, output_path, table_config)
        elif ext == ".json":
            result = interceptor.shield_json(input_path, output_path, table_config)
        else:
            return {"error": f"Unsupported file type: {ext}. Use .csv or .json"}

        return {"status": "success", **result}
    except Exception as e:
        return {"error": str(e)}


def _deploy_status(settings, **kwargs) -> Dict[str, Any]:
    try:
        service = _ensure_service(settings)
        project_id = kwargs.get("project_id")
        status = service.get_status(project_id or None)
        return {"status": "success", **status}
    except Exception as e:
        return {"error": str(e)}


# ============================================================================
# Action dispatch dictionaries
# ============================================================================

_PROJECT_ACTIONS = {
    "create": _project_create,
    "list": _project_list,
    "get": _project_get,
    "delete": _project_delete,
}

_TABLE_ACTIONS = {
    "auto_classify": _table_auto_classify,
    "add": _table_add,
    "remove": _table_remove,
    "preview": _table_preview,
}

_DEPLOY_ACTIONS = {
    "generate_ddl": _deploy_generate_ddl,
    "deploy": _deploy_deploy,
    "shield_file": _deploy_shield_file,
    "status": _deploy_status,
}


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

def dispatch_shield_project(settings, action: str, **kwargs) -> Dict[str, Any]:
    """Dispatch a shield_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"shield_project({action}) failed: {e}")
        return {"error": f"shield_project({action}) failed: {e}"}


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


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


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

def register_unified_datashield_tools(mcp, settings):
    """Register the 3 unified datashield MCP tools."""

    @mcp.tool()
    def shield_project(
        action: str,
        project_id: Optional[str] = None,
        name: Optional[str] = None,
        passphrase: Optional[str] = None,
        description: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Unified DataShield project management tool. Replaces 4 individual tools.

        Actions:
        - create: Create a new shield project (requires name, passphrase)
        - list: List all shield projects
        - get: Get project details (requires project_id)
        - delete: Delete a project and its encryption key (requires project_id)

        Args:
            action: The action to perform (create, list, get, delete)
            project_id: Shield project ID (for get, delete)
            name: Project name (for create)
            passphrase: Passphrase to protect the encryption keystore (for create)
            description: Optional project description (for create)

        Returns:
            Action-specific result dict
        """
        return dispatch_shield_project(settings, action, **{
            k: v for k, v in {
                "project_id": project_id, "name": name,
                "passphrase": passphrase, "description": description,
            }.items() if v is not None
        })

    @mcp.tool()
    def shield_table(
        action: str,
        project_id: Optional[str] = None,
        passphrase: Optional[str] = None,
        database: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
        columns: Optional[str] = None,
        sample_data: Optional[str] = None,
        row_count: int = 0,
        column_rules: Optional[str] = None,
        table_type: str = "unknown",
        key_columns: Optional[str] = None,
        skip_columns: Optional[str] = None,
        sample_rows: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Unified DataShield table configuration tool. Replaces 4 individual tools.

        Actions:
        - auto_classify: Auto-detect column classifications (requires columns JSON)
        - add: Add/update a table shield config (requires project_id, database, schema_name, table_name, column_rules)
        - remove: Remove a table from shield project (requires project_id, database, schema_name, table_name)
        - preview: Preview before/after scrambling (requires project_id, passphrase, database, schema_name, table_name, sample_rows)

        Args:
            action: The action to perform (auto_classify, add, remove, preview)
            project_id: Shield project ID (for add, remove, preview)
            passphrase: Passphrase to unlock the keystore (for preview)
            database: Database name (for add, remove, preview)
            schema_name: Schema name (for add, remove, preview)
            table_name: Table name (for add, remove, preview)
            columns: JSON array of column objects with "name" and "data_type" (for auto_classify)
            sample_data: JSON object of column name to sample value arrays (for auto_classify)
            row_count: Total row count for cardinality analysis (for auto_classify)
            column_rules: JSON array of column rule objects (for add)
            table_type: Table type - "fact", "dimension", or "unknown" (for add)
            key_columns: JSON array of PK/FK column names (for add)
            skip_columns: JSON array of columns to pass through unchanged (for add)
            sample_rows: JSON array of row objects to preview (for preview)

        Returns:
            Action-specific result dict
        """
        return dispatch_shield_table(settings, action, **{
            k: v for k, v in {
                "project_id": project_id, "passphrase": passphrase,
                "database": database, "schema_name": schema_name,
                "table_name": table_name, "columns": columns,
                "sample_data": sample_data, "row_count": row_count,
                "column_rules": column_rules, "table_type": table_type,
                "key_columns": key_columns, "skip_columns": skip_columns,
                "sample_rows": sample_rows,
            }.items() if v is not None
        })

    @mcp.tool()
    def shield_deploy(
        action: str,
        project_id: Optional[str] = None,
        passphrase: Optional[str] = None,
        connection_id: Optional[str] = None,
        key_ref: Optional[str] = None,
        input_path: Optional[str] = None,
        output_path: Optional[str] = None,
        database: Optional[str] = None,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Unified DataShield deployment and operations tool. Replaces 4 individual tools.

        Actions:
        - generate_ddl: Generate Snowflake UDFs and shielded views DDL (requires project_id)
        - deploy: Execute DDL on Snowflake (requires project_id, connection_id)
        - shield_file: Scramble a CSV/JSON file (requires project_id, passphrase, input_path, output_path)
        - status: Get DataShield status overview (optional project_id)

        Args:
            action: The action to perform (generate_ddl, deploy, shield_file, status)
            project_id: Shield project ID (for generate_ddl, deploy, shield_file, status)
            passphrase: Passphrase to unlock the keystore (for shield_file)
            connection_id: Backend connection ID for Snowflake (for deploy)
            key_ref: Key reference for UDF calls (for generate_ddl, deploy)
            input_path: Path to source file (for shield_file)
            output_path: Path to write the shielded output (for shield_file)
            database: Database reference for table lookup (for shield_file, default "LOCAL")
            schema_name: Schema reference for table lookup (for shield_file, default "FILES")
            table_name: Table name reference (for shield_file, uses filename if empty)

        Returns:
            Action-specific result dict
        """
        return dispatch_shield_deploy(settings, action, **{
            k: v for k, v in {
                "project_id": project_id, "passphrase": passphrase,
                "connection_id": connection_id, "key_ref": key_ref,
                "input_path": input_path, "output_path": output_path,
                "database": database, "schema_name": schema_name,
                "table_name": table_name,
            }.items() if v is not None
        })

    logger.info("Registered 3 unified datashield tools: shield_project, shield_table, shield_deploy")
    return _ensure_service(settings)
