"""Select query builder."""

from collections.abc import Generator
from textwrap import dedent
from typing import Any, Self, cast

from embar.column.base import ColumnBase
from embar.db.base import AllDbBase, AsyncDbBase, DbBase
from embar.model import (
    DataModel,
    generate_model,
    load_results,
)
from embar.query.clause_base import ClauseBase
from embar.query.order_by import Asc, BareColumn, Desc, OrderBy, RawSqlOrder
from embar.query.query import QuerySingle
from embar.sql import Sql
from embar.table import Table


class DeleteQueryReady[T: Table, Db: AllDbBase]:
    """
    `DeleteQueryReady` is used to delete data from a table.

    It is generic over the `Table` being inserted into, and the database being used.

    `DeleteQueryReady` is returned by `db.delete()`
    """

    table: type[T]
    _db: Db

    _where_clause: ClauseBase | None = None
    _order_clause: OrderBy | None = None
    _limit_value: int | None = None

    def __init__(self, table: type[T], db: Db):
        """
        Create a new SelectQueryReady instance.
        """
        self.table = table
        self._db = db

    def returning(self, use_pydantic: bool = True) -> DeleteQueryReturning[T, Db]:
        return DeleteQueryReturning(
            table=self.table,
            db=self._db,
            use_pydantic=use_pydantic,
            where_clause=self._where_clause,
            order_clause=self._order_clause,
            limit_value=self._limit_value,
        )

    def where(self, where_clause: ClauseBase) -> Self:
        """
        Add a WHERE clause to the query.
        """
        self._where_clause = where_clause
        return self

    def order_by(self, *clauses: ColumnBase | Asc | Desc | Sql) -> Self:
        """
        Add an ORDER BY clause to sort query results.

        Accepts multiple ordering clauses:
        - Bare column references (defaults to ASC): `User.id`
        - `Asc(User.id)` or `Asc(User.id, nulls="last")`
        - `Desc(User.id)` or `Desc(User.id, nulls="first")`
        - Raw SQL: `Sql(t"{User.id} DESC")`

        Can be called multiple times to add more sort columns.

        ```python
        from embar.db.pg import PgDb
        from embar.table import Table
        from embar.column.common import Integer, integer

        class User(Table):
            id: Integer = integer(primary=True)

        db = PgDb(None)

        query = db.delete(User).order_by(User.id)
        sql_result = query.sql()
        assert "ORDER BY" in sql_result.sql
        ```
        """
        # Convert each clause to an OrderByClause
        order_clauses: list[ClauseBase] = []
        for clause in clauses:
            if isinstance(clause, (Asc, Desc)):
                order_clauses.append(clause)
            elif isinstance(clause, Sql):
                order_clauses.append(RawSqlOrder(clause))
            else:
                order_clauses.append(BareColumn(clause))

        if self._order_clause is None:
            self._order_clause = OrderBy(tuple(order_clauses))
        else:
            # Add to existing ORDER BY clauses
            self._order_clause = OrderBy((*self._order_clause.clauses, *order_clauses))

        return self

    def limit(self, n: int) -> Self:
        """
        Add a LIMIT clause to the query.
        """
        self._limit_value = n
        return self

    def __await__(self):
        """
        Async users should construct their query and await it.

        Non-async users have the `run()` convenience method below.
        But this method will still work if called in an async context against a non-async db.

        The overrides provide for a few different cases:
        - A Model was passed, in which case that's the return type
        - `SelectAllPydantic` or `SelectAllDataclass` was passed (via `returning()`),
          in which case the return type is the `Table`
        - This is called with an async db, in which case an error is returned.
        """
        query = self.sql()

        async def awaitable():
            db = self._db
            if isinstance(db, AsyncDbBase):
                await db.execute(query)
            else:
                db = cast(DbBase, self._db)
                db.execute(query)

        return awaitable().__await__()

    def run(self) -> None:
        """
        Run the query against the underlying DB.

        Convenience method for those not using async.
        For async, use `await query` instead.
        """
        if isinstance(self._db, DbBase):
            query = self.sql()
            self._db.execute(query)

    def sql(self) -> QuerySingle:
        """
        Combine all the components of the query and build the SQL and bind parameters (psycopg format).
        """

        sql = f"""
        DELETE
        FROM {self.table.fqn()}
        """
        sql = dedent(sql).strip()

        count = -1

        def get_count() -> int:
            nonlocal count
            count += 1
            return count

        params: dict[str, Any] = {}

        if self._where_clause is not None:
            where_data = self._where_clause.sql(get_count)
            sql += f"\nWHERE {where_data.sql}"
            params = {**params, **where_data.params}

        if self._order_clause is not None:
            order_by_query = self._order_clause.sql(get_count)
            sql += f"\nORDER BY {order_by_query.sql}"
            params = {**params, **order_by_query.params}

        if self._limit_value is not None:
            sql += f"\nLIMIT {self._limit_value}"

        sql = sql.strip()

        return QuerySingle(sql, params=params)


class DeleteQueryReturning[T: Table, Db: AllDbBase]:
    """
    `DeleteQueryReturning` is used to delete data from a table and return it.
    """

    table: type[T]
    _db: Db
    _use_pydantic: bool

    _where_clause: ClauseBase | None = None
    _order_clause: OrderBy | None = None
    _limit_value: int | None = None

    def __init__(
        self,
        table: type[T],
        db: Db,
        use_pydantic: bool,
        where_clause: ClauseBase | None,
        order_clause: OrderBy | None,
        limit_value: int | None,
    ):
        """
        Create a new DeleteQueryReturning instance.
        """
        self.table = table
        self._db = db
        self._use_pydantic = use_pydantic
        self._where_clause = where_clause
        self._order_clause = order_clause
        self._limit_value = limit_value

    def __await__(self) -> Generator[Any, None, list[T]]:
        """
        Async users should construct their query and await it.

        Non-async users have the `run()` convenience method below.
        But this method will still work if called in an async context against a non-async db.

        The overrides provide for a few different cases:
        - A Model was passed, in which case that's the return type
        - `SelectAllPydantic` or `SelectAllDataclass` was passed (via `returning()`),
          in which case the return type is the `Table`
        - This is called with an async db, in which case an error is returned.
        """
        query = self.sql()
        model = self._get_model()
        model = cast(type[T], model)

        async def awaitable():
            db = self._db
            if isinstance(db, AsyncDbBase):
                data = await db.fetch(query)
            else:
                db = cast(DbBase, self._db)
                data = db.fetch(query)
            results = load_results(model, data)
            return results

        return awaitable().__await__()

    def run(self) -> list[T]:
        """
        Run the query against the underlying DB.

        Convenience method for those not using async.
        For async, use `await query` instead.
        """
        query = self.sql()
        model = self._get_model()
        model = cast(type[T], model)
        db = cast(DbBase, self._db)
        data = db.fetch(query)
        results = load_results(model, data)
        return results

    def sql(self) -> QuerySingle:
        """
        Combine all the components of the query and build the SQL and bind parameters (psycopg format).
        """

        sql = f"""
        DELETE
        FROM {self.table.fqn()}
        """
        sql = dedent(sql).strip()

        count = -1

        def get_count() -> int:
            nonlocal count
            count += 1
            return count

        params: dict[str, Any] = {}

        if self._where_clause is not None:
            where_data = self._where_clause.sql(get_count)
            sql += f"\nWHERE {where_data.sql}"
            params = {**params, **where_data.params}

        if self._order_clause is not None:
            order_by_query = self._order_clause.sql(get_count)
            sql += f"\nORDER BY {order_by_query.sql}"
            params = {**params, **order_by_query.params}

        if self._limit_value is not None:
            sql += f"\nLIMIT {self._limit_value}"

        # This is the only difference vs the regular DeleteQueryReady sql
        sql += "\nRETURNING *"

        sql = sql.strip()

        return QuerySingle(sql, params=params)

    def _get_model(self) -> type[DataModel]:
        """
        Generate the dataclass that will be used to deserialize (and validate) the query results.
        """
        model = generate_model(self.table, self._use_pydantic)
        return model
