from __future__ import annotations

from collections import defaultdict
import logging

import sqlglot
from sqlglot import exp

from auraindex.models import ParseFailure, ParsedQueryFeatures, QueryCandidate, StatementKind
from auraindex.utils.text import dedupe_preserve_order, sanitize_identifier

logging.getLogger("sqlglot").setLevel(logging.ERROR)


def analyze_sql_candidates(
    candidates: list[QueryCandidate],
    dialect: str = "postgres",
) -> tuple[list[ParsedQueryFeatures], list[ParseFailure]]:
    parsed: list[ParsedQueryFeatures] = []
    failures: list[ParseFailure] = []

    for candidate in candidates:
        try:
            tree, parsed_dialect = _parse_with_dialect_fallback(candidate.sql, preferred_dialect=dialect)
            if not isinstance(candidate.metadata, dict):
                candidate.metadata = {"raw_metadata": candidate.metadata}
            candidate.metadata["parsed_dialect"] = parsed_dialect
        except Exception as exc:  # parse exceptions are expected for malformed or pseudo SQL
            failures.append(ParseFailure(query=candidate, error=str(exc)))
            continue

        alias_map = _build_alias_map(tree)
        statement_kind = _detect_statement_kind(tree)
        tables = _extract_tables(tree)
        default_table = tables[0] if tables else "_unknown"

        where_columns = _extract_columns(tree.args.get("where"), alias_map, default_table)
        join_columns: list[tuple[str, str]] = []
        for join in tree.find_all(exp.Join):
            join_columns.extend(_extract_columns(join.args.get("on"), alias_map, default_table))
        join_columns = _dedupe_pairs(join_columns)

        order_columns = _extract_order_columns(tree.args.get("order"), alias_map, default_table)
        group_columns = _extract_group_columns(tree.args.get("group"), alias_map, default_table)
        equality_by_table = _extract_equality_columns(tree.args.get("where"), alias_map, default_table)

        parsed.append(
            ParsedQueryFeatures(
                query=candidate,
                statement_kind=statement_kind,
                tables=tables,
                where_columns=where_columns,
                join_columns=join_columns,
                order_by_columns=order_columns,
                group_by_columns=group_columns,
                equality_columns_by_table=equality_by_table,
            )
        )

    return parsed, failures


def _parse_with_dialect_fallback(sql: str, preferred_dialect: str) -> tuple[exp.Expression, str]:
    preferred = (preferred_dialect or "postgres").strip().lower()
    order = [preferred, "mysql", "postgres", "sqlite"]
    tried: set[str] = set()
    last_error: Exception | None = None

    for candidate_dialect in order:
        if candidate_dialect in tried:
            continue
        tried.add(candidate_dialect)
        try:
            return sqlglot.parse_one(sql, read=candidate_dialect), candidate_dialect
        except Exception as exc:
            last_error = exc

    if last_error is None:
        raise ValueError("SQL parse failed with unknown parser error.")
    raise last_error


def _build_alias_map(tree: exp.Expression) -> dict[str, str]:
    alias_map: dict[str, str] = {}
    for table in tree.find_all(exp.Table):
        table_name = sanitize_identifier(table.name or "")
        alias_name = sanitize_identifier(table.alias or "")
        if table_name and alias_name and alias_name != table_name:
            alias_map[alias_name] = table_name
    return alias_map


def _extract_tables(tree: exp.Expression) -> list[str]:
    tables: list[str] = []
    for table in tree.find_all(exp.Table):
        name = sanitize_identifier(table.name or "")
        if name:
            tables.append(name)
    return dedupe_preserve_order(tables)


def _extract_columns(
    expression: exp.Expression | None,
    alias_map: dict[str, str],
    default_table: str,
) -> list[tuple[str, str]]:
    if expression is None:
        return []
    pairs: list[tuple[str, str]] = []
    for column in expression.find_all(exp.Column):
        pairs.append(_column_to_pair(column, alias_map, default_table))
    return _dedupe_pairs(pairs)


def _extract_order_columns(
    order_expression: exp.Expression | None,
    alias_map: dict[str, str],
    default_table: str,
) -> list[tuple[str, str]]:
    if order_expression is None:
        return []
    pairs: list[tuple[str, str]] = []
    for ordered in order_expression.find_all(exp.Ordered):
        for column in ordered.find_all(exp.Column):
            pairs.append(_column_to_pair(column, alias_map, default_table))
    return _dedupe_pairs(pairs)


def _extract_group_columns(
    group_expression: exp.Expression | None,
    alias_map: dict[str, str],
    default_table: str,
) -> list[tuple[str, str]]:
    if group_expression is None:
        return []
    pairs: list[tuple[str, str]] = []
    for column in group_expression.find_all(exp.Column):
        pairs.append(_column_to_pair(column, alias_map, default_table))
    return _dedupe_pairs(pairs)


def _extract_equality_columns(
    where_expression: exp.Expression | None,
    alias_map: dict[str, str],
    default_table: str,
) -> dict[str, list[str]]:
    if where_expression is None:
        return {}

    grouped: dict[str, list[str]] = defaultdict(list)
    for eq_expr in where_expression.find_all(exp.EQ):
        for column in eq_expr.find_all(exp.Column):
            table, col = _column_to_pair(column, alias_map, default_table)
            grouped[table].append(col)

    return {table: dedupe_preserve_order(cols) for table, cols in grouped.items() if cols}


def _column_to_pair(
    column: exp.Column,
    alias_map: dict[str, str],
    default_table: str,
) -> tuple[str, str]:
    raw_table = sanitize_identifier(column.table or "")
    raw_column = sanitize_identifier(column.name or "")
    table = alias_map.get(raw_table, raw_table) if raw_table else default_table
    return table or default_table, raw_column or "_unknown_column"


def _dedupe_pairs(pairs: list[tuple[str, str]]) -> list[tuple[str, str]]:
    seen: set[tuple[str, str]] = set()
    ordered: list[tuple[str, str]] = []
    for pair in pairs:
        if pair in seen:
            continue
        seen.add(pair)
        ordered.append(pair)
    return ordered


def _detect_statement_kind(tree: exp.Expression) -> StatementKind:
    if isinstance(tree, exp.Select):
        return "select"
    if isinstance(tree, exp.Insert):
        return "insert"
    if isinstance(tree, exp.Update):
        return "update"
    if isinstance(tree, exp.Delete):
        return "delete"
    if isinstance(tree, exp.Merge):
        return "merge"
    if isinstance(tree, exp.Replace):
        return "replace"
    if isinstance(tree, exp.Create):
        return "ddl"
    return "other"
