from collections.abc import AsyncIterator, Awaitable, Callable
from typing import TypeVar

from aidial_sdk.exceptions import HTTPException as DialException
from pydantic import BaseModel

from aidial_interceptors_sdk.chat_completion.annotated_value import (
    AnnotatedException,
    AnnotatedValue,
    Annotation,
)
from aidial_interceptors_sdk.chat_completion.request_handler import (
    RequestHandler,
)
from aidial_interceptors_sdk.chat_completion.response_handler import (
    ResponseHandler,
)
from aidial_interceptors_sdk.dial_client import DialClient
from aidial_interceptors_sdk.utils.streaming import annotate_stream

RequestDict = dict

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


class ChatCompletionInterceptor(RequestHandler, ResponseHandler):
    dial_client: DialClient
    configuration: BaseModel | None

    @classmethod
    async def configuration_schema(cls) -> type[BaseModel] | None:
        return None

    def get_configuration(self, cls: type[_T]) -> _T:
        if self.configuration is None:
            raise ValueError("Can't find interceptor configuration")

        if not isinstance(self.configuration, cls):
            raise TypeError(
                f"The interceptor configuration is of type {type(self.configuration).__name__!r}, but expected to be of type {cls.__name__!r}"
            )

        return self.configuration

    async def call_upstreams(
        self,
        request: RequestDict,
        call_upstream: Callable[
            [Annotation, RequestDict],
            Awaitable[AsyncIterator[dict | DialException]],
        ],
    ) -> AsyncIterator[AnnotatedValue]:
        annotation = None
        return annotate_stream(
            annotation, await call_upstream(annotation, request)
        )

    async def on_stream_start(self) -> None:
        # TODO: it's probably worth to put all the chunks
        # generated by this method into a separate list.
        # And then merge them all with the first incoming chunk.
        # Otherwise, we may end up with choice being open *before*
        # its "assistant" role is reported.
        pass

    async def on_stream_error(self, error: AnnotatedException) -> None:
        raise error.error

    async def on_stream_end(self) -> None:
        # TODO: it's probably worth to withhold the last chunk generated by
        # on_stream_chunk and merge it with all the chunks reported by on_stream_end.
        pass


class ChatCompletionNoOpInterceptor(ChatCompletionInterceptor):
    pass
