import os
import datetime
from typing import AsyncIterator
from dataclasses import dataclass
from contextlib import asynccontextmanager
from google.cloud import bigquery
from mcp.server.fastmcp import FastMCP, Context


PROJECT_ID = os.getenv("PROJECT_ID")
DATASET = os.getenv("DATASET")


@dataclass
class AppContext:
    client: bigquery.Client


@asynccontextmanager
async def app_lifespan(server: FastMCP) -> AsyncIterator[AppContext]:
    """Manage application lifecycle with type-safe context"""
    client = bigquery.Client()
    yield AppContext(client=client)


mcp = FastMCP("Choose MCP Server", lifespan=app_lifespan)


@mcp.tool()
def get_tables(ctx: Context) -> str:
    """
    Get the list of tables in the database.

    Returns:
        str: The list of tables in the database.
    """
    client = ctx.request_context.lifespan_context.client
    query = f"SELECT table_name FROM `{PROJECT_ID}.{DATASET}.INFORMATION_SCHEMA.TABLES`"
    query_job = client.query(query)
    results = query_job.result()
    
    tables = [row.table_name for row in results]
    return "\n".join(tables)


@mcp.tool()
def get_table_schema(ctx: Context, table: str) -> str:
    """
    Get the schema of a table in the database, along with the description of each available field.
    
    Args:
        table (str): The name of the table in the database.

    Returns:
        str: The schema of the table in the database.
    """
    client = ctx.request_context.lifespan_context.client
    query = f"SELECT ddl FROM `{PROJECT_ID}.{DATASET}.INFORMATION_SCHEMA.TABLES` WHERE table_name = @table"
    job_config = bigquery.QueryJobConfig(
        query_parameters=[
            bigquery.ScalarQueryParameter("table", "STRING", table)
        ]
    )
    query_job = client.query(query, job_config=job_config)
    results = query_job.result()
    
    for row in results:
        return row.ddl
    return f"Table {table} not found"


@mcp.tool()
def query(ctx: Context, query: str) -> str:
    """
    Execute a SQL query against BigQuery and get back the result as dict.
    Input to this tool is a detailed and correct SQL query, output is a result from the database.
    If the query is not correct, an error message will be returned.
    If an error is returned, rewrite the query, check the query, and try again.
    If you encounter an issue with Unknown column 'xxxx' in 'field list', use get_table_schema to query the correct table fields.
    """
    client = ctx.request_context.lifespan_context.client
    try:
        query_job = client.query(query)
        results = [dict(row.items()) for row in query_job.result()]
        return str(results)
    except Exception as e:
        return f"Error: Query failed. {str(e)}. Please rewrite the query, check the query, and try again."


@mcp.tool()
def fetch_current_time() -> str:
    """
    Fetch the current time in UTC.

    Returns:
        str: The current time in UTC.
    """
    return datetime.datetime.now().isoformat()


if __name__ == "__main__":
    mcp.run()