from __future__ import annotations

import asyncio
import logging
from abc import ABC, abstractmethod
from collections.abc import Mapping, Sequence
from copy import copy
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Generic, Optional, Protocol, TypeVar, cast

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.sql.dml import Insert

from phoenix.db import models
from phoenix.db.insertion.constants import DEFAULT_RETRY_ALLOWANCE, DEFAULT_RETRY_DELAY_SEC
from phoenix.db.insertion.helpers import insert_on_conflict
from phoenix.server.dml_event import DmlEvent
from phoenix.server.types import DbSessionFactory

logger = logging.getLogger(__name__)


class Insertable(Protocol):
    @property
    def row(self) -> models.Base: ...


_AnyT = TypeVar("_AnyT")
_PrecursorT = TypeVar("_PrecursorT")
_InsertableT = TypeVar("_InsertableT", bound=Insertable)
_RowT = TypeVar("_RowT", bound=models.Base)
_DmlEventT = TypeVar("_DmlEventT", bound=DmlEvent)


@dataclass(frozen=True)
class Received(Generic[_AnyT]):
    item: _AnyT
    received_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    def postpone(self, retries_left: int = DEFAULT_RETRY_ALLOWANCE) -> Postponed[_AnyT]:
        return Postponed(item=self.item, received_at=self.received_at, retries_left=retries_left)


@dataclass(frozen=True)
class Postponed(Received[_AnyT]):
    retries_left: int = field(default=DEFAULT_RETRY_ALLOWANCE)


class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
    table: type[_RowT]
    unique_by: Sequence[str]

    def __init_subclass__(
        cls,
        table: type[_RowT],
        unique_by: Sequence[str],
    ) -> None:
        cls.table = table
        cls.unique_by = unique_by

    def __init__(
        self,
        db: DbSessionFactory,
        retry_delay_sec: float = DEFAULT_RETRY_DELAY_SEC,
        retry_allowance: int = DEFAULT_RETRY_ALLOWANCE,
    ) -> None:
        self._queue: list[Received[_PrecursorT]] = []
        self._db = db
        self._retry_delay_sec = retry_delay_sec
        self._retry_allowance = retry_allowance

    @property
    def empty(self) -> bool:
        return not bool(self._queue)

    async def enqueue(self, *items: _PrecursorT) -> None:
        self._queue.extend([Received(item) for item in items])

    @abstractmethod
    async def _partition(
        self,
        session: AsyncSession,
        *parcels: Received[_PrecursorT],
    ) -> tuple[
        list[Received[_InsertableT]],
        list[Postponed[_PrecursorT]],
        list[Received[_PrecursorT]],
    ]: ...

    async def insert(self) -> Optional[list[_DmlEventT]]:
        if not self._queue:
            return None
        self._queue, parcels = [], self._queue
        events: list[_DmlEventT] = []
        async with self._db() as session:
            to_insert, to_postpone, _ = await self._partition(session, *parcels)
            if to_insert:
                events, to_retry, _ = await self._insert(session, *to_insert)
                if to_retry:
                    to_postpone.extend(to_retry)
        if to_postpone:
            loop = asyncio.get_running_loop()
            loop.call_later(self._retry_delay_sec, self._queue.extend, to_postpone)
        return events

    def _insert_on_conflict(self, *records: Mapping[str, Any]) -> Insert:
        return insert_on_conflict(
            *records,
            table=self.table,
            unique_by=self.unique_by,
            dialect=self._db.dialect,
        )

    @abstractmethod
    async def _events(
        self,
        session: AsyncSession,
        *insertions: _InsertableT,
    ) -> list[_DmlEventT]: ...

    async def _insert(
        self,
        session: AsyncSession,
        *parcels: Received[_InsertableT],
    ) -> tuple[
        list[_DmlEventT],
        list[Postponed[_PrecursorT]],
        list[Received[_InsertableT]],
    ]:
        to_retry: list[Postponed[_PrecursorT]] = []
        failures: list[Received[_InsertableT]] = []
        events: list[_DmlEventT] = []
        try:
            async with session.begin_nested():
                events.extend(await self._events(session, *(p.item for p in parcels)))
        except BaseException:
            logger.exception(
                f"Failed to bulk insert for {self.table.__name__}. "
                f"Will try to insert ({len(parcels)} records) individually instead."
            )
            for p in parcels:
                try:
                    async with session.begin_nested():
                        events.extend(await self._events(session, p.item))
                except BaseException:
                    logger.exception(f"Failed to insert for {self.table.__name__}.")
                    if isinstance(p, Postponed) and p.retries_left == 1:
                        failures.append(p)
                    else:
                        to_retry.append(
                            Postponed(
                                item=cast(_PrecursorT, p.item),
                                received_at=p.received_at,
                                retries_left=(p.retries_left - 1)
                                if isinstance(p, Postponed)
                                else self._retry_allowance,
                            )
                        )
        return events, to_retry, failures


class Precursors(ABC):
    @dataclass(frozen=True)
    class SpanAnnotation:
        span_id: str
        obj: models.SpanAnnotation

        def as_insertable(
            self,
            span_rowid: int,
            id_: Optional[int] = None,
        ) -> Insertables.SpanAnnotation:
            return Insertables.SpanAnnotation(
                span_id=self.span_id,
                obj=self.obj,
                span_rowid=span_rowid,
                id_=id_,
            )

    @dataclass(frozen=True)
    class TraceAnnotation:
        trace_id: str
        obj: models.TraceAnnotation

        def as_insertable(
            self,
            trace_rowid: int,
            id_: Optional[int] = None,
        ) -> Insertables.TraceAnnotation:
            return Insertables.TraceAnnotation(
                trace_id=self.trace_id,
                obj=self.obj,
                trace_rowid=trace_rowid,
                id_=id_,
            )

    @dataclass(frozen=True)
    class DocumentAnnotation:
        span_id: str
        document_position: int
        obj: models.DocumentAnnotation

        def as_insertable(
            self,
            span_rowid: int,
            id_: Optional[int] = None,
        ) -> Insertables.DocumentAnnotation:
            return Insertables.DocumentAnnotation(
                span_id=self.span_id,
                document_position=self.document_position,
                obj=self.obj,
                span_rowid=span_rowid,
                id_=id_,
            )


class Insertables(ABC):
    @dataclass(frozen=True)
    class SpanAnnotation(Precursors.SpanAnnotation):
        span_rowid: int
        id_: Optional[int] = None

        @property
        def row(self) -> models.SpanAnnotation:
            obj = copy(self.obj)
            obj.span_rowid = self.span_rowid
            if self.id_ is not None:
                obj.id = self.id_
            return obj

    @dataclass(frozen=True)
    class TraceAnnotation(Precursors.TraceAnnotation):
        trace_rowid: int
        id_: Optional[int] = None

        @property
        def row(self) -> models.TraceAnnotation:
            obj = copy(self.obj)
            obj.trace_rowid = self.trace_rowid
            if self.id_ is not None:
                obj.id = self.id_
            return obj

    @dataclass(frozen=True)
    class DocumentAnnotation(Precursors.DocumentAnnotation):
        span_rowid: int
        id_: Optional[int] = None

        @property
        def row(self) -> models.DocumentAnnotation:
            obj = copy(self.obj)
            obj.span_rowid = self.span_rowid
            if self.id_ is not None:
                obj.id = self.id_
            return obj
