"""Database connection and utility functions for MSSQL MCP Server."""

import json
import re
import struct

import pyodbc
import sqlparse

from .config import DB_CONFIG, get_blocked_statements, is_readonly


def validate_query(sql: str) -> tuple[bool, str]:
    """Validate SQL query against permission restrictions."""
    blocked_statements = get_blocked_statements()

    # If nothing is blocked, allow all queries
    if not blocked_statements:
        return True, ""

    try:
        parsed = sqlparse.parse(sql)
        for statement in parsed:
            stmt_type = statement.get_type()
            if stmt_type and stmt_type.upper() in blocked_statements:
                if is_readonly():
                    return False, f"Blocked: {stmt_type} statements are not allowed in read-only mode."
                else:
                    return False, f"Blocked: {stmt_type} statements are disabled by configuration."

            tokens = [t for t in statement.tokens if not t.is_whitespace]
            if tokens:
                first_token = str(tokens[0]).upper().strip()
                for blocked in blocked_statements:
                    if first_token.startswith(blocked):
                        if is_readonly():
                            return False, f"Blocked: {blocked} statements are not allowed in read-only mode."
                        else:
                            return False, f"Blocked: {blocked} statements are disabled by configuration."
    except Exception as e:
        return False, f"Query validation failed: {str(e)}"

    return True, ""


def get_azure_token() -> bytes:
    """Get Azure AD token for database authentication."""
    try:
        from azure.identity import DefaultAzureCredential
    except ImportError:
        raise ImportError(
            "Azure AD authentication requires azure-identity. "
            "Install with: pip install mssql-agent-mcp[azure]"
        )

    credential = DefaultAzureCredential()
    token = credential.get_token("https://database.windows.net/.default")
    token_bytes = token.token.encode("utf-16-le")
    return struct.pack(f"<I{len(token_bytes)}s", len(token_bytes), token_bytes)


def get_connection():
    """Create and return a database connection."""
    auth_mode = DB_CONFIG["auth_mode"].lower()

    conn_str_parts = [
        f"DRIVER={{{DB_CONFIG['driver']}}}",
        f"SERVER={DB_CONFIG['server']},{DB_CONFIG['port']}",
        f"DATABASE={DB_CONFIG['database']}",
        f"Encrypt={DB_CONFIG['encrypt']}",
        f"TrustServerCertificate={DB_CONFIG['trust_server_certificate']}",
    ]

    if auth_mode == "windows":
        conn_str_parts.append("Trusted_Connection=yes")
        return pyodbc.connect(";".join(conn_str_parts))
    elif auth_mode == "azure":
        token_struct = get_azure_token()
        return pyodbc.connect(";".join(conn_str_parts), attrs_before={1256: token_struct})
    else:
        conn_str_parts.extend([
            f"UID={DB_CONFIG['user']}",
            f"PWD={DB_CONFIG['password']}",
        ])
        return pyodbc.connect(";".join(conn_str_parts))


def rows_to_dicts(cursor) -> list[dict]:
    """Convert pyodbc cursor results to list of dictionaries."""
    columns = [column[0] for column in cursor.description] if cursor.description else []
    return [dict(zip(columns, row)) for row in cursor.fetchall()]


def row_to_dict(cursor) -> dict | None:
    """Convert single pyodbc cursor result to dictionary."""
    if not cursor.description:
        return None
    columns = [column[0] for column in cursor.description]
    row = cursor.fetchone()
    return dict(zip(columns, row)) if row else None


def format_result(data) -> str:
    """Format query results as JSON string."""
    return json.dumps(data, indent=2, default=str)


def sanitize_filename(name: str) -> str:
    """Sanitize a string for use as a filename."""
    sanitized = re.sub(r'[<>:"/\\|?*]', "_", name)
    sanitized = re.sub(r"[\s_]+", "_", sanitized)
    return sanitized.strip("_ ")


def validate_sql_syntax(sql: str) -> dict:
    """Validate SQL syntax using SET PARSEONLY."""
    sql_stripped = sql.strip()
    if not sql_stripped:
        return {"valid": False, "error": "SQL command is empty"}

    try:
        conn = get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("SET PARSEONLY ON")
            try:
                cursor.execute(sql)
            except Exception as parse_error:
                return {"valid": False, "error": f"SQL syntax error: {str(parse_error)}"}
            finally:
                cursor.execute("SET PARSEONLY OFF")
        finally:
            conn.close()
        return {"valid": True, "error": None}
    except Exception as e:
        return {"valid": True, "error": None, "warning": f"Could not validate: {str(e)}"}
