"""
Nuuly BigQuery MCP Server

This server provides access to BigQuery data through the Model Context Protocol (MCP).

Available MCP endpoints:
- list_databases: List available databases (datasets)
- get_tables: Get list of tables in a dataset (with caching for better performance)
- get_schema: Get schema for a specific table in a database
- run_query: Run a SQL query against BigQuery

IMPORTANT: All endpoints require a project_id parameter. Valid project IDs are:
- rental-dev (default)
- rental-staging
- rental-prod
"""

import os
import httpx
import logging

# Configuration for BigQuery MCP Server
BIGQUERY_MCP_SERVER_URL = os.getenv("BIGQUERY_MCP_SERVER_URL", "https://www.example.com")
# API Key for authentication
BIGQUERY_MCP_API_KEY = os.getenv("BQ_API_KEY")

# Valid project IDs
VALID_PROJECTS = ["rental-dev", "rental-staging", "rental-prod"]
DEFAULT_PROJECT = "rental-dev"

logger = logging.getLogger(__name__)

from fastmcp import FastMCP, Context

mcp = FastMCP("Nuuly BigQuery MCP")

@mcp.tool()
async def get_tables(ctx: Context, database: str, project_id: str = DEFAULT_PROJECT) -> str:
    """Get list of tables in a database (dataset) with caching for better performance.

    Args:
        ctx: The MCP server provided context.
        database: The name of a database (BigQuery dataset)
        project_id: The GCP project ID to query (rental-dev, rental-staging, rental-prod)
    Returns:
        A JSON string representing the list of tables in the dataset.
    """
    try:
        # Validate project_id
        if project_id not in VALID_PROJECTS:
            logger.error(f"Invalid project_id: {project_id}. Must be one of {VALID_PROJECTS}")
            return {"error": f"Invalid project_id: {project_id}. Must be one of {VALID_PROJECTS}"}
            
        # Prepare request payload
        payload = {
            "function_name": "get_tables",
            "parameters": {
                "database": database,
                "project_id": project_id
            }
        }
        logger.info(f"Sending request to {BIGQUERY_MCP_SERVER_URL}")
        
        # Check if API key is available
        if not BIGQUERY_MCP_API_KEY:
            logger.error("BQ_API_KEY environment variable is not set")
            return {"error": "API key not configured. Set BQ_API_KEY environment variable."}
            
        # Include API key in headers
        headers = {"x-api-key": BIGQUERY_MCP_API_KEY}
        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(BIGQUERY_MCP_SERVER_URL, json=payload, headers=headers)
            response.raise_for_status()
            return response.json()
    except httpx.HTTPStatusError as e:
        logger.error(f"HTTPStatusError: status={e.response.status_code}, body={e.response.text}")
        return {"error": f"BigQuery MCP server error: {e.response.status_code}", "details": e.response.text}
    except Exception as e:
        logger.error(f"Exception reaching BigQuery MCP server: {e}")
        return {"error": "Failed to reach BigQuery MCP server", "details": str(e)}

@mcp.tool()
async def run_query(ctx: Context, database: str, sql: str, project_id: str = DEFAULT_PROJECT) -> str:
    """Run a SQL query against the database and return results.

    Args:
        ctx: The MCP server provided context.
        database: The name of a database (BigQuery dataset)
        sql: The SQL query string to execute. Accepts SELECT, SHOW, EXPLAIN queries only
        project_id: The GCP project ID to query (rental-dev, rental-staging, rental-prod)
    Returns:
        A JSON string representing the query result (columns and rows).
    """
    try:
        # Validate project_id
        if project_id not in VALID_PROJECTS:
            logger.error(f"Invalid project_id: {project_id}. Must be one of {VALID_PROJECTS}")
            return {"error": f"Invalid project_id: {project_id}. Must be one of {VALID_PROJECTS}"}
            
        # Prepare request payload
        payload = {
            "function_name": "run_query",
            "parameters": {
                "database": database,
                "sql": sql,
                "project_id": project_id
            }
        }
        logger.info(f"Sending request to {BIGQUERY_MCP_SERVER_URL}")
        
        # Check if API key is available
        if not BIGQUERY_MCP_API_KEY:
            logger.error("BQ_API_KEY environment variable is not set")
            return {"error": "API key not configured. Set BQ_API_KEY environment variable."}
            
        # Include API key in headers
        headers = {"x-api-key": BIGQUERY_MCP_API_KEY}
        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(BIGQUERY_MCP_SERVER_URL, json=payload, headers=headers)
            response.raise_for_status()
            return response.json()
    except httpx.HTTPStatusError as e:
        logger.error(f"HTTPStatusError: status={e.response.status_code}, body={e.response.text}")
        return {"error": f"BigQuery MCP server error: {e.response.status_code}", "details": e.response.text}
    except Exception as e:
        logger.error(f"Exception reaching BigQuery MCP server: {e}")
        return {"error": "Failed to reach BigQuery MCP server", "details": str(e)}

@mcp.tool()
async def get_schema(ctx: Context, database: str, table: str = None, project_id: str = DEFAULT_PROJECT) -> str:
    """Get schema for a database or specific table in BigQuery.

    Args:
        ctx: The MCP server provided context.
        database: The name of a database (BigQuery dataset) - this can 
            be an instance name or an alias. A full list of databases is available via the
            `list_databases` tool.
        table: Optional table name. If provided, only returns schema for this table.
        project_id: The GCP project ID to query (rental-dev, rental-staging, rental-prod)
    Returns:
        A JSON string describing the tables and their columns, foreign keys, primary keys,
            and constraints.
    """
    try:
        # Validate project_id
        if project_id not in VALID_PROJECTS:
            logger.error(f"Invalid project_id: {project_id}. Must be one of {VALID_PROJECTS}")
            return {"error": f"Invalid project_id: {project_id}. Must be one of {VALID_PROJECTS}"}
            
        payload = {
            "function_name": "get_schema",
            "parameters": {
                "database": database,
                "project_id": project_id
            }
        }
        
        # Add table parameter if provided
        if table:
            payload["parameters"]["table"] = table
            logger.info(f"Fetching schema for specific table: {database}.{table}")
        else:
            logger.info(f"Fetching schema for all tables in dataset: {database}")
            
        logger.info(f"Sending request to {BIGQUERY_MCP_SERVER_URL}")
        
        # Check if API key is available
        if not BIGQUERY_MCP_API_KEY:
            logger.error("BQ_API_KEY environment variable is not set")
            return {"error": "API key not configured. Set BQ_API_KEY environment variable."}
            
        # Include API key in headers
        headers = {"x-api-key": BIGQUERY_MCP_API_KEY}
        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(BIGQUERY_MCP_SERVER_URL, json=payload, headers=headers)
            response.raise_for_status()
            return response.json()
    except httpx.HTTPStatusError as e:
        logger.error(f"HTTPStatusError: status={e.response.status_code}, body={e.response.text}")
        return {"error": f"BigQuery MCP server error: {e.response.status_code}", "details": e.response.text}
    except Exception as e:
        logger.error(f"Exception reaching BigQuery MCP server: {e}")
        return {"error": "Failed to reach BigQuery MCP server", "details": str(e)}

@mcp.tool()
async def list_databases(ctx: Context, project_id: str = DEFAULT_PROJECT) -> str:
    """List all databases and their aliases accessible via BigQuery.

    Args:
        ctx: The MCP server provided context.
        project_id: The GCP project ID to query (rental-dev, rental-staging, rental-prod)
    Returns:
        A JSON string representing the list of databases, their aliases, and instance info.
    """
    try:
        # Validate project_id
        if project_id not in VALID_PROJECTS:
            logger.error(f"Invalid project_id: {project_id}. Must be one of {VALID_PROJECTS}")
            return {"error": f"Invalid project_id: {project_id}. Must be one of {VALID_PROJECTS}"}
            
        payload = {
            "function_name": "list_databases",
            "parameters": {
                "project_id": project_id
            }
        }
        logger.info(f"Sending request to {BIGQUERY_MCP_SERVER_URL}")
        
        # Check if API key is available
        if not BIGQUERY_MCP_API_KEY:
            logger.error("BQ_API_KEY environment variable is not set")
            return {"error": "API key not configured. Set BQ_API_KEY environment variable."}
            
        # Include API key in headers
        headers = {"x-api-key": BIGQUERY_MCP_API_KEY}
        async with httpx.AsyncClient(timeout=60) as client:
            response = await client.post(BIGQUERY_MCP_SERVER_URL, json=payload, headers=headers)
            response.raise_for_status()
            return response.json()
    except httpx.HTTPStatusError as e:
        logger.error(f"HTTPStatusError: status={e.response.status_code}, body={e.response.text}")
        return {"error": f"BigQuery MCP server error: {e.response.status_code}", "details": e.response.text}
    except Exception as e:
        logger.error(f"Exception reaching BigQuery MCP server: {e}")
        return {"error": "Failed to reach BigQuery MCP server", "details": str(e)}
