import json
from datetime import date, datetime, time, timedelta
from decimal import Decimal
from typing import Any
from uuid import UUID

from dbslice.config import DatabaseType
from dbslice.models import Table


class SQLGenerator:
    """Generates SQL INSERT statements from extracted data."""

    def __init__(
        self,
        db_type: DatabaseType = DatabaseType.POSTGRESQL,
        include_transaction: bool = True,
        include_truncate: bool = False,
        disable_fk_checks: bool = False,
        schema: str | None = None,
    ):
        self.db_type = db_type
        self.include_transaction = include_transaction
        self.include_truncate = include_truncate
        self.disable_fk_checks = disable_fk_checks
        self.schema = schema

    def generate(
        self,
        tables_data: dict[str, list[dict[str, Any]]],
        insert_order: list[str],
        tables_schema: dict[str, Table],
        broken_fks: list[Any] | None = None,
        deferred_updates: list[Any] | None = None,
    ) -> str:
        """
        Generate complete SQL output with cycle handling support.

        Args:
            tables_data: dict mapping table name to list of row dicts
            insert_order: Tables in topologically sorted order for insertion
            tables_schema: dict mapping table name to Table schema
            broken_fks: List of ForeignKey objects that were broken to handle cycles
            deferred_updates: List of DeferredUpdate objects for restoring FK values

        Returns:
            Complete SQL string with all INSERT statements
        """
        lines: list[str] = []

        broken_fks = broken_fks or []
        deferred_updates = deferred_updates or []

        total_rows = sum(len(rows) for rows in tables_data.values())
        lines.append("-- Generated by dbslice")
        lines.append(f"-- Tables: {len(tables_data)}, Rows: {total_rows}")
        if broken_fks:
            lines.append(f"-- Circular references detected: {len(broken_fks)} FK(s) broken")
        lines.append("")

        if self.schema and self.schema != "public":
            lines.append(f"SET search_path TO {self._quote_identifier(self.schema)}, public;")
            lines.append("")

        if self.disable_fk_checks:
            lines.append(self._disable_fk_checks())
            lines.append("")

        if self.include_transaction:
            lines.append("BEGIN;")
            lines.append("")

        if self.include_truncate:
            for table in insert_order:
                if table in tables_data:
                    lines.append(f"TRUNCATE TABLE {self._quote_identifier(table)} CASCADE;")
            lines.append("")

        broken_fk_cols = self._build_broken_fk_map(broken_fks)

        # INSERT statements in topological order (with NULLs for broken FKs)
        for table in insert_order:
            if table not in tables_data:
                continue

            rows = tables_data[table]
            if not rows:
                continue

            table_schema = tables_schema.get(table)
            lines.append(f"-- {table} ({len(rows)} rows)")

            for row in rows:
                null_cols = broken_fk_cols.get(table)
                lines.append(self._generate_insert(table, row, table_schema, null_cols))

            lines.append("")

        # Deferred UPDATE statements to restore FK values
        if deferred_updates:
            lines.append("-- Restore circular foreign key references")
            for update in deferred_updates:
                lines.append(self._generate_deferred_update(update, tables_schema))
            lines.append("")

        if self.include_transaction:
            lines.append("COMMIT;")

        if self.disable_fk_checks:
            lines.append("")
            lines.append(self._enable_fk_checks())

        return "\n".join(lines)

    def _generate_insert(
        self,
        table: str,
        data: dict[str, Any],
        table_schema: Table | None,
        null_columns: set[str] | None = None,
    ) -> str:
        """
        Generate a single INSERT statement.

        Args:
            table: Table name
            data: Row data dict
            table_schema: Table schema
            null_columns: Optional set of column names to set to NULL (for breaking circular deps)

        Returns:
            INSERT statement
        """
        null_columns = null_columns or set()
        columns = list(data.keys())

        col_types: dict[str, str] = {}
        if table_schema:
            for col in table_schema.columns:
                col_types[col.name] = col.data_type

        values = [
            "NULL" if col in null_columns else self._format_value(data[col], col_types.get(col))
            for col in columns
        ]

        cols_str = ", ".join(self._quote_identifier(c) for c in columns)
        vals_str = ", ".join(values)

        return f"INSERT INTO {self._quote_identifier(table)} ({cols_str}) VALUES ({vals_str});"

    def _format_value(self, value: Any, column_type: str | None = None) -> str:
        """Format a Python value as SQL literal.

        Args:
            value: The Python value to format
            column_type: Optional column data type for proper array/JSON handling
        """
        if value is None:
            return "NULL"

        if isinstance(value, bool):
            if self.db_type == DatabaseType.MYSQL:
                return "1" if value else "0"
            return "TRUE" if value else "FALSE"

        if isinstance(value, int | float | Decimal):
            return str(value)

        if isinstance(value, bytes):
            if self.db_type == DatabaseType.POSTGRESQL:
                return f"E'\\\\x{value.hex()}'"
            elif self.db_type == DatabaseType.MYSQL:
                return f"X'{value.hex()}'"
            else:  # SQLite
                return f"X'{value.hex()}'"

        if isinstance(value, datetime):
            return self._escape_string(value.isoformat())

        if isinstance(value, date):
            return self._escape_string(value.isoformat())

        if isinstance(value, time):
            return self._escape_string(value.isoformat())

        if isinstance(value, timedelta):
            # Convert to interval string
            total_seconds = int(value.total_seconds())
            hours, remainder = divmod(total_seconds, 3600)
            minutes, seconds = divmod(remainder, 60)
            return self._escape_string(f"{hours:02d}:{minutes:02d}:{seconds:02d}")

        if isinstance(value, UUID):
            return self._escape_string(str(value))

        if isinstance(value, list):
            # Check if this is a PostgreSQL array column
            if column_type and self._is_array_type(column_type):
                return self._format_pg_array(value)
            # Otherwise treat as JSON
            json_str = json.dumps(value, default=str, ensure_ascii=False)
            return self._escape_json_string(json_str)

        if isinstance(value, dict):
            # JSON/JSONB data
            json_str = json.dumps(value, default=str, ensure_ascii=False)
            return self._escape_json_string(json_str)

        # String value - check if it's going into a JSON column
        str_value = str(value)
        if column_type and self._is_json_type(column_type):
            # String going into JSON column - need to ensure it's valid JSON
            # If it looks like JSON, re-serialize to ensure proper escaping
            try:
                parsed = json.loads(str_value)
                json_str = json.dumps(parsed, default=str, ensure_ascii=False)
                return self._escape_json_string(json_str)
            except (json.JSONDecodeError, TypeError):
                # Not valid JSON, wrap as JSON string
                json_str = json.dumps(str_value, ensure_ascii=False)
                return self._escape_json_string(json_str)

        return self._escape_string(str_value)

    def _is_array_type(self, column_type: str) -> bool:
        """Check if column type is a PostgreSQL array type."""
        if not column_type:
            return False
        col_upper = column_type.upper()
        # PostgreSQL array types: text[], varchar[], integer[], etc.
        # Also internal types like _text, _varchar, _int4
        return "[]" in column_type or col_upper.startswith("_") or "ARRAY" in col_upper

    def _is_json_type(self, column_type: str) -> bool:
        """Check if column type is a JSON/JSONB type."""
        if not column_type:
            return False
        col_upper = column_type.upper()
        return col_upper in ("JSON", "JSONB")

    def _format_pg_array(self, value: list) -> str:
        """Format a Python list as PostgreSQL array literal.

        PostgreSQL array syntax: '{item1,item2,item3}'
        For strings: '{"item1","item2"}'
        """
        if not value:
            return "'{}'"

        def format_element(elem: Any) -> str:
            if elem is None:
                return "NULL"
            if isinstance(elem, bool):
                return "t" if elem else "f"
            if isinstance(elem, int | float | Decimal):
                return str(elem)
            # String elements need to be double-quoted and escaped
            s = str(elem)
            # Escape backslashes and double quotes
            s = s.replace("\\", "\\\\").replace('"', '\\"')
            return f'"{s}"'

        elements = [format_element(e) for e in value]
        array_str = "{" + ",".join(elements) + "}"
        return self._escape_string(array_str)

    def _escape_string(self, s: str) -> str:
        """Escape a string for SQL."""
        escaped = s.replace("'", "''")

        if self.db_type == DatabaseType.POSTGRESQL:
            # Handle backslashes in PostgreSQL
            if "\\" in escaped:
                escaped = escaped.replace("\\", "\\\\")
                return f"E'{escaped}'"

        return f"'{escaped}'"

    def _escape_json_string(self, json_str: str) -> str:
        """Escape a JSON string for SQL insertion.

        JSON strings need special handling because they contain backslash
        escape sequences (like \\n, \\t) that must be preserved.
        """
        # For JSON, we need to:
        # 1. Escape single quotes for SQL
        # 2. Double the backslashes for PostgreSQL E'' strings so JSON escapes are preserved
        escaped = json_str.replace("'", "''")

        if self.db_type == DatabaseType.POSTGRESQL:
            # Double backslashes so \n becomes \\n in E'' string,
            # which PostgreSQL will interpret as literal \n (the JSON escape)
            escaped = escaped.replace("\\", "\\\\")
            return f"E'{escaped}'"

        return f"'{escaped}'"

    def _quote_identifier(self, name: str) -> str:
        """Quote a table or column name."""
        if self.db_type == DatabaseType.MYSQL:
            return f"`{name}`"
        # PostgreSQL and SQLite use double quotes
        return f'"{name}"'

    def _disable_fk_checks(self) -> str:
        """Generate SQL to disable FK checks."""
        if self.db_type == DatabaseType.MYSQL:
            return "SET FOREIGN_KEY_CHECKS = 0;"
        if self.db_type == DatabaseType.SQLITE:
            return "PRAGMA foreign_keys = OFF;"
        # PostgreSQL uses deferred constraints
        return "SET CONSTRAINTS ALL DEFERRED;"

    def _enable_fk_checks(self) -> str:
        """Generate SQL to re-enable FK checks."""
        if self.db_type == DatabaseType.MYSQL:
            return "SET FOREIGN_KEY_CHECKS = 1;"
        if self.db_type == DatabaseType.SQLITE:
            return "PRAGMA foreign_keys = ON;"
        return "SET CONSTRAINTS ALL IMMEDIATE;"

    def _build_broken_fk_map(self, broken_fks: list[Any]) -> dict[str, set[str]]:
        """
        Build a map of table -> set of FK column names that were broken.

        Args:
            broken_fks: List of ForeignKey objects that were broken

        Returns:
            Dict mapping table name to set of FK column names
        """
        broken_fk_cols: dict[str, set[str]] = {}

        for fk in broken_fks:
            table = fk.source_table
            if table not in broken_fk_cols:
                broken_fk_cols[table] = set()

            for col in fk.source_columns:
                broken_fk_cols[table].add(col)

        return broken_fk_cols

    def _generate_deferred_update(
        self,
        update: Any,  # DeferredUpdate
        tables_schema: dict[str, Table],
    ) -> str:
        """
        Generate UPDATE statement to restore a broken FK value.

        Args:
            update: DeferredUpdate object with UPDATE details
            tables_schema: Table schemas for formatting values

        Returns:
            UPDATE statement to set FK column to its correct value
        """
        # Build WHERE clause
        where_conditions = []
        for col, val in zip(update.pk_columns, update.pk_values):
            col_quoted = self._quote_identifier(col)
            val_formatted = self._format_value(val)
            where_conditions.append(f"{col_quoted} = {val_formatted}")

        where_clause = " AND ".join(where_conditions)

        # Build SET clause
        fk_col_quoted = self._quote_identifier(update.fk_column)
        fk_val_formatted = self._format_value(update.fk_value)

        return (
            f"UPDATE {self._quote_identifier(update.table)} "
            f"SET {fk_col_quoted} = {fk_val_formatted} "
            f"WHERE {where_clause};"
        )


def generate_sql(
    tables_data: dict[str, list[dict[str, Any]]],
    insert_order: list[str],
    tables_schema: dict[str, Table],
    db_type: DatabaseType = DatabaseType.POSTGRESQL,
) -> str:
    """
    Convenience function to generate SQL output.

    Args:
        tables_data: dict mapping table name to list of row dicts
        insert_order: Tables in topologically sorted order
        tables_schema: dict mapping table name to Table schema
        db_type: Target database type

    Returns:
        Complete SQL string
    """
    generator = SQLGenerator(db_type=db_type)
    return generator.generate(tables_data, insert_order, tables_schema)
