"""
Databricks MCP Server - Unity Catalog Explorer

Provides MCP tools for exploring Databricks Unity Catalog:
- Who am I (current user)
- List catalogs
- List schemas in a catalog
- List tables in a schema
- Get table details (columns, metadata)
- Get sample data from a table
"""

import os
import json
import logging
from functools import wraps
from typing import Optional
from mcp.server.fastmcp import FastMCP

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize MCP server
mcp = FastMCP("databricks-mcp-server")

# Global workspace client - will be initialized lazily
_workspace_client = None
_client_error_message = None


def _get_workspace_client():
    """
    Get or create the Databricks WorkspaceClient.
    Uses U2M auth via databricks-cli profile.
    Handles token expiration gracefully.
    """
    global _workspace_client, _client_error_message
    
    try:
        from databricks.sdk import WorkspaceClient
        from databricks.sdk.core import DatabricksError
    except ImportError:
        _client_error_message = (
            "databricks-sdk is not installed. "
            "Please install it with: pip install databricks-sdk"
        )
        raise RuntimeError(_client_error_message)
    
    # Get profile from environment or use default
    profile = os.getenv("DATABRICKS_PROFILE", "DEFAULT")
    
    try:
        # Create a new client each time to handle token refresh
        # The SDK handles token refresh automatically for U2M auth
        client = WorkspaceClient(profile=profile)
        
        # Test the connection by making a simple API call
        # This will trigger OAuth token refresh if needed
        client.current_user.me()
        
        _workspace_client = client
        _client_error_message = None
        return client
        
    except DatabricksError as e:
        error_str = str(e)
        if "token" in error_str.lower() or "expired" in error_str.lower() or "auth" in error_str.lower():
            _client_error_message = (
                f"Databricks authentication failed. Your token may have expired.\n"
                f"Please re-authenticate using: databricks auth login --profile {profile}\n"
                f"Original error: {error_str}"
            )
        else:
            _client_error_message = f"Databricks API error: {error_str}"
        raise RuntimeError(_client_error_message)
        
    except Exception as e:
        _client_error_message = (
            f"Failed to connect to Databricks.\n"
            f"Make sure you have run: databricks auth login --profile {profile}\n"
            f"Error: {str(e)}"
        )
        raise RuntimeError(_client_error_message)


def _handle_auth_error(func):
    """Decorator to handle authentication errors gracefully."""
    @wraps(func)
    async def wrapper(*args, **kwargs):
        try:
            return await func(*args, **kwargs)
        except RuntimeError as e:
            return str(e)
        except Exception as e:
            error_str = str(e)
            if "token" in error_str.lower() or "expired" in error_str.lower() or "auth" in error_str.lower():
                profile = os.getenv("DATABRICKS_PROFILE", "DEFAULT")
                return (
                    f"❌ Authentication failed. Your token may have expired.\n"
                    f"Please re-authenticate using:\n"
                    f"  databricks auth login --profile {profile}\n\n"
                    f"Original error: {error_str}"
                )
            return f"❌ Error: {error_str}"
    return wrapper


@mcp.tool()
@_handle_auth_error
async def who_am_i() -> str:
    """
    Get the current authenticated Databricks user.
    
    Returns the username of the currently authenticated user.
    This is useful to verify that authentication is working correctly.
    """
    client = _get_workspace_client()
    current_user = client.current_user.me()
    
    result = {
        "username": current_user.user_name,
        "display_name": current_user.display_name,
        "active": current_user.active,
        "id": current_user.id
    }
    
    if current_user.emails:
        result["emails"] = [email.value for email in current_user.emails]
    
    return json.dumps(result, indent=2)


@mcp.tool()
@_handle_auth_error
async def list_catalogs() -> str:
    """
    List all available catalogs in Unity Catalog.
    
    Returns a list of catalog names that the current user has access to.
    """
    client = _get_workspace_client()
    
    catalogs = list(client.catalogs.list())
    
    catalog_list = []
    for catalog in catalogs:
        catalog_info = {
            "name": catalog.name,
            "comment": catalog.comment if catalog.comment else None,
            "owner": catalog.owner if catalog.owner else None,
            "created_at": str(catalog.created_at) if catalog.created_at else None
        }
        catalog_list.append(catalog_info)
    
    return json.dumps({
        "total_count": len(catalog_list),
        "catalogs": catalog_list
    }, indent=2)


@mcp.tool()
@_handle_auth_error
async def list_schemas(catalog_name: str) -> str:
    """
    List all schemas in a specified catalog.
    
    Args:
        catalog_name: The name of the catalog to list schemas from
    
    Returns a list of schema names within the specified catalog.
    """
    client = _get_workspace_client()
    
    schemas = list(client.schemas.list(catalog_name=catalog_name))
    
    schema_list = []
    for schema in schemas:
        schema_info = {
            "name": schema.name,
            "full_name": schema.full_name,
            "comment": schema.comment if schema.comment else None,
            "owner": schema.owner if schema.owner else None,
            "created_at": str(schema.created_at) if schema.created_at else None
        }
        schema_list.append(schema_info)
    
    return json.dumps({
        "catalog": catalog_name,
        "total_count": len(schema_list),
        "schemas": schema_list
    }, indent=2)


@mcp.tool()
@_handle_auth_error
async def list_tables(catalog_name: str, schema_name: str) -> str:
    """
    List all tables in a specified catalog and schema.
    
    Args:
        catalog_name: The name of the catalog
        schema_name: The name of the schema within the catalog
    
    Returns a list of tables within the specified catalog.schema.
    """
    client = _get_workspace_client()
    
    tables = list(client.tables.list(
        catalog_name=catalog_name,
        schema_name=schema_name
    ))
    
    table_list = []
    for table in tables:
        table_info = {
            "name": table.name,
            "full_name": table.full_name,
            "table_type": str(table.table_type) if table.table_type else None,
            "data_source_format": str(table.data_source_format) if table.data_source_format else None,
            "comment": table.comment if table.comment else None,
            "owner": table.owner if table.owner else None,
            "created_at": str(table.created_at) if table.created_at else None
        }
        table_list.append(table_info)
    
    return json.dumps({
        "catalog": catalog_name,
        "schema": schema_name,
        "total_count": len(table_list),
        "tables": table_list
    }, indent=2)


@mcp.tool()
@_handle_auth_error
async def get_table_details(catalog_name: str, schema_name: str, table_name: str) -> str:
    """
    Get detailed information about a table including column names, data types, and metadata.
    
    Args:
        catalog_name: The name of the catalog
        schema_name: The name of the schema
        table_name: The name of the table
    
    Returns detailed metadata about the table including all columns with their data types.
    """
    client = _get_workspace_client()
    
    full_name = f"{catalog_name}.{schema_name}.{table_name}"
    table_info = client.tables.get(full_name=full_name)
    
    columns = []
    if table_info.columns:
        for i, col in enumerate(table_info.columns, 1):
            col_info = {
                "position": i,
                "name": col.name,
                "type": str(col.type_name) if col.type_name else col.type_text,
                "nullable": col.nullable if col.nullable is not None else True,
                "comment": col.comment if col.comment else None
            }
            if col.type_precision:
                col_info["precision"] = col.type_precision
            if col.type_scale:
                col_info["scale"] = col.type_scale
            columns.append(col_info)
    
    result = {
        "full_name": full_name,
        "table_type": str(table_info.table_type) if table_info.table_type else None,
        "data_source_format": str(table_info.data_source_format) if table_info.data_source_format else None,
        "storage_location": table_info.storage_location if table_info.storage_location else None,
        "owner": table_info.owner if table_info.owner else None,
        "comment": table_info.comment if table_info.comment else None,
        "created_at": str(table_info.created_at) if table_info.created_at else None,
        "updated_at": str(table_info.updated_at) if table_info.updated_at else None,
        "column_count": len(columns),
        "columns": columns
    }
    
    # Add properties if available
    if table_info.properties:
        result["properties"] = dict(table_info.properties)
    
    return json.dumps(result, indent=2)


@mcp.tool()
@_handle_auth_error
async def get_sample_data(
    catalog_name: str, 
    schema_name: str, 
    table_name: str,
    limit: int = 10
) -> str:
    """
    Get sample data from a table.
    
    Args:
        catalog_name: The name of the catalog
        schema_name: The name of the schema
        table_name: The name of the table
        limit: Number of sample rows to return (min: 5, max: 50, default: 10)
    
    Returns sample rows from the table in a structured format.
    """
    client = _get_workspace_client()
    
    # Enforce limit bounds
    limit = max(5, min(50, limit))
    
    full_name = f"{catalog_name}.{schema_name}.{table_name}"
    
    # First, find available SQL warehouses
    try:
        warehouses = list(client.warehouses.list())
    except Exception as e:
        return json.dumps({
            "error": "Failed to list SQL warehouses",
            "message": str(e),
            "suggestion": "Ensure you have access to at least one SQL warehouse"
        }, indent=2)
    
    if not warehouses:
        return json.dumps({
            "error": "No SQL warehouses available",
            "message": "Cannot execute sample query without a SQL warehouse",
            "suggestion": "Please create or request access to a SQL warehouse in your Databricks workspace"
        }, indent=2)
    
    # Find a running warehouse, or use the first available
    running_warehouses = [wh for wh in warehouses if str(wh.state) == "State.RUNNING"]
    warehouse = running_warehouses[0] if running_warehouses else warehouses[0]
    
    logger.info(f"Using warehouse: {warehouse.name} ({warehouse.id})")
    
    # Execute the sample query
    try:
        result = client.statement_execution.execute_statement(
            warehouse_id=warehouse.id,
            statement=f"SELECT * FROM {full_name} LIMIT {limit}",
            wait_timeout="50s"
        )
        
        if result.status and result.status.state:
            state = str(result.status.state)
            if "FAILED" in state:
                error_msg = result.status.error.message if result.status.error else "Unknown error"
                return json.dumps({
                    "error": "Query execution failed",
                    "state": state,
                    "message": error_msg,
                    "table": full_name
                }, indent=2)
        
        if result.result and result.result.data_array:
            # Get column names from manifest
            col_names = [col.name for col in result.manifest.schema.columns]
            
            # Convert data to list of dictionaries
            rows = []
            for row_data in result.result.data_array:
                row_dict = {}
                for i, val in enumerate(row_data):
                    col_name = col_names[i] if i < len(col_names) else f"col_{i}"
                    row_dict[col_name] = val
                rows.append(row_dict)
            
            return json.dumps({
                "table": full_name,
                "warehouse_used": warehouse.name,
                "row_count": len(rows),
                "columns": col_names,
                "data": rows
            }, indent=2)
        else:
            return json.dumps({
                "table": full_name,
                "warehouse_used": warehouse.name,
                "row_count": 0,
                "message": "No data returned - table may be empty",
                "columns": [],
                "data": []
            }, indent=2)
            
    except Exception as e:
        error_str = str(e)
        
        # Handle common errors gracefully
        if "WAREHOUSE_STARTING" in error_str or "starting" in error_str.lower():
            return json.dumps({
                "error": "SQL warehouse is starting",
                "message": "The SQL warehouse is currently starting up. Please try again in a few moments.",
                "warehouse": warehouse.name,
                "table": full_name
            }, indent=2)
        
        if "permission" in error_str.lower() or "access" in error_str.lower():
            return json.dumps({
                "error": "Permission denied",
                "message": f"You don't have permission to query table: {full_name}",
                "details": error_str
            }, indent=2)
        
        if "does not exist" in error_str.lower() or "not found" in error_str.lower():
            return json.dumps({
                "error": "Table not found",
                "message": f"Table does not exist: {full_name}",
                "suggestion": "Check the catalog, schema, and table names"
            }, indent=2)
        
        return json.dumps({
            "error": "Query execution error",
            "message": error_str,
            "table": full_name,
            "warehouse": warehouse.name
        }, indent=2)


async def main():
    """Main entry point for the MCP server."""
    logger.info("Starting Databricks MCP Server")
    
    # Log configuration
    profile = os.getenv("DATABRICKS_PROFILE", "DEFAULT")
    logger.info(f"Using Databricks profile: {profile}")
    
    # Run the MCP server
    await mcp.run_async()


if __name__ == "__main__":
    import asyncio
    asyncio.run(main())
