"""Database query and schema inspection tools."""

from typing import Annotated

from ..config import is_readonly
from ..utils import format_result, get_connection, rows_to_dicts, validate_query


def query(
    sql: Annotated[str, "The SQL query to execute (SELECT statements)"]
) -> str:
    """Execute a SQL query and return results. Use for SELECT statements."""
    is_valid, error_msg = validate_query(sql)
    if not is_valid:
        return format_result({"error": error_msg})

    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(sql)
        if cursor.description:
            rows = rows_to_dicts(cursor)
            return format_result(rows) if rows else "Query executed successfully. No rows returned."
        return "Query executed successfully. No rows returned."
    finally:
        conn.close()


def execute(
    sql: Annotated[str, "The SQL statement to execute (INSERT, UPDATE, DELETE, CREATE, etc.)"]
) -> str:
    """Execute a SQL statement (INSERT, UPDATE, DELETE, CREATE, etc.) and return affected rows count."""
    # Use validate_query for granular permission checking
    is_valid, error_msg = validate_query(sql)
    if not is_valid:
        return format_result({"error": error_msg})

    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(sql)
        conn.commit()
        return f"Statement executed successfully. Rows affected: {cursor.rowcount}"
    finally:
        conn.close()


def list_tables(
    schema: Annotated[str | None, "Schema name to filter tables (optional, defaults to all schemas)"] = None
) -> str:
    """List all tables in the current database."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        if schema:
            sql = """
                SELECT TABLE_SCHEMA as [schema], TABLE_NAME as [table], TABLE_TYPE as [type]
                FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ?
                ORDER BY TABLE_SCHEMA, TABLE_NAME
            """
            cursor.execute(sql, (schema,))
        else:
            sql = """
                SELECT TABLE_SCHEMA as [schema], TABLE_NAME as [table], TABLE_TYPE as [type]
                FROM INFORMATION_SCHEMA.TABLES ORDER BY TABLE_SCHEMA, TABLE_NAME
            """
            cursor.execute(sql)
        return format_result(rows_to_dicts(cursor))
    finally:
        conn.close()


def describe_table(
    table: Annotated[str, "The table name to describe"],
    schema: Annotated[str, "The schema name (defaults to 'dbo')"] = "dbo"
) -> str:
    """Get the schema/structure of a specific table including columns, data types, and constraints."""
    sql = """
        SELECT
            c.COLUMN_NAME as [column], c.DATA_TYPE as [type],
            c.CHARACTER_MAXIMUM_LENGTH as [max_length],
            c.NUMERIC_PRECISION as [precision], c.NUMERIC_SCALE as [scale],
            c.IS_NULLABLE as [nullable], c.COLUMN_DEFAULT as [default],
            CASE WHEN pk.COLUMN_NAME IS NOT NULL THEN 'YES' ELSE 'NO' END as [primary_key]
        FROM INFORMATION_SCHEMA.COLUMNS c
        LEFT JOIN (
            SELECT ku.TABLE_SCHEMA, ku.TABLE_NAME, ku.COLUMN_NAME
            FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS tc
            JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE ku ON tc.CONSTRAINT_NAME = ku.CONSTRAINT_NAME
            WHERE tc.CONSTRAINT_TYPE = 'PRIMARY KEY'
        ) pk ON c.TABLE_SCHEMA = pk.TABLE_SCHEMA AND c.TABLE_NAME = pk.TABLE_NAME AND c.COLUMN_NAME = pk.COLUMN_NAME
        WHERE c.TABLE_NAME = ? AND c.TABLE_SCHEMA = ?
        ORDER BY c.ORDINAL_POSITION
    """
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(sql, (table, schema))
        return format_result(rows_to_dicts(cursor))
    finally:
        conn.close()


def list_databases() -> str:
    """List all databases on the SQL Server instance."""
    sql = """
        SELECT name as [database], state_desc as [state],
               recovery_model_desc as [recovery_model], create_date as [created]
        FROM sys.databases ORDER BY name
    """
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(sql)
        return format_result(rows_to_dicts(cursor))
    finally:
        conn.close()


def get_table_sample(
    table: Annotated[str, "The table name"],
    schema: Annotated[str, "The schema name (defaults to 'dbo')"] = "dbo",
    limit: Annotated[int, "Number of rows to return (defaults to 10, max 1000)"] = 10
) -> str:
    """Get a sample of rows from a table."""
    safe_schema = "".join(c for c in schema if c.isalnum() or c == "_")
    safe_table = "".join(c for c in table if c.isalnum() or c == "_")
    safe_limit = min(max(1, int(limit)), 1000)

    sql = f"SELECT TOP {safe_limit} * FROM [{safe_schema}].[{safe_table}]"
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(sql)
        return format_result(rows_to_dicts(cursor))
    finally:
        conn.close()


def get_table_indexes(
    table: Annotated[str, "The table name"],
    schema: Annotated[str, "The schema name (defaults to 'dbo')"] = "dbo"
) -> str:
    """Get indexes defined on a table."""
    sql = """
        SELECT i.name as [index_name], i.type_desc as [index_type],
               i.is_unique as [is_unique], i.is_primary_key as [is_primary_key],
               STRING_AGG(c.name, ', ') WITHIN GROUP (ORDER BY ic.key_ordinal) as [columns]
        FROM sys.indexes i
        JOIN sys.index_columns ic ON i.object_id = ic.object_id AND i.index_id = ic.index_id
        JOIN sys.columns c ON ic.object_id = c.object_id AND ic.column_id = c.column_id
        JOIN sys.tables t ON i.object_id = t.object_id
        JOIN sys.schemas s ON t.schema_id = s.schema_id
        WHERE t.name = ? AND s.name = ?
        GROUP BY i.name, i.type_desc, i.is_unique, i.is_primary_key
        ORDER BY i.name
    """
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(sql, (table, schema))
        return format_result(rows_to_dicts(cursor))
    finally:
        conn.close()


def get_foreign_keys(
    table: Annotated[str, "The table name"],
    schema: Annotated[str, "The schema name (defaults to 'dbo')"] = "dbo"
) -> str:
    """Get foreign key relationships for a table."""
    sql = """
        SELECT fk.name as [constraint_name],
               OBJECT_SCHEMA_NAME(fk.parent_object_id) as [from_schema],
               OBJECT_NAME(fk.parent_object_id) as [from_table],
               COL_NAME(fkc.parent_object_id, fkc.parent_column_id) as [from_column],
               OBJECT_SCHEMA_NAME(fk.referenced_object_id) as [to_schema],
               OBJECT_NAME(fk.referenced_object_id) as [to_table],
               COL_NAME(fkc.referenced_object_id, fkc.referenced_column_id) as [to_column]
        FROM sys.foreign_keys fk
        JOIN sys.foreign_key_columns fkc ON fk.object_id = fkc.constraint_object_id
        JOIN sys.tables t ON fk.parent_object_id = t.object_id
        JOIN sys.schemas s ON t.schema_id = s.schema_id
        WHERE t.name = ? AND s.name = ?
        ORDER BY fk.name
    """
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(sql, (table, schema))
        return format_result(rows_to_dicts(cursor))
    finally:
        conn.close()
