"""MCP tools for exploring Snowflake database schemas."""

import logging
from typing import Annotated, Any, Literal

import httpx
from pydantic import Field
from snowflake.sqlalchemy.snowdialect import SnowflakeDialect

from komodo.auth.Session import Session
from komodo.config import SDKSettings
from komodo.snowflake import get_snowflake_connection


logger = logging.getLogger(__name__)

# Lazy connection management - uses Any type to avoid import path issues
_connection: Any = None


def _quote_identifier(name: str) -> str:
    """Safely quote a Snowflake identifier using SQLAlchemy's IdentifierPreparer."""
    preparer = SnowflakeDialect().identifier_preparer
    return preparer.quote_identifier(name)


def _serialize_value(value: Any) -> Any:
    """Convert a value to a JSON-serializable type.

    Handles numpy types, datetime objects, and other non-serializable types
    by converting them to native Python types.
    """
    if value is None:
        return None
    if isinstance(value, str | bool):
        return value
    # Handle numpy integer types (numpy.int64, numpy.int32, etc.)
    if hasattr(value, "item") and hasattr(value, "dtype"):
        return value.item()
    if isinstance(value, int | float):
        return value
    # Convert everything else to string
    return str(value)


def _get_connection() -> Any:
    """Get or create a Snowflake connection.

    Uses lazy initialization - connection is created on first tool call.
    Credentials are read from ~/.komodo/credentials (set via `komodo login`).
    """
    global _connection
    if _connection is None:
        try:
            _connection = get_snowflake_connection()
        except Exception as e:
            raise RuntimeError(f"Failed to connect to Snowflake. Have you run 'uv run komodo login' " f"and 'uv run komodo account set'? Error: {e}") from e
    return _connection


def _fetch_sfma_prompt() -> str | None:
    """Fetch the SFMA exploration prompt from the connector gateway API.

    Returns:
        The SFMA prompt string, or None if fetch fails.
    """
    try:
        settings = SDKSettings()
        session = Session(environment=settings.environment)
        credentials = session.load_credentials()

        if not credentials or not credentials.access_token:
            logger.warning("No access token found for fetching SFMA prompt")
            return None

        url = f"{settings.proxy_domain}/prompts/sfma"
        headers = {"Authorization": f"Bearer {credentials.access_token}", "accept": "text/plain"}

        with httpx.Client() as client:
            response = client.get(url, headers=headers, timeout=30.0)
            response.raise_for_status()
            return response.text
    except Exception as e:
        logger.warning(f"Failed to fetch SFMA prompt from API: {e}")
        return None


def _execute_query(query: str, params: dict[str, Any] | None = None) -> list[dict[str, Any]]:
    """Execute a query and return results as a list of dicts."""
    conn = _get_connection()
    cursor = conn.cursor()
    try:
        if params:
            cursor.execute(query, params)
        else:
            cursor.execute(query)

        columns = [desc[0] for desc in cursor.description] if cursor.description else []
        rows = cursor.fetchall()
        return [dict(zip(columns, row, strict=False)) for row in rows]
    finally:
        cursor.close()


def _list_databases() -> list[str]:
    """List all accessible databases in Snowflake."""
    results = _execute_query("SHOW DATABASES")
    return [row["name"] for row in results]


def _list_schemas(database: str) -> list[str]:
    """List all schemas within a database."""
    quoted_db = _quote_identifier(database)
    results = _execute_query(f"SHOW SCHEMAS IN DATABASE {quoted_db}")
    return [row["name"] for row in results]


def _list_tables_and_views(database: str, schema: str) -> list[dict[str, Any]]:
    """List all tables and views within a schema with metadata."""
    quoted_db = _quote_identifier(database)
    quoted_schema = _quote_identifier(schema)
    results = _execute_query(f"SHOW OBJECTS IN {quoted_db}.{quoted_schema}")
    return [
        {
            "name": row["name"],
            "rows": _serialize_value(row.get("rows")),
            "created_on": str(row.get("created_on", "")),
            "kind": row.get("kind", ""),
        }
        for row in results
    ]


def _list_columns(database: str, schema: str, table: str) -> list[dict[str, Any]]:
    """List all columns in a table with their types."""
    quoted_db = _quote_identifier(database)
    # S608: False positive - quoted_db is safely quoted via IdentifierPreparer,
    # schema/table use parameterized queries
    query = f"""
    SELECT 
        COLUMN_NAME,
        DATA_TYPE,
        IS_NULLABLE,
        COLUMN_DEFAULT,
        CHARACTER_MAXIMUM_LENGTH,
        NUMERIC_PRECISION,
        NUMERIC_SCALE
    FROM {quoted_db}.INFORMATION_SCHEMA.COLUMNS
    WHERE TABLE_SCHEMA = %(schema)s AND TABLE_NAME = %(table)s
    ORDER BY ORDINAL_POSITION
    """  # noqa: S608
    results = _execute_query(query, {"schema": schema, "table": table})
    return [
        {
            "column_name": row["COLUMN_NAME"],
            "data_type": row["DATA_TYPE"],
            "nullable": row["IS_NULLABLE"] == "YES",
            "default": _serialize_value(row.get("COLUMN_DEFAULT")),
            "max_length": _serialize_value(row.get("CHARACTER_MAXIMUM_LENGTH")),
            "precision": _serialize_value(row.get("NUMERIC_PRECISION")),
            "scale": _serialize_value(row.get("NUMERIC_SCALE")),
        }
        for row in results
    ]


def list_snowflake(
    list_type: Literal["database", "schema", "table", "column"],
    database: str | None = None,
    schema: str | None = None,
    table: str | None = None,
    prompt_loaded: Annotated[
        bool,
        Field(
            default=False,
            description="Whether the SFMA exploration prompt has already been loaded. "
            "On the first call (False), the prompt will be fetched and included in the response. "
            "Subsequent calls should pass True to skip fetching the prompt.",
        ),
    ] = False,
) -> dict[str, Any]:
    """List Snowflake objects based on the specified type.

    Args:
        list_type: Type of objects to list ("database", "schema", "table", "column")
        database: Required for schema/table/column listing
        schema: Required for table/column listing
        table: Required for column listing
        prompt_loaded: Whether the SFMA exploration prompt has already been loaded.
            On the first call (prompt_loaded=False), the prompt will be fetched from
            the API and included in the response. Subsequent calls should pass
            prompt_loaded=True to skip fetching the prompt.

    Returns:
        A dict containing:
        - 'data': A list of database/schema names (strings), or table/column info (dicts).
        - 'prompt' (optional): The SFMA exploration prompt, included only on first call.
    """
    # Get the data based on list_type
    match list_type:
        case "database":
            data = _list_databases()
        case "schema":
            if not database:
                raise ValueError("database is required when listing schemas")
            data = _list_schemas(database)
        case "table":
            if not database or not schema:
                raise ValueError("database and schema are required when listing tables")
            data = _list_tables_and_views(database, schema)
        case "column":
            if not database or not schema or not table:
                raise ValueError("database, schema, and table are required when listing columns")
            data = _list_columns(database, schema, table)
        case _:
            raise ValueError(f"Invalid list_type: {list_type}. Must be one of: database, schema, table, column")

    result: dict[str, Any] = {"data": data}

    # Include prompt if not already loaded
    if not prompt_loaded:
        prompt = _fetch_sfma_prompt()
        if prompt:
            result["prompt"] = prompt

    return result
