from __future__ import annotations

from dataclasses import asdict
import sqlite3
from typing import Any

import sqlglot
from sqlglot import exp

from auraindex.analyzers.schema_checks import ExistingIndex
from auraindex.utils.text import dedupe_preserve_order, sanitize_identifier


def introspect_existing_indexes_from_live_db(
    *,
    dialect: str,
    host: str | None = None,
    port: int | None = None,
    user: str | None = None,
    password: str | None = None,
    database: str | None = None,
    db_path: str | None = None,
    schema: str | None = None,
    timeout_seconds: int = 8,
) -> tuple[list[ExistingIndex], list[str], dict[str, object]]:
    summary: dict[str, object] = {
        "enabled": True,
        "dialect": dialect,
        "connected": False,
        "host": host,
        "port": port,
        "database": database,
        "db_path": db_path,
        "schema": schema,
        "discovered_indexes": 0,
    }
    warnings: list[str] = []

    if dialect == "sqlite":
        indexes, warnings = _introspect_sqlite_indexes(db_path=db_path, timeout_seconds=timeout_seconds)
        summary["connected"] = bool(indexes or not warnings)
    elif dialect == "postgres":
        indexes, warnings = _introspect_postgres_indexes(
            host=host,
            port=port,
            user=user,
            password=password,
            database=database,
            schema=schema,
            timeout_seconds=timeout_seconds,
        )
        summary["connected"] = bool(indexes or not warnings)
    elif dialect == "mysql":
        indexes, warnings = _introspect_mysql_indexes(
            host=host,
            port=port,
            user=user,
            password=password,
            database=database,
            timeout_seconds=timeout_seconds,
        )
        summary["connected"] = bool(indexes or not warnings)
    else:
        indexes = []
        warnings.append(f"Unsupported db introspection dialect: {dialect}")

    deduped = _dedupe_existing_indexes(indexes)
    summary["discovered_indexes"] = len(deduped)
    summary["sample_indexes"] = [asdict(index) for index in deduped[:40]]
    return deduped, warnings, summary


def _introspect_sqlite_indexes(
    *,
    db_path: str | None,
    timeout_seconds: int,
) -> tuple[list[ExistingIndex], list[str]]:
    warnings: list[str] = []
    if not db_path:
        return [], ["SQLite db introspection requires --db-path."]

    indexes: list[ExistingIndex] = []
    try:
        connection = sqlite3.connect(db_path, timeout=max(1, timeout_seconds))
    except Exception as exc:
        return [], [f"Could not connect to sqlite database: {exc}"]

    try:
        table_rows = connection.execute(
            "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
        ).fetchall()
        tables = [str(row[0]) for row in table_rows if row and row[0]]

        for table in tables:
            safe_table = table.replace("'", "''")
            index_list = connection.execute(f"PRAGMA index_list('{safe_table}')").fetchall()
            for row in index_list:
                if len(row) < 3:
                    continue
                index_name = str(row[1])
                unique = bool(int(row[2]))

                safe_index = index_name.replace("'", "''")
                index_info = connection.execute(f"PRAGMA index_info('{safe_index}')").fetchall()
                columns = [sanitize_identifier(str(info_row[2])) for info_row in index_info if len(info_row) >= 3]
                columns = [col for col in columns if col]
                columns = dedupe_preserve_order(columns)
                if not columns:
                    continue

                sql_row = connection.execute(
                    "SELECT sql FROM sqlite_master WHERE type='index' AND name=?",
                    (index_name,),
                ).fetchone()
                source_sql = str(sql_row[0]) if sql_row and sql_row[0] else f"-- sqlite index: {index_name}"

                partial_where = None
                if source_sql and " where " in source_sql.lower():
                    parts = source_sql.split("WHERE", maxsplit=1)
                    if len(parts) == 2:
                        partial_where = parts[1].strip()

                indexes.append(
                    ExistingIndex(
                        name=sanitize_identifier(index_name) or index_name,
                        table=sanitize_identifier(table),
                        columns=columns,
                        unique=unique,
                        partial_where=partial_where,
                        source_file="live_db:sqlite",
                        source_line=0,
                        source_sql=source_sql[:500],
                    )
                )
    except Exception as exc:
        warnings.append(f"SQLite introspection query failed: {exc}")
    finally:
        connection.close()

    return indexes, warnings


def _introspect_postgres_indexes(
    *,
    host: str | None,
    port: int | None,
    user: str | None,
    password: str | None,
    database: str | None,
    schema: str | None,
    timeout_seconds: int,
) -> tuple[list[ExistingIndex], list[str]]:
    warnings: list[str] = []
    if not database:
        return [], ["Postgres db introspection requires --db-name."]

    connect_kwargs: dict[str, Any] = {
        "host": host or "localhost",
        "port": port or 5432,
        "user": user,
        "password": password,
        "dbname": database,
        "connect_timeout": max(1, timeout_seconds),
    }
    connect_kwargs = {key: value for key, value in connect_kwargs.items() if value is not None}

    connection = None
    driver_name = None
    try:
        import psycopg  # type: ignore

        connection = psycopg.connect(**connect_kwargs)
        driver_name = "psycopg"
    except Exception:
        try:
            import psycopg2  # type: ignore

            connection = psycopg2.connect(**connect_kwargs)
            driver_name = "psycopg2"
        except Exception as exc:
            warnings.append(
                "Could not connect to postgres database. Install `psycopg[binary]` or `psycopg2` "
                f"and verify credentials. Detail: {exc}"
            )
            return [], warnings

    indexes: list[ExistingIndex] = []
    try:
        if driver_name == "psycopg":
            with connection.cursor() as cursor:
                cursor.execute(
                    """
                    SELECT schemaname, tablename, indexname, indexdef
                    FROM pg_indexes
                    WHERE schemaname NOT IN ('pg_catalog', 'information_schema')
                      AND (%s IS NULL OR schemaname = %s)
                    """,
                    (schema, schema),
                )
                rows = cursor.fetchall()
        else:
            cursor = connection.cursor()
            cursor.execute(
                """
                SELECT schemaname, tablename, indexname, indexdef
                FROM pg_indexes
                WHERE schemaname NOT IN ('pg_catalog', 'information_schema')
                  AND (%s IS NULL OR schemaname = %s)
                """,
                (schema, schema),
            )
            rows = cursor.fetchall()
            cursor.close()

        for schemaname, tablename, indexname, indexdef in rows:
            parsed = _parse_index_sql(
                str(indexdef),
                dialect="postgres",
                fallback_table=str(tablename),
                fallback_index_name=str(indexname),
            )
            if parsed is None:
                continue
            parsed["source_file"] = f"live_db:postgres:{schemaname}"
            parsed["source_line"] = 0
            parsed["source_sql"] = str(indexdef)[:500]
            indexes.append(ExistingIndex(**parsed))
    except Exception as exc:
        warnings.append(f"Postgres introspection query failed: {exc}")
    finally:
        connection.close()

    return indexes, warnings


def _introspect_mysql_indexes(
    *,
    host: str | None,
    port: int | None,
    user: str | None,
    password: str | None,
    database: str | None,
    timeout_seconds: int,
) -> tuple[list[ExistingIndex], list[str]]:
    warnings: list[str] = []
    if not database:
        return [], ["MySQL db introspection requires --db-name."]

    try:
        import pymysql  # type: ignore
    except Exception:
        return [], ["MySQL db introspection requires `PyMySQL`. Install with `pip install auraindex[db]`."]

    indexes: list[ExistingIndex] = []
    try:
        connection = pymysql.connect(
            host=host or "localhost",
            port=int(port or 3306),
            user=user,
            password=password,
            database=database,
            connect_timeout=max(1, timeout_seconds),
            read_timeout=max(1, timeout_seconds),
            write_timeout=max(1, timeout_seconds),
            charset="utf8mb4",
        )
    except Exception as exc:
        return [], [f"Could not connect to mysql database: {exc}"]

    try:
        with connection.cursor() as cursor:
            cursor.execute(
                """
                SELECT TABLE_NAME, INDEX_NAME, NON_UNIQUE, SEQ_IN_INDEX, COLUMN_NAME
                FROM INFORMATION_SCHEMA.STATISTICS
                WHERE TABLE_SCHEMA = %s
                ORDER BY TABLE_NAME, INDEX_NAME, SEQ_IN_INDEX
                """,
                (database,),
            )
            rows = cursor.fetchall()

        grouped: dict[tuple[str, str], dict[str, Any]] = {}
        for table_name, index_name, non_unique, seq_in_index, column_name in rows:
            key = (str(table_name), str(index_name))
            entry = grouped.setdefault(
                key,
                {
                    "table": sanitize_identifier(str(table_name)),
                    "name": sanitize_identifier(str(index_name)) or str(index_name),
                    "unique": int(non_unique) == 0,
                    "columns": [],
                },
            )
            col = sanitize_identifier(str(column_name))
            if col:
                entry["columns"].append((int(seq_in_index), col))

        for (table_name, index_name), entry in grouped.items():
            ordered_columns = [col for _, col in sorted(entry["columns"], key=lambda item: item[0])]
            ordered_columns = dedupe_preserve_order(ordered_columns)
            if not ordered_columns:
                continue
            indexes.append(
                ExistingIndex(
                    name=entry["name"],
                    table=entry["table"],
                    columns=ordered_columns,
                    unique=bool(entry["unique"]),
                    partial_where=None,
                    source_file=f"live_db:mysql:{database}",
                    source_line=0,
                    source_sql=f"-- mysql index {index_name} on {table_name}",
                )
            )
    except Exception as exc:
        warnings.append(f"MySQL introspection query failed: {exc}")
    finally:
        connection.close()

    return indexes, warnings


def _parse_index_sql(
    sql_text: str,
    *,
    dialect: str,
    fallback_table: str,
    fallback_index_name: str,
) -> dict[str, Any] | None:
    try:
        tree = sqlglot.parse_one(sql_text, read=dialect)
    except Exception:
        tree = None

    if isinstance(tree, exp.Create):
        kind = str(tree.args.get("kind") or "").lower()
        if "index" in kind and isinstance(tree.this, exp.Index):
            index_expr = tree.this
            table = _extract_index_table(index_expr) or sanitize_identifier(fallback_table)
            columns = _extract_index_columns(index_expr)
            if table and columns:
                index_name = _extract_index_name(index_expr) or sanitize_identifier(fallback_index_name)
                params = index_expr.args.get("params")
                where_expr = params.args.get("where") if isinstance(params, exp.IndexParameters) else None
                return {
                    "name": index_name or sanitize_identifier(fallback_index_name),
                    "table": table,
                    "columns": columns,
                    "unique": bool(tree.args.get("unique") or index_expr.args.get("unique")),
                    "partial_where": where_expr.sql(dialect=dialect) if where_expr is not None else None,
                }

    table = sanitize_identifier(fallback_table)
    index_name = sanitize_identifier(fallback_index_name) or fallback_index_name
    open_paren = sql_text.find("(")
    close_paren = sql_text.rfind(")")
    if open_paren == -1 or close_paren == -1 or close_paren <= open_paren:
        return None
    raw_columns = sql_text[open_paren + 1 : close_paren]
    columns = []
    for token in raw_columns.split(","):
        cleaned = sanitize_identifier(token.split()[0])
        if cleaned:
            columns.append(cleaned)
    columns = dedupe_preserve_order(columns)
    if not table or not columns:
        return None
    return {
        "name": index_name,
        "table": table,
        "columns": columns,
        "unique": "unique index" in sql_text.lower(),
        "partial_where": None,
    }


def _extract_index_table(index_expr: exp.Index) -> str:
    table_expr = index_expr.args.get("table")
    if isinstance(table_expr, exp.Table):
        return sanitize_identifier(table_expr.name or "")
    if isinstance(table_expr, exp.Expression):
        return sanitize_identifier(table_expr.sql())
    return ""


def _extract_index_name(index_expr: exp.Index) -> str:
    this_expr = index_expr.args.get("this")
    if isinstance(this_expr, exp.Identifier):
        return sanitize_identifier(this_expr.name or "")
    if isinstance(this_expr, exp.Expression):
        return sanitize_identifier(this_expr.sql())
    return ""


def _extract_index_columns(index_expr: exp.Index) -> list[str]:
    params = index_expr.args.get("params")
    if not isinstance(params, exp.IndexParameters):
        return []
    columns = params.args.get("columns")
    if not isinstance(columns, list):
        return []

    parsed: list[str] = []
    for column_expr in columns:
        inner = column_expr.this if isinstance(column_expr, exp.Ordered) else column_expr
        if isinstance(inner, exp.Column):
            col_name = sanitize_identifier(inner.name or "")
        elif isinstance(inner, exp.Identifier):
            col_name = sanitize_identifier(inner.name or "")
        else:
            col_name = sanitize_identifier(inner.sql() if isinstance(inner, exp.Expression) else "")
        if col_name:
            parsed.append(col_name)
    return dedupe_preserve_order(parsed)


def _dedupe_existing_indexes(indexes: list[ExistingIndex]) -> list[ExistingIndex]:
    seen: set[tuple[str, tuple[str, ...], bool, str | None]] = set()
    deduped: list[ExistingIndex] = []
    for index in indexes:
        key = (index.table, tuple(index.columns), index.unique, index.partial_where)
        if key in seen:
            continue
        seen.add(key)
        deduped.append(index)
    return deduped
