import json
import typing as t
from contextlib import nullcontext
from uuid import UUID

import sqlalchemy as sa
import sqlmodel as sm
import uuid_utils
from sqlalchemy.dialects.postgresql import insert as postgres_insert
from sqlalchemy.orm import declared_attr
from sqlalchemy.orm.attributes import flag_modified as sa_flag_modified
from sqlmodel import Column, Field, Session, SQLModel, inspect, select
from typeid import TypeID

from activemodel.mixins.pydantic_json import PydanticJSONMixin

# NOTE: this patches a core method in sqlmodel to support db comments
from .patches import get_column_from_field_patch  # noqa: F401
from .query_wrapper import QueryWrapper
from .utils import to_snake_case
from .session_manager import get_session

POSTGRES_INDEXES_NAMING_CONVENTION = {
    "ix": "%(column_0_label)s_idx",
    "uq": "%(table_name)s_%(column_0_name)s_key",
    "ck": "%(table_name)s_%(constraint_name)s_check",
    "fk": "%(table_name)s_%(column_0_name)s_fkey",
    "pk": "%(table_name)s_pkey",
}
"""
By default, the foreign key naming convention in sqlalchemy do not create unique identifiers when there are multiple
foreign keys in a table. This naming convention is a workaround to fix this issue:

- https://github.com/zhanymkanov/fastapi-best-practices?tab=readme-ov-file#set-db-keys-naming-conventions
- https://github.com/fastapi/sqlmodel/discussions/1213
- Implementation lifted from: https://github.com/AlexanderZharyuk/billing-service/blob/3c8aaf19ab7546b97cc4db76f60335edec9fc79d/src/models.py#L24
"""

SQLModel.metadata.naming_convention = POSTGRES_INDEXES_NAMING_CONVENTION


class BaseModel(SQLModel):
    """
    Base model class to inherit from so we can hate python less.

    Some notes:

    - Inspired by https://github.com/woofz/sqlmodel-basecrud/blob/main/sqlmodel_basecrud/basecrud.py
    - lifecycle hooks are modeled after Rails.
    - class docstrings are converted to table-level comments
    - save(), delete(), select(), where(), and other easy methods you would expect in a real ORM
    - Fixes foreign key naming conventions
    - Sane table names

    Here's how hooks work:

        Create/Update: before_create, after_create, before_update, after_update, before_save, after_save, around_save
        Delete: before_delete, after_delete, around_delete
        Read: after_find, after_initialize

    around_* hooks must be context managers (method returning a CM or a CM attribute).
    Ordering (create): before_create -> before_save -> (enter around_save) -> persist -> after_create -> after_save -> (exit around_save)
    Ordering (update): before_update -> before_save -> (enter around_save) -> persist -> after_update -> after_save -> (exit around_save)
    Delete: before_delete -> (enter around_delete) -> delete -> after_delete -> (exit around_delete)
    Read: finder/query method -> after_find -> after_initialize
    Construction: Model(...) -> after_initialize

        # TODO document this in activemodel, this is an interesting edge case
    # https://claude.ai/share/f09e4f70-2ff7-4cd0-abff-44645134693a

    """

    __table_args__ = None

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._call_hook("after_initialize")

    @classmethod
    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)

        # Enables field-level docstrings on the pydantic `description` field, which we
        # copy into table/column comments by patching SQLModel internals elsewhere.
        cls.model_config["use_attribute_docstrings"] = True

        cls._apply_class_doc()

    @classmethod
    def _apply_class_doc(cls):
        """
        Pull class-level docstring into a table comment.

        This will help AI SQL writers like: https://github.com/iloveitaly/sql-ai-prompt-generator
        """

        doc = cls.__doc__.strip() if cls.__doc__ else None

        if doc:
            table_args = getattr(cls, "__table_args__", None)

            if table_args is None:
                cls.__table_args__ = {"comment": doc}
            elif isinstance(table_args, dict):
                table_args.setdefault("comment", doc)
            elif isinstance(table_args, tuple):
                # If it's a tuple, we need to convert it to a list and add the comment
                table_args = list(table_args)
                table_args.append({"comment": doc})
                cls.__table_args__ = tuple(table_args)
            else:
                raise ValueError(
                    f"Unexpected __table_args__ type {type(table_args)}, expected dictionary."
                )

    # TODO no type check decorator here
    @declared_attr
    def __tablename__(cls) -> str:
        """
        Automatically generates the table name for the model by converting the model's class name from PascalCase to snake_case.
        This is the recommended text case style for table names:

        https://wiki.postgresql.org/wiki/Don%27t_Do_This#Don.27t_use_upper_case_table_or_column_names

        By default, the model's class name is lower cased which makes it harder to read.

        This implementation properly handles acronyms: e.g. `LLMCache` -> `llm_cache`.
        https://stackoverflow.com/questions/1175208/elegant-python-function-to-convert-camelcase-to-snake-case
        """
        return to_snake_case(cls.__name__)

    @classmethod
    def foreign_key(cls, **kwargs) -> t.Any:
        """
        Returns a `Field` object referencing the foreign key of the model.

        Helps quickly build a many-to-one or one-to-one relationship.

        >>> other_model_id: TypeID = OtherModel.foreign_key()
        >>> other_model = Relationship()
        """

        field_options = {"nullable": False} | kwargs

        return Field(
            # TODO id field is hard coded, should pick the PK field in case it's different
            sa_type=cls.model_fields["id"].sa_column.type,  # type: ignore
            foreign_key=f"{cls.__tablename__}.id",
            **field_options,
        )

    @classmethod
    def select(cls, *args: t.Any):
        "create a query wrapper to easily run sqlmodel queries on this model"
        return QueryWrapper(cls, *args)

    @classmethod
    def where(cls, *args: t.Any):
        "convenience method to avoid having to write .select().where() in order to add conditions"
        return cls.select().where(*args)

    # TODO we should add an instance method for this as well
    @classmethod
    def upsert(
        cls,
        data: dict[str, t.Any],
        unique_by: str | list[str],
    ) -> t.Self:
        """
        This method will insert a new record if it doesn't exist, or update the existing record if it does.

        It uses SQLAlchemy's `on_conflict_do_update` and does not yet support MySQL. Some implementation details below.

        ---

        - `index_elements=["name"]`: Specifies the column(s) to check for conflicts (e.g., unique constraint or index). If a row with the same "name" exists, it triggers the update instead of an insert.
        - `values`: Defines the data to insert (e.g., `name="example", value=123`). If no conflict occurs, this data is inserted as a new row.

        The `set_` parameter (e.g., `set_=dict(value=123)`) then dictates what gets updated on conflict, overriding matching fields in `values` if specified.
        """
        index_elements = [unique_by] if isinstance(unique_by, str) else unique_by

        stmt = (
            postgres_insert(cls)
            .values(**data)
            .on_conflict_do_update(index_elements=index_elements, set_=data)
            .returning(cls)
        )

        with get_session() as session:
            result = session.exec(stmt)
            session.commit()

            # TODO this is so ugly:
            result = result.one()[0]

        return result

    def delete(self):
        """Delete instance running delete hooks and optional around_delete context manager."""

        cm = self._get_around_context_manager("around_delete") or nullcontext()

        with get_session() as session:
            if (
                old_session := Session.object_session(self)
            ) and old_session is not session:
                old_session.expunge(self)
            session.delete(self)

            self._call_hook("before_delete")
            with cm:
                session.commit()
            self._call_hook("after_delete")

        return True

    def save(self):
        """Persist instance running create/update hooks and optional around_save context manager."""

        is_new = self.is_new()
        cm = self._get_around_context_manager("around_save") or nullcontext()

        with get_session() as session:
            if (
                old_session := Session.object_session(self)
            ) and old_session is not session:
                old_session.expunge(self)

            session.add(self)

            # the order and placement of these hooks is really important
            # we need the current object to be in a session otherwise it will not be able to
            # load any relationships.
            self._call_hook("before_create" if is_new else "before_update")
            self._call_hook("before_save")

            with cm:
                session.commit()
                session.refresh(self)

            self._call_hook("after_create" if is_new else "after_update")
            self._call_hook("after_save")

            # Only call the transform method if the class is a subclass of PydanticJSONMixin
            if issubclass(self.__class__, PydanticJSONMixin):
                self.__class__.__transform_dict_to_pydantic__(self)
        return self

    def _call_hook(self, hook_name: str) -> None:
        method = getattr(self, hook_name, None)
        if callable(method):
            if method.__code__.co_argcount != 1:
                raise TypeError(
                    f"Hook '{hook_name}' must accept exactly 1 positional argument (self)"
                )
            method()

    @classmethod
    @t.overload
    def _run_after_find_hook(cls, instance: None) -> None: ...

    @classmethod
    @t.overload
    def _run_after_find_hook[T](cls, instance: T) -> T: ...

    @classmethod
    def _run_after_find_hook(cls, instance: t.Any):
        """Run the Rails-style `after_find` hook for model instances only.

        This helper exists because `after_find` is different from the save/delete
        hooks above:

        - it is not tied to a single persistence entrypoint, so multiple finder and
          query methods need the same behavior
        - some query paths return `None` or scalar values, which should pass
          through untouched
        - the hook should run while the SQLAlchemy session is still active so model
          callbacks can safely touch lazy relationships
        """
        if not isinstance(instance, BaseModel):
            return instance

        instance._call_hook("after_find")
        return instance

    @classmethod
    @t.overload
    def _run_after_load_hooks(cls, instance: None) -> None: ...

    @classmethod
    @t.overload
    def _run_after_load_hooks[T](cls, instance: T) -> T: ...

    @classmethod
    def _run_after_load_hooks(cls, instance: t.Any | None):
        if not isinstance(instance, BaseModel):
            return instance

        instance = cls._run_after_find_hook(instance)
        instance._call_hook("after_initialize")
        return instance

    def _get_around_context_manager(self, name: str) -> t.ContextManager | None:
        obj = getattr(self, name, None)
        if obj is None:
            return None

        # If it's a callable (method/function), call it to obtain the CM
        if callable(obj):
            obj = obj()

        cm = obj
        if not (hasattr(cm, "__enter__") and hasattr(cm, "__exit__")):
            raise TypeError(
                f"{name} must return or be a context manager implementing __enter__/__exit__"
            )
        return t.cast(t.ContextManager, cm)

    def refresh(self):
        "Refreshes an object from the database"

        with get_session() as session:
            if (
                old_session := Session.object_session(self)
            ) and old_session is not session:
                old_session.expunge(self)

            session.add(self)
            session.refresh(self)

            # Only call the transform method if the class is a subclass of PydanticJSONMixin
            if issubclass(self.__class__, PydanticJSONMixin):
                self.__class__.__transform_dict_to_pydantic__(self)

        return self

    # TODO shouldn't this be handled by pydantic?
    # TODO where is this actually used? shoudl prob remove this
    # TODO should we even do this? Can we specify a better json rendering class?
    def json(self, **kwargs):
        return json.dumps(self.model_dump(), default=str, **kwargs)

    # TODO should move this to the wrapper
    @classmethod
    def count(cls) -> int:
        """
        Returns the number of records in the database.
        """
        with get_session() as session:
            return session.scalar(sm.select(sm.func.count()).select_from(cls))

    # TODO got to be a better way to fwd these along...
    @classmethod
    def first(cls):
        # TODO should use dynamic pk
        return cls.select().order_by(sa.desc(cls.id)).first()

    # @classmethod
    # def last(cls):
    #     return cls.select().first()

    # TODO throw an error if this field is set on the model
    def is_new(self) -> bool:
        return not self._sa_instance_state.has_identity

    def flag_modified(self, *args: str) -> None:
        """
        Flag one or more fields as modified/mutated/dirty. Useful for marking a field containing sub-objects as modified.

        Will throw an error if an invalid field is passed.
        """

        assert len(args) > 0, "Must pass at least one field name"

        for field_name in args:
            if field_name not in type(self).model_fields:
                raise ValueError(f"Field '{field_name}' does not exist in the model.")

            # check if the field exists
            sa_flag_modified(self, field_name)

    def modified_fields(self) -> set[str]:
        "set of fields that are modified"

        insp = inspect(self)
        assert insp is not None

        return {attr.key for attr in insp.attrs if attr.history.has_changes()}

    @classmethod
    def find_or_create_by(cls, **kwargs):
        """
        Find record or create it with the passed args if it doesn't exist.
        """

        result = cls.get(**kwargs)

        if result:
            return result

        new_model = cls(**kwargs)
        new_model.save()

        return new_model

    @classmethod
    def find_or_initialize_by(cls, **kwargs):
        """
        Unfortunately, unlike ruby, python does not have a great lambda story. This makes writing convenience methods
        like this a bit more difficult.
        """

        result = cls.get(**kwargs)

        if result:
            return result

        new_model = cls(**kwargs)
        return new_model

    @classmethod
    def primary_key_column(cls) -> Column:
        """
        Returns the primary key column of the model by inspecting SQLAlchemy field information.

        >>> ExampleModel.primary_key_field().name
        """

        # TODO note_schema.__class__.__table__.primary_key
        # TODO no reason why this couldn't be cached

        pk_columns = list(cls.__table__.primary_key.columns)

        if not pk_columns:
            raise ValueError("No primary key defined for the model.")

        if len(pk_columns) > 1:
            raise ValueError(
                "Multiple primary keys defined. This method supports only single primary key models."
            )

        return pk_columns[0]

    # TODO what's super dangerous here is you pass a kwarg which does not map to a specific
    #      field it will result in `True`, which will return all records, and not give you any typing
    #      errors. Dangerous when iterating on structure quickly
    # TODO can we pass the generic of the superclass in?
    # TODO can we type the method signature a bit better?
    # def get(cls, *args: sa.BinaryExpression, **kwargs: t.Any):
    @classmethod
    def get(cls, *args: t.Any, **kwargs: t.Any):
        """
        Gets a single record (or None) from the database. Pass an PK ID or kwargs to filter by.
        """
        # TODO id is hardcoded, not good! Need to dynamically pick the best uid field
        id_field_name = "id"

        # special case for getting by ID
        if len(args) == 1 and isinstance(
            args[0], (int, TypeID, str, UUID, uuid_utils.UUID)
        ):
            kwargs[id_field_name] = args[0]
            args = ()

        statement = select(cls).filter(*args).filter_by(**kwargs)

        with get_session() as session:
            result = session.exec(statement).first()
            return cls._run_after_load_hooks(result)

    @classmethod
    def one_or_none(cls, *args: t.Any, **kwargs: t.Any):
        """
        Gets a single record from the database. Pass an PK ID or a kwarg to filter by.
        Returns None if no record is found. Throws an error if more than one record is found.
        """

        args, kwargs = cls.__process_filter_args__(*args, **kwargs)
        statement = select(cls).filter(*args).filter_by(**kwargs)

        with get_session() as session:
            result = session.exec(statement).one_or_none()
            return cls._run_after_load_hooks(result)

    @classmethod
    def one(cls, *args: t.Any, **kwargs: t.Any):
        """
        Gets a single record from the database. Pass an PK ID or a kwarg to filter by.
        """

        args, kwargs = cls.__process_filter_args__(*args, **kwargs)
        statement = select(cls).filter(*args).filter_by(**kwargs)

        with get_session() as session:
            result = session.exec(statement).one()
            return cls._run_after_load_hooks(result)

    @classmethod
    def __process_filter_args__(cls, *args: t.Any, **kwargs: t.Any):
        """
        Helper method to process filter arguments and implement some nice DX for our devs.
        """

        id_field_name = cls.primary_key_column().name

        # special case for getting by ID without having to specify the field name
        # TODO should dynamically add new pk types based on column definition
        if len(args) == 1 and isinstance(
            args[0], (int, TypeID, str, UUID, uuid_utils.UUID)
        ):
            kwargs[id_field_name] = args[0]
            args = ()

        return args, kwargs

    @classmethod
    def all(cls):
        "get a generator for all records in the database"
        with get_session() as session:
            results = session.exec(sm.select(cls))

            # TODO do we need this or can we just return results?
            for result in results:
                yield cls._run_after_load_hooks(result)

    @classmethod
    def sample(cls):
        """
        Pick a random record from the database. Raises if none exist.

        Helpful for testing and console debugging.
        """

        query = sm.select(cls).order_by(sa.sql.func.random()).limit(1)

        with get_session() as session:
            result = session.exec(query).one()
            return cls._run_after_load_hooks(result)
