#!/usr/bin/env python

import logging
import os
import json
from contextlib import asynccontextmanager
from collections.abc import AsyncIterator
from dataclasses import dataclass
from dotenv import load_dotenv
from typing import Dict, List, Any, Optional, Tuple
from mcp.server.fastmcp import FastMCP, Context
from legion_query_runner import QueryRunner
import sys
import argparse

# Load environment variables
load_dotenv()
logger = logging.getLogger("server")

# Structure for database configuration
@dataclass
class DbConfig:
    id: str
    db_type: str
    configuration: Dict[str, Any]
    description: str
    schema: Optional[List[Dict[str, Any]]] = None
    query_runner: Optional[QueryRunner] = None

# Initialize the Legion Query Runners
def init_config(testing=False, test_db_type=None, test_db_config=None, test_db_configs=None) -> Dict[str, DbConfig]:
    """
    Initialize the Legion Query Runners
    
    Args:
        testing: If True, use test arguments instead of parsing command line
        test_db_type: Test database type (for testing)
        test_db_config: Test database config (for testing)
        test_db_configs: Test database configs (for testing)
    """
    if not testing:
    # Use command line arguments for direct execution
        parser = argparse.ArgumentParser(description='Legion MCP Server')
        parser.add_argument('--db-type', required=False, help='Database type (e.g., mysql, postgresql)')
        parser.add_argument('--db-config', required=False, help='JSON string containing database configuration')
        parser.add_argument('--db-configs', required=False, help='JSON string array containing multiple database configurations')

        args = parser.parse_args()
        
        # Check for multiple database configuration first
        db_configs_str = args.db_configs
        if not db_configs_str:
            db_configs_str = os.getenv("DB_CONFIGS", "")
        
        # Get single database config if needed
        db_type = args.db_type
        db_config_str = args.db_config
    else:
        # For testing, use provided test values
        db_configs_str = test_db_configs
        db_type = test_db_type
        db_config_str = test_db_config
        
        # Fallback to environment variables
        if not db_configs_str:
            db_configs_str = os.getenv("DB_CONFIGS", "")
        if not db_type:
            db_type = os.getenv("DB_TYPE", "")
        if not db_config_str:
            db_config_str = os.getenv("DB_CONFIG", "")
    
    # If we have multiple database configs
    if db_configs_str:
        try:
            db_configs_list = json.loads(db_configs_str) if isinstance(db_configs_str, str) else db_configs_str
            if not isinstance(db_configs_list, list) or len(db_configs_list) == 0:
                raise ValueError("DB_CONFIGS must be a non-empty JSON array")
            
            db_configs = {}
            for i, config in enumerate(db_configs_list):
                if not isinstance(config, dict) or "db_type" not in config or "configuration" not in config or "description" not in config:
                    raise ValueError("Each DB_CONFIG must contain db_type, configuration, and description")
                
                # Generate a meaningful ID if not provided
                db_id = config.get("id", "")
                if not db_id:
                    # Create an ID based on database type and description
                    db_type_abbr = config["db_type"][:2].lower()
                    desc_part = ''.join(c for c in config["description"] if c.isalnum())[:8].lower()
                    db_id = f"{db_type_abbr}_{desc_part}_{i}"
                
                db_config = DbConfig(
                    db_type=config["db_type"],
                    configuration=config["configuration"],
                    description=config["description"],
                    id=db_id
                )
                
                # print(f"Initializing query runner for {db_config.db_type} database: {db_config.description}")
                db_config.query_runner = QueryRunner(db_type=db_config.db_type, configuration=db_config.configuration)
                db_configs[db_id] = db_config
                
            print(f"Initialized {len(db_configs)} database connections")
            return db_configs
        except json.JSONDecodeError as e:
            print(f"Error parsing DB_CONFIGS: {str(e)}")
            raise
    
    # Fallback to single database configuration
    # Only parse args if we're not in MCP CLI mode
    if not db_type:
        db_type = os.getenv("DB_TYPE", "pg")
    if not db_config_str:
        db_config_str = os.getenv("DB_CONFIG", "")
    
    try:
        db_config = json.loads(db_config_str) if isinstance(db_config_str, str) else db_config_str
        if isinstance(db_config, str):
            db_config = json.loads(db_config)  # Handle double-string JSON
    except json.JSONDecodeError as e:
        print(f"Error parsing DB_CONFIG: {str(e)}")
        raise

    # Allowing for empty configuration for now, this will render this MCP server useless
    # but it will not crash
    if not db_type or not db_config:
        print("Database type and configuration are required...")
        return {}

    # print(f"Initializing single query runner for {db_type} database")
    
    # Generate a meaningful ID for the single database
    db_id = f"{db_type.lower()}_default"
    
    db_config_obj = DbConfig(
        db_type=db_type,
        configuration=db_config,
        description="Default database connection",
        id=db_id
    )
    db_config_obj.query_runner = QueryRunner(db_type=db_type, configuration=db_config)
    
    db_configs = {db_id: db_config_obj}
    return db_configs

# Initialize query_runners as a dictionary
config_map: Dict[str, DbConfig] = {} 

# Initialize only if not being imported by test
_is_test = 'pytest' in sys.modules
if not _is_test:
    try:
        config_map = init_config()
    except Exception as e:
        print(f"Error initializing query runners: {str(e)}")
        print("\nUsage:")
        print("1. For MCP CLI mode with single database:")
        print("   Set environment variables: DB_TYPE, DB_CONFIG")
        print("   Or for multiple databases: DB_CONFIGS='[{\"db_type\":\"pg\",\"configuration\":{...},\"description\":\"My PostgreSQL DB\"}]'")
        print("   Then run: mcp install mcp_server.py")
        print("   Or: mcp dev mcp_server.py")
        print("\n2. For direct execution with single database:")
        print("   python mcp_server.py --db-type <db_type> --db-config '<json_config>'")
        print("   Example: python mcp_server.py --db-type mysql --db-config '{\"host\":\"localhost\",\"port\":3306,\"user\":\"root\",\"password\":\"pass\",\"database\":\"test\"}'")
        print("\n3. For direct execution with multiple databases:")
        print("   python mcp_server.py --db-configs '[{\"db_type\":\"pg\",\"configuration\":{\"host\":\"localhost\"},\"description\":\"My PostgreSQL DB\"}]'")
        sys.exit(1)

    # Fetch schema information for all databases
    for db_config in config_map.values():
        try:
            db_config.schema = db_config.query_runner.get_schema()
        except Exception as e:
            print(f"Warning: Could not fetch schema for {db_config.description}: {str(e)}")


# Define database context that will be available to all handlers
@dataclass
class DbContext:
    db_configs: Dict[str, DbConfig]
    last_query: Optional[str] = None
    last_result: Optional[Dict[str, Any]] = None
    query_history: List[str] = None
    
    def __post_init__(self):
        if self.query_history is None:
            self.query_history = []
    
    def get_default_query_runner(self) -> QueryRunner:
        """Gets the default query runner (first one if multiple exist)"""
        if not self.db_configs:
            raise ValueError("No database connections available")
        # Get the first database in the dictionary
        first_db_id = next(iter(self.db_configs))
        return self.db_configs[first_db_id].query_runner

# Server lifespan manager
@asynccontextmanager
async def db_lifespan(server: FastMCP) -> AsyncIterator[DbContext]:
    """Initialize database connections on startup and provide context to handlers"""
    
    # Initialize context
    db_context = DbContext(
        db_configs=config_map,
    )
    
    try:
        # Test connection on startup for all databases
        for db_config in db_context.db_configs.values():
            try:
                db_config.query_runner.test_connection()
                print(f"Successfully connected to {db_config.description}")
            except Exception as e:
                print(f"Warning: Connection test failed for {db_config.description}: {str(e)}")
        
        yield db_context
    finally:
        # Cleanup could happen here if needed
        pass

# Pass lifespan to server
mcp = FastMCP("Legion Multi-Database Access", lifespan=db_lifespan)

# Define resources
@mcp.resource("resource://schema/{database_id}")
def get_schema(database_id: Optional[str] = None) -> str:
    """Get the database schemas for all databases or a specific database
    
    Args:
        database_id: Optional database ID to get schema for. If None, get schemas for all databases.
    """
    try:
        schemas = []
        
        if database_id is not None and database_id != "all":
            # Get schema for specific database
            if database_id not in config_map:
                return f"Error: Invalid database ID {database_id}"
            
            db_config = config_map[database_id]
            try:
                schema = db_config.query_runner.get_schema()
                schemas.append({
                    "id": database_id,
                    "description": db_config.description,
                    "db_type": db_config.db_type,
                    "schema": schema
                })
            except Exception as e:
                schemas.append({
                    "id": database_id,
                    "description": db_config.description,
                    "db_type": db_config.db_type,
                    "error": str(e)
                })
        else:
            # Get schema for all databases
            for db_id, db_config in config_map.items():
                try:
                    schema = db_config.query_runner.get_schema()
                    schemas.append({
                        "id": db_id,
                        "description": db_config.description,
                        "db_type": db_config.db_type,
                        "schema": schema
                    })
                except Exception as e:
                    schemas.append({
                        "id": db_id,
                        "description": db_config.description,
                        "db_type": db_config.db_type,
                        "error": str(e)
                    })
        
        return json.dumps(schemas)
    except Exception as e:
        return f"Error getting schemas: {str(e)}"

@mcp.tool()
def get_query_history(ctx: Context) -> str:
    """Get the recent query history"""
    db_context = ctx.request_context.lifespan_context
    
    if not db_context.query_history:
        return "No query history available"
    
    result = "Query History:\n"
    for entry in db_context.query_history:
        if isinstance(entry, str):
            result += f"- {entry}\n"
        else:
            result += f"- {entry['query']} (Database: {entry['db_name']})\n"
    
    return result

@mcp.tool()
def list_databases(ctx: Context) -> str:
    """List all available database connections"""
    db_context: DbContext = ctx.request_context.lifespan_context
    
    if not db_context.db_configs:
        return "No database connections available."
    
    db_list = []
    for db_id, db_config in db_context.db_configs.items():
        db_info = f"ID: {db_id} - {db_config.description} (Type: {db_config.db_type})"
        
        # Add table count if schema is available
        if db_config.schema:
            table_count = len(db_config.schema)
            db_info += f" - {table_count} tables"
        
        db_list.append(db_info)
    
    return "Available databases:\n" + "\n".join(db_list)

@mcp.prompt()
def select_database() -> str:
    """Help user select which database to use"""
    return "I need to determine which database to use for your query. Please use the list_databases tool first, then tell me which database ID to use."

def get_database_schema_summary(db_config: DbConfig) -> str:
    """Create a summary of the database schema for display"""
    if db_config.schema is None:
        return "Schema information not available"
    
    tables = db_config.schema
    if not tables:
        return "No tables found in schema"
    
    table_summaries = []
    for table in tables[:10]:  # Limit to first 10 tables
        table_name = table.get("name", "")
        if not table_name:
            continue
        
        columns = table.get("columns", [])
        column_names = [col.get("name", "") for col in columns[:5]]  # Limit to first 5 columns
        column_str = ", ".join(column_names)
        if len(columns) > 5:
            column_str += ", ..."
        
        table_summaries.append(f"- {table_name} ({column_str})")
    
    if len(tables) > 10:
        table_summaries.append(f"... and {len(tables) - 10} more tables")
    
    return "\n".join(table_summaries) if table_summaries else "No tables found in schema"

@mcp.tool()
def get_database_info(ctx: Context, db_id: Optional[str] = None) -> str:
    """Get detailed information about a database including schema summary"""
    db_context: DbContext = ctx.request_context.lifespan_context
    
    if db_id is None:
        # Return info for all databases
        all_info = []
        for curr_db_id, db_config in db_context.db_configs.items():
            info = f"Database ID: {curr_db_id}\n"
            info += f"Description: {db_config.description}\n"
            info += f"Type: {db_config.db_type}\n"
            info += f"Schema Summary:\n{get_database_schema_summary(db_config)}"
            all_info.append(info)
        
        return "\n\n".join(all_info)
    
    # Return info for specific database
    if db_id not in db_context.db_configs:
        return f"Invalid database ID: {db_id}"
    
    db_config = db_context.db_configs[db_id]
    info = f"Database ID: {db_id}\n"
    info += f"Description: {db_config.description}\n"
    info += f"Type: {db_config.db_type}\n"
    info += f"Schema Summary:\n{get_database_schema_summary(db_config)}"
    
    return info

def _execute_and_get_results(query: str, ctx: Context, db_id: str) -> Dict[str, Any]:
    """Helper function to execute query and get formatted results"""
    db_context = ctx.request_context.lifespan_context
    
    if db_id not in db_context.db_configs:
        raise ValueError(f"Invalid database ID: {db_id}")
    
    # Get the selected query runner
    db_config = db_context.db_configs[db_id]
    query_runner = db_config.query_runner
    
    # Execute query
    result = query_runner.run_query(query)
    
    # Update query history in the context
    db_context.last_query = query
    db_context.last_result = result
    db_context.query_history.append(f"[{db_id}] [{db_config.description}] {query}")
    
    # Extract column info
    columns = result.get('columns', [])
    column_names = [col.get('friendly_name', col.get('name', '')) for col in columns]
    
    # Extract row data
    rows = result.get('rows', [])
    row_count = len(rows)
    
    # Process rows - each row is a dictionary with column names as keys
    processed_rows = []
    for row_dict in rows:
        # Create a row with values in the same order as column_names
        processed_row = [row_dict.get(col.get('name', '')) for col in columns]
        processed_rows.append(processed_row)
    
    return {
        'column_names': column_names,
        'columns': columns,
        'rows': processed_rows,
        'raw_rows': rows,
        'row_count': row_count,
        'database': {
            'id': db_id,
            'description': db_config.description,
            'db_type': db_config.db_type
        }
    }

# Define tools
@mcp.tool()
def execute_query(query: str, ctx: Context, db_id: str) -> str:
    """Execute a SQL query and return results as a markdown table"""
    try:
        result = _execute_and_get_results(query, ctx, db_id)
        
        # Add database info to the output
        db_info = f"Database: {result['database']['description']} (Type: {result['database']['db_type']})"
        
        # Build a markdown table for output
        header = " | ".join(result['column_names'])
        separator = " | ".join(["---"] * len(result['column_names']))
        
        table_rows = []
        for row in result['rows'][:10]:  # Limit to first 10 rows for display
            table_rows.append(" | ".join(str(cell) for cell in row))
        
        result_table = f"Query executed on Database: {result['database']['description']}\n\n"
        result_table += f"{header}\n{separator}\n" + "\n".join(table_rows)
        
        if result['row_count'] > 10:
            result_table += f"\n\n... and {result['row_count'] - 10} more rows (total: {result['row_count']})"
            
        return f"{db_info}\n\n{result_table}"
    except ValueError as e:
        if "Invalid database ID" in str(e):
            return f"Error: {str(e)}"
        return f"Error executing query: {str(e)}"
    except Exception as e:
        return f"Error executing query: {str(e)}"

@mcp.tool()
def execute_query_json(query: str, ctx: Context, db_id: str) -> str:
    """Execute a SQL query and return results as JSON"""
    try:
        result = _execute_and_get_results(query, ctx, db_id)
        
        # Create a more compact representation for JSON output
        output = {
            'database': result['database'],
            'columns': result['column_names'],
            'rows': result['raw_rows'],  # Return the original row dictionaries for JSON output
            'row_count': result['row_count']
        }
        return json.dumps(output, indent=2)
    except Exception as e:
        return f"Error executing query: {str(e)}"

@mcp.tool()
def get_table_columns(table_name: str, ctx: Context, db_id: str) -> str:
    """Get column names for a specific table"""
    try:
        db_context: DbContext = ctx.request_context.lifespan_context
        
        if db_id not in db_context.db_configs:
            raise ValueError(f"Invalid database ID: {db_id}")
        
        db_config = db_context.db_configs[db_id]
        columns = db_config.query_runner.get_table_columns(table_name)
        
        return json.dumps({
            'database': {
                'id': db_id,
                'description': db_config.description,
                'db_type': db_config.db_type
            },
            'table': table_name,
            'columns': columns
        })
    except Exception as e:
        return f"Error getting columns for table {table_name}: {str(e)}"

@mcp.tool()
def get_table_types(table_name: str, ctx: Context, db_id: str) -> str:
    """Get column types for a specific table"""
    try:
        db_context: DbContext = ctx.request_context.lifespan_context
        
        if db_id not in db_context.db_configs:
            raise ValueError(f"Invalid database ID: {db_id}")
        
        db_config = db_context.db_configs[db_id]
        types = db_config.query_runner.get_table_types(table_name)
        
        return json.dumps({
            'database': {
                'id': db_id,
                'description': db_config.description,
                'db_type': db_config.db_type
            },
            'table': table_name,
            'types': types
        })
    except Exception as e:
        return f"Error getting types for table {table_name}: {str(e)}"

@mcp.tool()
def describe_table(ctx: Context, table_name: str, db_id: str) -> str:
    """Get detailed description of a table including column names and types"""
    try:
        db_context: DbContext = ctx.request_context.lifespan_context
        
        if db_id not in db_context.db_configs:
            return f"Error: Invalid database ID {db_id}"
        
        db_config = db_context.db_configs[db_id]
        
        # Get column names and types
        columns = db_config.query_runner.get_table_columns(table_name)
        types = db_config.query_runner.get_table_types(table_name)
        
        # Build description
        description = f"Table: {table_name} in Database: {db_config.description} (ID: {db_id})\n\n"
        description += "Columns:\n"
        
        for column in columns:
            column_type = types.get(column, "unknown")
            description += f"- {column} ({column_type})\n"
        
        return description
    except Exception as e:
        return f"Error describing table {table_name}: {str(e)}"

@mcp.tool()
def get_table_sample(ctx: Context, table_name: str, db_id: str, limit: int = 10) -> str:
    """Get a sample of data from a table"""
    try:
        db_context: DbContext = ctx.request_context.lifespan_context
        
        if db_id not in db_context.db_configs:
            return f"Error: Invalid database ID {db_id}"
        
        db_config = db_context.db_configs[db_id]
        query_runner = db_config.query_runner
        
        # Construct a safe query to sample data
        query = f"SELECT * FROM {table_name} LIMIT {min(limit, 100)}"
        result = query_runner.run_query(query)
        
        # Extract column info
        columns = result.get('columns', [])
        column_names = []
        for col in columns:
            if isinstance(col, dict):
                column_names.append(col.get('friendly_name', col.get('name', '')))
            else:
                column_names.append(col)
        
        # Format as markdown table
        header = " | ".join(column_names)
        separator = " | ".join(["---"] * len(column_names))
        
        # Extract row data
        rows = result.get('rows', [])
        
        table_rows = []
        for row in rows:
            if isinstance(row, dict):
                # Handle dict rows
                row_values = [str(row.get(col, '')) for col in column_names]
            else:
                # Handle list rows
                row_values = [str(cell) for cell in row]
            table_rows.append(" | ".join(row_values))
        
        sample_data = f"Sample data from table '{table_name}' in Database: {db_config.description} (ID: {db_id})\n\n"
        sample_data += f"{header}\n{separator}\n" + "\n".join(table_rows)
        
        if not rows:
            sample_data += "\n\nNo data found in table."
        
        return sample_data
    except Exception as e:
        return f"Error getting sample data from table {table_name}: {str(e)}"

@mcp.tool()
def find_table(table_name: str, ctx: Context) -> str:
    """Find which database contains a specific table"""
    db_context: DbContext = ctx.request_context.lifespan_context
    found_in = []
    
    for db_id, db_config in db_context.db_configs.items():
        if not db_config.schema:
            continue
        
        for table in db_config.schema:
            if table.get("name") == table_name:
                found_in.append({
                    "db_id": db_id,
                    "db_name": db_config.description,
                    "db_type": db_config.db_type
                })
                break
    
    if not found_in:
        return f"Table '{table_name}' was not found in any database schema."
    
    result = f"Table '{table_name}' was found in the following databases:\n"
    for db in found_in:
        result += f"- Database ID: {db['db_id']} - {db['db_name']} (Type: {db['db_type']})\n"
    
    return result

# Define prompts
@mcp.prompt()
def sql_query() -> str:
    """Create an SQL query against the database"""
    return "Please help me write a SQL query for the following question:\n\n"

@mcp.prompt()
def explain_query(query: str) -> str:
    """Explain what a SQL query does"""
    return f"Can you explain what the following SQL query does?\n\n```sql\n{query}\n```"

@mcp.prompt()
def optimize_query(query: str) -> str:
    """Optimize a SQL query for better performance"""
    return f"Can you optimize the following SQL query for better performance?\n\n```sql\n{query}\n```"

def main():
    # print(f"Starting Legion Multi-Database MCP server with {len(query_runners)} database connections...")
    mcp.run()

if __name__ == "__main__":
    main() 