"""SQL parsing utilities for analyzing query types and commands."""

import logging

import sqlparse

from data_exchange_agent.data_sources.sql_command_type import SQLCommandType


def get_permitted_sql_command_type(query: str) -> SQLCommandType | None:
    """
    Classify a SQL statement and return its type if it is a permitted operation.

    Permitted operations include read-only queries (SELECT, WITH, DESCRIBE,
    SHOW, EXPLAIN) and data-export commands (UNLOAD).

    Args:
        query: The SQL query to parse.

    Returns:
        SQLCommandType | None: The command type if permitted, None otherwise.

    """
    try:
        statements = sqlparse.parse(query)

        if len(statements) != 1:
            logging.error("Multiple statements are not allowed.")
            return None

        statement = statements[0]
        first_token = statement.token_first()

        if first_token:
            # Get the command word - may be the full token value or just the first word
            token_value = first_token.value.upper()
            # Extract just the command word (first word before space or parenthesis)
            command = token_value.split()[0] if " " in token_value else token_value.split("(")[0]

            result_type = None

            # Special case: Check for UNLOAD first (may have ttype=None from sqlparse)
            if command == "UNLOAD":
                result_type = SQLCommandType.UNLOAD
            elif first_token.ttype:
                match first_token.ttype:
                    case sqlparse.tokens.DML:
                        match command:
                            case "SELECT":
                                result_type = SQLCommandType.SELECT
                            case _:
                                logging.error(f"The DML command '{command}' is not a permitted read-only operation.")
                    case sqlparse.tokens.CTE:
                        match command:
                            case "WITH":
                                result_type = SQLCommandType.WITH
                            case _:
                                logging.error(f"The CTE command '{command}' is not a permitted read-only operation.")
                    case sqlparse.tokens.DDL:
                        match command:
                            case "DESCRIBE":
                                result_type = SQLCommandType.DESCRIBE
                            case "DESC":
                                result_type = SQLCommandType.DESC
                            case "SHOW":
                                result_type = SQLCommandType.SHOW
                            case "EXPLAIN":
                                result_type = SQLCommandType.EXPLAIN
                            case _:
                                logging.error(f"The DDL command '{command}' is not a permitted read-only operation.")
                    case sqlparse.tokens.Keyword:
                        # Handle other keyword commands
                        logging.error(f"The keyword command '{command}' is not a permitted read-only operation.")
                    case _:
                        logging.error(f"The command '{command}' is not a permitted read-only operation.")
            else:
                # No ttype - only allow if we already identified it (like UNLOAD above)
                if result_type is None:
                    logging.error(f"The command '{command}' is not a permitted read-only operation.")

            return result_type
        else:
            logging.error("The query is empty or invalid.")
            return None

    except Exception as e:
        logging.error(f"Error parsing the query: {e}")
        return None
