from __future__ import annotations

import re


SQL_TYPE_ALIASES: dict[str, str] = {
    "TEXT": "VARCHAR",
    "STRING": "VARCHAR",
    "BOOL": "BOOLEAN",
    "FLOAT": "DOUBLE",
    "FLOAT8": "DOUBLE",
}


FUNCTION_RETURN_SQL_TYPE_BY_NAME: dict[str, str] = {
    # Text-producing functions
    "MD5": "VARCHAR",
    "CONCAT": "VARCHAR",
    "TRIM": "VARCHAR",
    "UUID_TO_STRING": "VARCHAR",
    # Floating-point math functions
    "POWER": "DOUBLE",
    "LN": "DOUBLE",
    "LOG": "DOUBLE",
    "LOG10": "DOUBLE",
    "SQRT": "DOUBLE",
    "EXP": "DOUBLE",
    "CBRT": "DOUBLE",
    "COS": "DOUBLE",
    "SIN": "DOUBLE",
    "TAN": "DOUBLE",
    "ACOS": "DOUBLE",
    "ASIN": "DOUBLE",
    "ATAN": "DOUBLE",
    "COSH": "DOUBLE",
    "SINH": "DOUBLE",
    "TANH": "DOUBLE",
    "ACOSH": "DOUBLE",
    "ASINH": "DOUBLE",
    "ATANH": "DOUBLE",
    "RADIANS": "DOUBLE",
    "DEGREES": "DOUBLE",
    "COT": "DOUBLE",
    # Boolean predicates
    "ISINF": "BOOLEAN",
    "ISNAN": "BOOLEAN",
    "LIKE": "BOOLEAN",
    "REGEXP_MATCHES": "BOOLEAN",
    # Window ranking functions
    "ROW_NUMBER": "INTEGER",
    "RANK": "INTEGER",
    "DENSE_RANK": "INTEGER",
}


def normalize_sql_type(sql_type: str) -> str:
    normalized = re.sub(r"\s+", "", sql_type).upper()
    return SQL_TYPE_ALIASES.get(normalized, normalized)


def same_sql_type(left: str, right: str) -> bool:
    return normalize_sql_type(left) == normalize_sql_type(right)


def is_text_sql_type(sql_type: str) -> bool:
    return normalize_sql_type(sql_type) == "VARCHAR"


def infer_function_return_sql_type(name: str, *, arg_types: tuple[str | None, ...] = ()) -> str | None:
    normalized_name = name.strip().upper()
    mapped = FUNCTION_RETURN_SQL_TYPE_BY_NAME.get(normalized_name)
    if mapped is not None:
        return mapped

    # COALESCE returns a common supertype. We only infer when all known args agree.
    if normalized_name == "COALESCE":
        known_arg_types = [normalize_sql_type(t) for t in arg_types if t is not None]
        if known_arg_types and all(t == known_arg_types[0] for t in known_arg_types):
            return known_arg_types[0]

    # MIN/MAX preserve input type for a single known argument.
    if normalized_name in {"MIN", "MAX"} and arg_types:
        first = arg_types[0]
        if first is not None:
            return normalize_sql_type(first)

    return None
