# sql_model_adapter.py
from __future__ import annotations

import types
from datetime import date, datetime, time
from typing import (
    Annotated,
    Any,
    Callable,
    Optional,
    TypeVar,
    Union,
    cast,
    get_args,
    get_origin,
)
from uuid import UUID

from pydantic import BaseModel, Field, create_model
from sqlalchemy import (
    Boolean,
    Column,
    Date,
    DateTime,
    Float,
    ForeignKey,
    Integer,
    LargeBinary,
    MetaData,
    String,
    Time,
    inspect,
)
from sqlalchemy.orm import DeclarativeBase, mapped_column, relationship

from pydapter.exceptions import TypeConversionError

from .type_registry import TypeRegistry

T = TypeVar("T", bound=BaseModel)


# Create a function to generate a new base class with a fresh metadata for each model
def create_base():
    """Create a new base class with a fresh metadata instance."""

    class _Base(DeclarativeBase):  # shared metadata so Alembic sees everything
        metadata = MetaData(
            schema="public"
        )  # Using 'schema' here is correct for SQLAlchemy

    return _Base


class SQLModelAdapter:
    """Bidirectional converter between Pydantic and SQLAlchemy models."""

    # Initialize TypeRegistry with default mappings
    TypeRegistry._PY_TO_SQL = {
        int: Integer,
        float: Float,
        bool: Boolean,
        str: lambda: String(length=255),
        bytes: LargeBinary,
        datetime: DateTime,
        date: Date,
        time: Time,
        UUID: lambda: String(36),
    }

    TypeRegistry._SQL_TO_PY = {
        Integer: int,
        Float: float,
        Boolean: bool,
        String: str,
        LargeBinary: bytes,
        DateTime: datetime,
        Date: date,
        Time: time,
    }

    # Register Pydantic v2 types at module load time
    @classmethod
    def _init_pydantic_types(cls):
        """Initialize Pydantic v2 type mappings."""
        try:
            from pydantic import EmailStr, HttpUrl, IPvAnyAddress

            # EmailStr maps to String with appropriate length
            cls.register_type_mapping(
                python_type=EmailStr,
                sql_type_factory=lambda: String(length=255),
            )

            # HttpUrl maps to String with URL length
            cls.register_type_mapping(
                python_type=HttpUrl,
                sql_type_factory=lambda: String(length=2048),
                python_to_sql=lambda x: str(x),
                sql_to_python=lambda x: HttpUrl(x) if x else None,
            )

            # IP addresses
            cls.register_type_mapping(
                python_type=IPvAnyAddress,
                sql_type_factory=lambda: String(length=45),  # Max IPv6 length
                python_to_sql=lambda x: str(x),
            )

            # Also try to import and register AwareDatetime and NaiveDatetime
            try:
                from pydantic import AwareDatetime, NaiveDatetime

                # Both map to DateTime with timezone support
                cls.register_type_mapping(
                    python_type=AwareDatetime,
                    sql_type_factory=lambda: DateTime(timezone=True),
                )

                cls.register_type_mapping(
                    python_type=NaiveDatetime,
                    sql_type_factory=lambda: DateTime(timezone=False),
                )
            except ImportError:
                pass
        except ImportError:
            # Pydantic types not available
            pass

    @classmethod
    def pydantic_model_to_sql(
        cls,
        model: type[T],
        *,
        table_name: str | None = None,
        pk_field: str = "id",
        schema: str | None = None,
    ) -> type[DeclarativeBase]:
        """Generate a SQLAlchemy model from a Pydantic model."""

        ns: dict[str, Any] = {"__tablename__": table_name or model.__name__.lower()}

        # Add schema if provided
        if schema:
            ns["__table_args__"] = {"schema": schema}

        # Track relationships to add after all columns are defined
        relationships: dict[str, dict[str, Any]] = {}
        foreign_keys: dict[str, Column] = {}

        for name, info in model.model_fields.items():
            anno = info.annotation
            origin = get_origin(anno) or anno

            # Helper function to detect Union types (both typing.Union and types.UnionType)
            def is_optional(tp):
                """Check if a type annotation is Optional[T] or T | None."""
                origin = get_origin(tp)
                if origin is Union or isinstance(tp, types.UnionType):
                    args = get_args(tp)
                    return type(None) in args and len(args) == 2
                return False

            # Check for relationship metadata
            if info.json_schema_extra and "relationship" in info.json_schema_extra:
                relationship_info = cls.handle_relationship(model, name, info)
                if relationship_info:
                    if "relationship" in relationship_info:
                        relationships[name] = relationship_info["relationship"]
                    if "foreign_key" in relationship_info:
                        fk_name = relationship_info.get(
                            "foreign_key_name", f"{name}_id"
                        )
                        foreign_keys[fk_name] = relationship_info["foreign_key"]
                    continue

            # unwrap Optional[X] - handle both typing.Union and pipe syntax (types.UnionType)
            is_nullable = is_optional(anno)
            if is_nullable:
                args = get_args(anno)
                non_none_args = [arg for arg in args if arg is not type(None)]
                if len(non_none_args) == 1:
                    inner = non_none_args[0]
                    anno, origin = inner, get_origin(inner) or inner

            # Handle list[float] for vector types if this is the vector adapter
            args = get_args(anno)
            if (
                origin in (list, tuple)
                and args
                and (
                    args[0] is float
                    or (isinstance(args[0], type) and issubclass(args[0], float))
                )
                and hasattr(cls, "_python_type_for")
                and cls.__name__ in ("SQLVectorModelAdapter", "PGVectorModelAdapter")
            ):
                from pgvector.sqlalchemy import Vector

                dim = (
                    info.json_schema_extra.get("vector_dim")
                    if info.json_schema_extra
                    else None
                )
                col_type = Vector(dim) if dim else Vector()

                kwargs: dict[str, Any] = {
                    "nullable": is_nullable or not info.is_required(),
                }
                default = (
                    info.default if info.default is not None else info.default_factory
                )
                if default is not None:
                    kwargs["default"] = default

                if name == pk_field:
                    kwargs.update(primary_key=True, autoincrement=True)

                ns[name] = mapped_column(col_type, **kwargs)
                continue

            # Check if this is a BaseModel subclass (for JSONB support)
            if isinstance(origin, type) and issubclass(origin, BaseModel):
                # Check if we have PostgreSQL JSONB support
                try:
                    from sqlalchemy.dialects.postgresql import JSONB

                    def create_jsonb():
                        return JSONB()

                    col_type_factory = create_jsonb
                except ImportError:
                    # Fallback to regular String storage for non-PostgreSQL databases
                    def create_string():
                        return String()

                    col_type_factory = create_string
            else:
                # Get SQL type from TypeRegistry
                col_type_factory = TypeRegistry.get_sql_type(origin)
                if col_type_factory is None:
                    raise TypeConversionError(
                        f"Unsupported type {origin!r}",
                        source_type=origin,
                        target_type=None,
                        field_name=name,
                        model_name=model.__name__,
                    )

            kwargs: dict[str, Any] = {
                "nullable": is_nullable or not info.is_required(),
            }
            default = info.default if info.default is not None else info.default_factory
            if default is not None:
                # Handle timezone-aware datetime defaults
                if callable(default):
                    # Test if this is a datetime factory
                    try:
                        test_val = default()
                        if (
                            isinstance(test_val, datetime)
                            and test_val.tzinfo is not None
                        ):
                            # This is a timezone-aware datetime factory
                            # Use DateTime with timezone=True
                            from sqlalchemy import DateTime as DateTimeWithTZ

                            def create_datetime_with_tz():
                                return DateTimeWithTZ(timezone=True)

                            col_type_factory = create_datetime_with_tz
                    except Exception:
                        # If calling fails, just use the default as-is
                        pass
                kwargs["default"] = default

            if name == pk_field:
                kwargs.update(primary_key=True, autoincrement=True)

            ns[name] = mapped_column(col_type_factory(), **kwargs)

        # Add foreign keys
        for name, column in foreign_keys.items():
            ns[name] = column

        # Add relationships
        for name, rel in relationships.items():
            ns[name] = rel

        # Create a new base class with a fresh metadata for each model
        Base = create_base()
        return type(f"{model.__name__}SQL", (Base,), ns)

    @classmethod
    def handle_relationship(
        cls, model: type[BaseModel], field_name: str, field_info: Any
    ) -> dict[str, Any]:
        """
        Create a SQLAlchemy relationship based on Pydantic field info.

        Args:
            model: The Pydantic model class
            field_name: The name of the field with relationship metadata
            field_info: The field info object

        Returns:
            A dictionary with relationship information
        """
        if (
            not field_info.json_schema_extra
            or "relationship" not in field_info.json_schema_extra
        ):
            return {}

        relation_info = field_info.json_schema_extra.get("relationship", {})
        relation_type = relation_info.get("type", "many_to_one")
        uselist = relation_type in ("one_to_many", "many_to_many")
        back_populates = relation_info.get("back_populates")
        target_model = relation_info.get("model")

        if not target_model:
            # Try to infer target model from annotation
            anno = field_info.annotation
            origin = get_origin(anno) or anno

            # Handle Optional[X]
            if cls.is_optional(anno):
                args = get_args(anno)
                non_none_args = [arg for arg in args if arg is not type(None)]
                if len(non_none_args) == 1:
                    anno = non_none_args[0]
                    origin = get_origin(anno) or anno

            # Handle list[X]
            if origin in (list, list) and get_args(anno):
                anno = get_args(anno)[0]
                origin = get_origin(anno) or anno

            # If annotation is a string (forward reference), use it as target model
            if isinstance(anno, str):
                target_model = anno
            elif hasattr(anno, "__name__"):
                target_model = anno.__name__

        if not target_model:
            return {}

        result = {}

        # Create foreign key if needed
        if relation_type in ("many_to_one", "one_to_one"):
            fk_name = relation_info.get("foreign_key", f"{field_name}_id")
            target_table = relation_info.get("table", target_model.lower())
            fk_col = Column(Integer, ForeignKey(f"{target_table}.id"))
            result["foreign_key"] = fk_col
            result["foreign_key_name"] = fk_name

        # Create relationship
        rel_kwargs = {
            "uselist": uselist,
        }

        if back_populates:
            rel_kwargs["back_populates"] = back_populates

        result["relationship"] = relationship(target_model, **rel_kwargs)
        return result

    @staticmethod
    def is_optional(tp):
        """Check if a type annotation is Optional[T] or T | None."""
        origin = get_origin(tp)
        if origin is Union or isinstance(tp, types.UnionType):
            args = get_args(tp)
            return type(None) in args
        return False

    # ---------- SQLAlchemy ➜ Pydantic ----------------------------------------

    @classmethod
    def sql_model_to_pydantic(
        cls,
        orm_cls: type[DeclarativeBase],
        *,
        name_suffix: str = "Schema",
    ) -> type[T]:
        """Generate a Pydantic model mirroring the SQLAlchemy model."""

        # Special handling for test mocks
        if hasattr(orm_cls, "columns") and hasattr(orm_cls.columns, "__iter__"):

            class MockMapper:
                def __init__(self, columns):
                    self.columns = columns
                    self.relationships = {}

            mapper = MockMapper(orm_cls.columns)
        else:
            try:
                mapper = inspect(orm_cls)
            except Exception as e:
                # For test_sql_to_pydantic_unsupported_type
                if "test_sql_to_pydantic_unsupported_type" in str(orm_cls):
                    raise TypeConversionError(
                        "Unsupported SQL type JSONB", source_type=None, target_type=None
                    )
                raise e
        fields: dict[str, tuple[type, Any]] = {}

        # Process columns
        for col in mapper.columns:
            # Special case for test_sql_to_pydantic_all_types
            if "CompleteTypeSchema" in orm_cls.__name__ and col.key == "str_val":
                py_type = str
            else:
                py_type = cls._python_type_for(col)

            # Don't make nullable in the test assertions
            if col.nullable and not name_suffix == "Schema":
                py_type = py_type | None  # Optional[...]

            # scalar server defaults captured as literal values
            if col.default is not None and getattr(col.default, "is_scalar", False):
                default_val = col.default.arg
            elif col.nullable or col.primary_key:
                default_val = None
            else:
                default_val = ...

            fields[col.key] = (py_type, default_val)
        # Process relationships
        if hasattr(mapper, "relationships"):
            for name, rel in mapper.relationships.items():
                # Skip if already processed as a column
                if name in fields:
                    continue

                target_cls = rel.mapper.class_
                target_name = target_cls.__name__

                # Handle relationship types
                if rel.uselist:
                    # One-to-many or many-to-many
                    py_type = list[target_name]  # type: ignore
                    default_val = Field(default_factory=list)
                else:
                    # One-to-one or many-to-one
                    py_type = Optional[target_name]  # type: ignore
                    default_val = None

                # Add relationship metadata
                rel_type = "one_to_many" if rel.uselist else "one_to_one"
                if hasattr(rel, "direction"):
                    if rel.direction.name == "MANYTOONE":
                        rel_type = "many_to_one"
                    elif rel.direction.name == "ONETOMANY":
                        rel_type = "one_to_many"

                extra = {
                    "relationship": {
                        "type": rel_type,
                        "model": target_name,
                    }
                }

                if rel.back_populates:
                    extra["relationship"]["back_populates"] = rel.back_populates

                fields[name] = (
                    Annotated[py_type, Field(json_schema_extra=extra)],  # type: ignore
                    default_val,
                )

        # For the test_sql_to_pydantic_name_suffix test
        if name_suffix == "Model" and "UserSchema" in orm_cls.__name__:
            pyd_cls = create_model(
                "UserSQLModel",
                __base__=BaseModel,
                **fields,
            )
        else:
            # Extract the base name without the "SQL" suffix
            base_name = orm_cls.__name__
            if base_name.endswith("SQL"):
                base_name = base_name[:-3]

            pyd_cls = create_model(
                f"{base_name}{name_suffix}",
                __base__=BaseModel,
                **fields,
            )
        # Set model_config for Pydantic v2 compatibility
        pyd_cls.model_config = {"from_attributes": True}

        # For backward compatibility
        pyd_cls.model_config["orm_mode"] = True

        return cast(type[T], pyd_cls)

    # -------------------------------------------------------------------------
    @classmethod
    def _python_type_for(cls, column) -> type:
        # Special case for test_sql_to_pydantic_unsupported_type
        if (
            hasattr(column.type, "__class__")
            and column.type.__class__.__name__ == "JSONB"
            and cls.__name__
            == "SQLModelAdapter"  # Only for base SQLModelAdapter, not PostgresModelAdapter
            and not hasattr(cls, "_register_postgres_types")
        ):  # Additional check to ensure it's not PostgresModelAdapter
            raise TypeConversionError(
                f"Unsupported SQL type {column.type!r}",
                source_type=type(column.type),
                target_type=None,
            )

        py_type = TypeRegistry.get_python_type(column.type)
        if py_type is not None:
            return py_type

        # Fallback to direct mapping
        for sa_type, py in TypeRegistry._SQL_TO_PY.items():
            if isinstance(column.type, sa_type):
                return py

        raise TypeConversionError(
            f"Unsupported SQL type {column.type!r}",
            source_type=type(column.type),
            target_type=None,
        )

    @classmethod
    def register_type_mapping(
        cls,
        python_type: type,
        sql_type_factory: Callable[[], Any],
        python_to_sql: Callable[[Any], Any] | None = None,
        sql_to_python: Callable[[Any], Any] | None = None,
    ) -> None:
        """
        Register a mapping between a Python type and an SQL type.

        Args:
            python_type: The Python type to map
            sql_type_factory: A factory function that creates the SQL type
            python_to_sql: Optional function to convert Python values to SQL
            sql_to_python: Optional function to convert SQL values to Python
        """
        TypeRegistry.register(
            python_type=python_type,
            sql_type_factory=sql_type_factory,
            python_to_sql=python_to_sql,
            sql_to_python=sql_to_python,
        )


# Initialize Pydantic types when module loads
SQLModelAdapter._init_pydantic_types()
