# --- Standard Library Imports ---
import sys
import logging
import json
import re
from pathlib import Path
from typing import Dict, Optional, List, Tuple, Any


import argparse

# --- Third-party Library Imports ---
from mysql.connector import Error
from mcp.server.fastmcp import FastMCP

# --- Local Application/Library Specific Imports ---
from mcp_db.db import DatabaseCredentials, SelectQuery
from mcp_db.connection_pool import ConnectionPool
from mcp_db.security import SecurityManager, enforce_security
from mcp_db.metrics import MetricsTracker, track_query_metrics

# ==============================================================================
# Logging Configuration
# ==============================================================================
log_dir = Path.home() / ".mcp"
log_dir.mkdir(parents=True, exist_ok=True)
log_file = log_dir / "mcp-db.log"
logging.basicConfig(
    level=logging.INFO, 
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout), # Log to console
        logging.FileHandler(log_file, encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)

# ==============================================================================
# Custom Error Class
# ==============================================================================
class MCPError(Exception):
    """Custom exception class for MCP-specific errors."""
    def __init__(self, message: str, details: Optional[Dict] = None):
        self.message = message
        self.details = details or {}
        super().__init__(self.message)

    def to_dict(self) -> Dict:
        return {"error": self.message, "details": self.details}

# ==============================================================================
# Error Handling Decorator
# ==============================================================================
def handle_error(func):
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except MCPError as e:
            logger.error(f"MCP Error in {func.__name__}: {e.message}", exc_info=False)
            logger.debug(f"MCP Error Details: {e.details}")
            try:
                return json.dumps(e.to_dict())
            except Exception as serialization_error:
                logger.critical(f"Failed to serialize MCPError: {serialization_error}", exc_info=True)
                return json.dumps({"error": "Internal Server Error: Failed to report specific error."})
        except Error as db_error:
            logger.error(f"Database Error in {func.__name__}: {db_error}", exc_info=True)
            error_dict = {
                "error": f"Database operation failed: {db_error.msg}",
                "details": {
                    "type": type(db_error).__name__,
                    "errno": db_error.errno,
                    "sqlstate": db_error.sqlstate,
                }
            }
            try:
                return json.dumps(error_dict)
            except Exception as serialization_error:
                logger.critical(f"Failed to serialize Database Error details: {serialization_error}", exc_info=True)
                return json.dumps({"error": "Internal Server Error: Failed to report database error details."})
        except Exception as e:
            logger.error(f"Unexpected Error in {func.__name__}: {e}", exc_info=True)
            error_dict = {
                "error": "An unexpected internal server error occurred.",
                "details": {
                    "type": type(e).__name__,
                    "message": str(e),
                }
            }
            try:
                return json.dumps(error_dict)
            except Exception as serialization_error:
                logger.critical(f"CRITICAL: Failed to serialize unexpected error details: {serialization_error}", exc_info=True)
                return json.dumps({"error": "Internal Server Error: Failed to report error details."})
    wrapper.__name__ = func.__name__
    wrapper.__doc__ = func.__doc__
    wrapper.__annotations__ = func.__annotations__
    return wrapper

# ==============================================================================
# Environment Loading
# ==============================================================================
def load_env_file(env_path: str) -> Dict[str, str]:
    if not env_path:
        logger.critical("No .env file path provided. Use --env to specify the path.")
        raise MCPError("Configuration Error: .env file path is required.")
    env_file = Path(env_path)
    if not env_file.is_file():
        logger.critical(f".env file not found at specified path: {env_path}")
        raise MCPError(f"Configuration Error: .env file not found at {env_path}")
    env_vars = {}
    with open(env_file, "r") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith('#'):
                continue
            if '=' not in line:
                continue
            key, value = line.split('=', 1)
            env_vars[key.strip()] = value.strip()
    logger.info(f"Successfully loaded environment variables from {env_path}")
    return env_vars

# ==============================================================================
# Main Application Setup
# ==============================================================================
def run_server(env_file_path: str):
    try:
        env_vars = load_env_file(env_file_path)
        logger.info("Initializing MCP server...")
        mcp = FastMCP(
            "MySQL Navigator MCP",
            description="A powerful MySQL/MariaDB database navigation tool using MCP.",
            version="1.0.0"
        )
        logger.info("MCP server framework initialized.")
        pool = ConnectionPool()
        security_manager = SecurityManager()
        metrics = MetricsTracker()
        logger.info("Core services (Connection Pool, Security, Metrics) initialized.")
        try:
            db_port = env_vars.get("DB_PORT")
            db_credentials = DatabaseCredentials(
                host=env_vars.get("DB_HOST", "localhost"),
                port=int(db_port) if db_port and db_port.isdigit() else 3306,
                database=env_vars.get("DB_NAME"),
                username=env_vars.get("DB_USER"),
                password=env_vars.get("DB_PASSWORD"),
                ssl_ca=env_vars.get("DB_SSL_CA"),
                max_retries=int(env_vars.get("DB_MAX_RETRIES", "3")),
                retry_delay=float(env_vars.get("DB_RETRY_DELAY", "1.0"))
            )
            security_manager.sanitize_connection_params(db_credentials)
            logger.info(f"Database credentials loaded and validated for initial database: {db_credentials.database}")
        except Exception as e:
            logger.critical(f"Failed to load or validate initial database credentials from environment: {e}", exc_info=True)
            raise MCPError(f"Configuration Error: Invalid database credentials in .env: {e}")
        try:
            pool.initialize_pool(db_credentials)
            logger.info(f"Connection pool initialized for database: {db_credentials.database}")
        except Exception as e:
            logger.critical(f"Failed to initialize connection pool: {e}", exc_info=True)
            raise MCPError(f"Initialization Error: Could not establish initial database connection pool: {e}")

        # MCP Tool Definitions
        @mcp.tool()
        @handle_error
        def connect_to_database() -> str:
            """Connect to a mysql or mariadb database

            Returns:
                The an object with connection status, database name, and host
            """
            global db_credentials
            logger.info(f"Attempting to connect/verify connection to database: {db_credentials.database}")
            try:
                pool.close_all()
                pool.initialize_pool(db_credentials)
                conn = pool.get_connection()
                conn.ping(reconnect=True, attempts=1, delay=0)
                conn.close()
                logger.info(f"Successfully connected/verified connection to database: {db_credentials.database}")
                status = {"status": "connected", "database": db_credentials.database, "host": db_credentials.host}
                return json.dumps(status)
            except Error as db_err:
                logger.error(f"Database connection/verification failed: {db_err}", exc_info=True)
                raise MCPError(f"Failed to connect to database '{db_credentials.database}'", {"details": str(db_err), "errno": db_err.errno})
            except Exception as e:
                logger.error(f"Unexpected error during connect_to_database: {e}", exc_info=True)
                raise MCPError("Failed to connect to database due to an unexpected error.", {"details": str(e)})

        @mcp.tool()
        @handle_error
        def switch_database(database: str) -> str:
            """Switch to a different mysql or mariadb database

            Returns:
                The an object with connection status, database name, and host
            """
            global db_credentials
            logger.info(f"Attempting to switch database connection to: {database}")
            if not database or not isinstance(database, str):
                raise MCPError("Invalid input: 'database' parameter must be a non-empty string.")
            try:
                new_credentials = DatabaseCredentials(
                    host=db_credentials.host,
                    port=db_credentials.port,
                    database=database.strip(),
                    username=db_credentials.username,
                    password=db_credentials.password,
                    ssl_ca=db_credentials.ssl_ca,
                    max_retries=db_credentials.max_retries,
                    retry_delay=db_credentials.retry_delay
                )
                security_manager.sanitize_connection_params(new_credentials)
                pool.close_all()
                pool.initialize_pool(new_credentials)
                conn = pool.get_connection()
                conn.ping(reconnect=True, attempts=1, delay=0)
                conn.close()
                db_credentials = new_credentials
                logger.info(f"Successfully switched connection pool to database: {database}")
                status = {"status": "switched", "database": database, "host": db_credentials.host}
                return json.dumps(status)
            except Error as db_err:
                logger.error(f"Failed to switch to database '{database}': {db_err}", exc_info=True)
                raise MCPError(f"Failed to switch to database '{database}'", {"details": str(db_err), "errno": db_err.errno})
            except Exception as e:
                logger.error(f"Unexpected error during switch_database: {e}", exc_info=True)
                raise MCPError(f"Failed to switch database due to an unexpected error.", {"details": str(e)})

        @mcp.tool()
        @handle_error
        @enforce_security
        def load_database_schema() -> str:
            """Load the schema for the current database

            Returns:
                The schema for the current database
            """
            logger.info(f"Loading schema for database: {db_credentials.database}")
            connection = None
            try:
                connection = pool.get_connection()
                cursor = connection.cursor()
                cursor.execute("SHOW TABLES")
                tables = [table[0] for table in cursor.fetchall()]
                logger.debug(f"Found tables: {tables}")
                schema = {}
                for table_name in tables:
                    if not re.match(r'^[a-zA-Z0-9_]+$', table_name):
                        logger.warning(f"Skipping potentially unsafe table name for DESCRIBE: {table_name}")
                        continue
                    cursor.execute(f"DESCRIBE `{table_name}`")
                    columns = cursor.fetchall()
                    schema[table_name] = [{"name": col[0], "type": col[1]} for col in columns]
                    logger.debug(f"Schema for table '{table_name}': {schema[table_name]}")
                logger.info(f"Successfully loaded schema for {len(schema)} tables.")
                return json.dumps({"schema": schema})
            except Error as db_err:
                logger.error(f"Error getting database schema: {db_err}", exc_info=True)
                raise MCPError("Error getting database schema", {"details": str(db_err), "errno": db_err.errno})
            except Exception as e:
                logger.error(f"Unexpected error loading schema: {e}", exc_info=True)
                raise MCPError("Unexpected error loading schema.", {"details": str(e)})
            finally:
                if connection and connection.is_connected():
                    cursor.close()
                    connection.close()
                    logger.debug("Database connection closed after loading schema.")

        @handle_error
        @enforce_security
        @track_query_metrics
        @mcp.tool()
        def query_database(query_params: dict) -> List[Tuple[Any, ...]]:
            """Query a mysql or mariadb database

            Args:
                query_params: A dictionary containing:
                    - select_fields: List of fields to select (default: ["*"])
                    - table_name: Name of the table to query from (required)
                    - where_conditions: Dictionary of field-value pairs for WHERE clause (optional)
                    - order_by: List of fields to order by (optional)
                    - order_direction: Sort direction "ASC" or "DESC" (default: "ASC")
                    - limit: Number of records to return (optional)
                    - offset: Number of records to skip (optional)
                    - group_by: List of fields to group by (optional)
                    - having: Dictionary of field-value pairs for HAVING clause (optional)
                    - join_table: Name of the table to join with (optional)
                    - join_type: Type of JOIN operation (default: "INNER")
                    - join_conditions: Dictionary of join conditions (optional)

            Returns:
                The results of the query
            """
            # Validate query_params using SelectQuery
            try:
                validated_params = SelectQuery(**query_params)
            except Exception as e:
                logger.error(f"SelectQuery validation failed: {e}")
                raise MCPError(f"Query parameter validation failed: {e}")
            # Use validated_params.dict() instead of query_params from here on
            query_params = validated_params.model_dump()
            logger.info(f"Received query request for database: {db_credentials.database}")
            logger.debug(f"Query parameters received: {query_params}")
            connection = None
            try:
                table_name = query_params.get("table_name")
                if not table_name or not isinstance(table_name, str) or not re.match(r'^[a-zA-Z0-9_]+$', table_name):
                    raise MCPError("Invalid or missing 'table_name'. Must be alphanumeric/underscore.")
                select_fields_input = query_params.get("select_fields", ["*"])
                if not isinstance(select_fields_input, list) or not select_fields_input:
                    raise MCPError("'select_fields' must be a non-empty list.")
                select_fields = []
                for field in select_fields_input:
                    if field == "*":
                        select_fields.append("*")
                    elif isinstance(field, str) and re.match(r'^[a-zA-Z0-9_]+$', field):
                        select_fields.append(f"`{field}`")
                    else:
                        raise MCPError(f"Invalid character in select field: {field}")
                select_clause = ", ".join(select_fields)
                where_conditions = query_params.get("where_conditions")
                order_by_input = query_params.get("order_by")
                order_direction_input = query_params.get("order_direction", "ASC").upper()
                limit_input = query_params.get("limit")
                offset_input = query_params.get("offset")
                query_parts = [f"SELECT {select_clause} FROM `{table_name}`"]
                parameters = []
                if isinstance(where_conditions, dict) and where_conditions:
                    where_clauses = []
                    for key, value in where_conditions.items():
                        if not isinstance(key, str) or not re.match(r'^[a-zA-Z0-9_]+$', key):
                            raise MCPError(f"Invalid character in where condition key: {key}")
                        where_clauses.append(f"`{key}` = %s")
                        parameters.append(value)
                    if where_clauses:
                        query_parts.append("WHERE " + " AND ".join(where_clauses))
                if isinstance(order_by_input, list) and order_by_input:
                    order_by_clauses = []
                    for field in order_by_input:
                        if not isinstance(field, str) or not re.match(r'^[a-zA-Z0-9_]+$', field):
                            raise MCPError(f"Invalid character in order_by field: {field}")
                        order_by_clauses.append(f"`{field}`")
                    if order_by_clauses:
                        if order_direction_input not in ["ASC", "DESC"]:
                            raise MCPError("Invalid 'order_direction'. Must be 'ASC' or 'DESC'.")
                        query_parts.append(f"ORDER BY {', '.join(order_by_clauses)} {order_direction_input}")
                effective_limit = 100
                if limit_input is not None:
                    try:
                        limit_val = int(limit_input)
                        if limit_val < 1:
                            raise MCPError("'limit' must be a positive integer.")
                        effective_limit = min(limit_val, 500)
                    except (ValueError, TypeError):
                        raise MCPError("'limit' must be an integer.")
                query_parts.append("LIMIT %s")
                parameters.append(effective_limit)
                if offset_input is not None:
                    try:
                        offset_val = int(offset_input)
                        if offset_val < 0:
                            raise MCPError("'offset' must be a non-negative integer.")
                        query_parts.append("OFFSET %s")
                        parameters.append(offset_val)
                    except (ValueError, TypeError):
                        raise MCPError("'offset' must be an integer.")
                final_query = " ".join(query_parts)
                logger.info(f"Executing parameterized query: {final_query}")
                logger.debug(f"Parameters: {tuple(parameters)}")
                connection = pool.get_connection()
                cursor = connection.cursor(dictionary=True)
                cursor.execute(final_query, tuple(parameters))
                results = cursor.fetchall()
                logger.info(f"Query executed successfully. Retrieved {len(results)} rows.")
                return json.dumps({"results": results, "row_count": len(results)})
            except Error as db_err:
                logger.error(f"Database query execution failed: {db_err}", exc_info=True)
                raise MCPError(f"Database query failed: {db_err.msg}", {"details": str(db_err), "errno": db_err.errno, "sqlstate": db_err.sqlstate})
            except MCPError:
                raise
            except Exception as e:
                logger.error(f"Unexpected error during query execution: {e}", exc_info=True)
                raise MCPError("An unexpected error occurred while executing the query.", {"details": str(e)})
            finally:
                if connection and connection.is_connected():
                    if 'cursor' in locals() and cursor:
                        cursor.close()
                    connection.close()
                    logger.debug("Database connection closed after query.")

        logger.info("Starting MySQL Navigator MCP...")
        try:
            if pool._pool is None:
                logger.critical("Connection pool not available. Aborting server start.")
                sys.exit("Error: Database connection pool failed to initialize.")
            logger.info("MCP server setup complete. Ready to accept connections.")
            mcp.run()
        except KeyboardInterrupt:
            logger.info("Received KeyboardInterrupt (Ctrl+C). Initiating shutdown...")
        except MCPError as e:
            logger.critical(f"MCP Initialization Error: {e.message} - Details: {e.details}")
            sys.exit(f"Error: {e.message}")
        except Exception as e:
            logger.critical(f"Failed to start MCP server: {e}", exc_info=True)
            sys.exit(f"Error: Failed to start server - {e}")
        finally:
            logger.info("Initiating server shutdown sequence...")
            if 'pool' in locals() and pool:
                pool.close_all()
                logger.info("Database connection pool closed.")
            else:
                logger.warning("Connection pool object not found during shutdown.")
            logger.info("MySQL Navigator MCP server shutdown complete.")
    except Exception as e:
        logger.critical(f"A critical error occurred outside the main execution block: {e}", exc_info=True)
        sys.exit(f"Critical Error: {e}")

def main_cli():
    parser = argparse.ArgumentParser(description="MySQL Navigator MCP Server")
    parser.add_argument('--config', type=str, required=True, help="Path to .env file with DB credentials")
    args = parser.parse_args()
    run_server(args.config)