from typing import Any, Optional, List, TYPE_CHECKING

from .base import DatabaseBackend
from atomsql.orm.fields import (
    IntegerField,
    StringField,
    CharField,
    TextField,
    ForeignKey,
    FloatField,
    DateField,
    DateTimeField,
    DateTimeTZField,
    TimeStampField,
    BooleanField,
    EmailField,
    URLField,
    JSONField,
    UUIDField,
    ManyToManyField,
)

if TYPE_CHECKING:
    import psycopg


class PostgresBackend(DatabaseBackend):
    def __init__(self, db_path: str) -> None:
        self.db_path = db_path
        self.connection: Optional["psycopg.Connection"] = None
        self.cursor: Optional["psycopg.Cursor"] = None

    def connect(self, **kwargs: Any) -> None:
        import psycopg

        self.connection = psycopg.connect(self.db_path, **kwargs)
        self.cursor = self.connection.cursor()

    def disconnect(self) -> None:
        if self.connection:
            self.connection.close()

    def commit(self) -> None:
        if self.connection:
            self.connection.commit()

    def rollback(self) -> None:
        self.connection.rollback()

    def close(self) -> None:
        self.disconnect()

    def get_bulk_insert_ids(self, cursor, num_records: int) -> List[int]:
        """
        Extract inserted IDs after a bulk insert for PostgreSQL.
        PostgreSQL uses RETURNING clause to return IDs directly.
        """
        # Fetch all returned IDs from RETURNING clause
        return [row[0] for row in cursor.fetchall()]

    def bulk_insert(self, sql: str, values: List[Any], num_records: int) -> List[int]:
        """
        Execute bulk insert for PostgreSQL with RETURNING clause.
        PostgreSQL supports RETURNING id to get all inserted IDs.
        """
        # Add RETURNING clause if not present
        if "RETURNING" not in sql.upper():
            sql = sql + " RETURNING id"

        cursor = self.execute(sql, values)
        return self.get_bulk_insert_ids(cursor, num_records)

    def get_field_type(self, field) -> str:
        if isinstance(field, BooleanField):
            return "BOOLEAN"
        elif isinstance(field, TimeStampField):
            return "BIGINT"
        elif isinstance(field, IntegerField):
            return "INTEGER"
        elif isinstance(field, CharField):
            return f"VARCHAR({field.max_length})"
        elif isinstance(field, (TextField, EmailField, URLField, StringField)):
            return "TEXT"
        elif isinstance(field, ForeignKey):
            return f'INTEGER REFERENCES "{field.to._table_name}"(id)'
        elif isinstance(field, ManyToManyField):
            # ManyToManyField doesn't have a direct column type
            return None
        elif isinstance(field, FloatField):
            return "NUMERIC"
        elif isinstance(field, DateField):
            return "DATE"
        elif isinstance(field, DateTimeTZField):
            # Check DateTimeTZ before DateTimeField since it inherits from DateTimeField
            return "TIMESTAMP WITH TIME ZONE"
        elif isinstance(field, DateTimeField):
            return "TIMESTAMP WITHOUT TIME ZONE"
        elif isinstance(field, JSONField):
            return "JSONB"
        elif isinstance(field, UUIDField):
            return "UUID"
        return "TEXT"

    def get_primary_key_constraint(self, column_name, field_type):
        return f'"{column_name}" SERIAL PRIMARY KEY'

    # PostgreSQL supports both CHECK constraints and ENUM types
    # For now, we use CHECK constraints from the base class
    # Future enhancement: could create ENUM types for better performance

    @property
    def placeholder_char(self) -> str:
        return "%s"

    def get_tables(self, exclude_migrations: bool = False) -> list:
        """Get list of tables in database."""
        cursor = self.connection.cursor()
        cursor.execute(
            """
            SELECT tablename
            FROM pg_tables
            WHERE schemaname = 'public'
            """
        )
        tables = [row[0] for row in cursor.fetchall()]

        if exclude_migrations:
            tables = [t for t in tables if t != "_migrations"]

        return tables

    def get_table_columns(self, table_name: str) -> list:
        """Get columns for a table."""
        cursor = self.connection.cursor()
        cursor.execute(
            """
            SELECT column_name, data_type, is_nullable
            FROM information_schema.columns
            WHERE table_name = %s
            ORDER BY ordinal_position
            """,
            (table_name,),
        )
        return cursor.fetchall()

    def prepare_value(self, value: Any) -> Any:
        """Prepare a value for insertion into the database."""
        if isinstance(value, (dict, list)):
            import json

            return json.dumps(value)
        return value

    def get_json_extract_sql(
        self, column: str, path: List[str], cast_type: str = None
    ) -> str:
        """
        Postgres JSONB extraction.
        Uses -> for intermediate steps and ->> for the final step to get text.
        """
        if not path:
            return column

        sql = column
        for i, key in enumerate(path):
            clean_key = key.replace("'", "''")

            # Use ->> for the last element to get the value as text
            # Use -> for intermediate elements (returns json object)
            op = "->>" if i == len(path) - 1 else "->"
            sql = f"{sql}{op}'{clean_key}'"

        if cast_type:
            sql = f"({sql}::{cast_type})"

        return sql
