"""adapters for converting chat interface to completion interface.

This module provides a adapters class that implements completion methods by
delegating to chat methods. This is useful for LLM providers that primarily
support a chat interface but need to implement the completion interface as well.
"""

from __future__ import annotations

from typing import Any, Literal, overload

from serapeum.core.base.llms.types import (
    ChatResponse,
    CompletionResponse,
    CompletionResponseAsyncGen,
    CompletionResponseGen,
    MessageList,
)

__all__ = ["ChatToCompletion"]


class ChatToCompletion:
    """adapters that provides completion methods by delegating to chat methods.

    This adapters adds completion method implementations for classes that have
    chat methods. It uses duck typing and assumes the class it's mixed into
    implements the following methods:

    - ``chat(messages, stream, **kwargs) -> ChatResponse | ChatResponseGen``
    - ``achat(messages, stream, **kwargs) -> ChatResponse | ChatResponseAsyncGen``

    The adapters provides implementations for:

    - ``complete(prompt, *, stream=False, **kwargs)``
    - ``stream_complete(prompt, **kwargs)`` (shim → ``complete(stream=True)``)
    - ``acomplete(prompt, *, stream=False, **kwargs)``
    - ``astream_complete(prompt, **kwargs)`` (shim → ``acomplete(stream=True)``)

    Examples:
        Basic usage with a custom LLM class:

        ```python
        >>> from serapeum.core.llms import FunctionCallingLLM, ChatResponse, Message, MessageRole
        >>> from serapeum.core.base.llms.types import TextChunk
        >>> from serapeum.core.llms.abstractions.adapters import ChatToCompletion
        >>>
        >>> class MyLLM(ChatToCompletion, FunctionCallingLLM):
        ...     def chat(self, messages, *, stream=False, **kwargs):
        ...         if stream:
        ...             return self._stream_chat(messages, **kwargs)
        ...         return ChatResponse(
        ...             message=Message(
        ...                 role=MessageRole.ASSISTANT,
        ...                 chunks=[TextChunk(content="Response")],
        ...             )
        ...         )
        ...     def _stream_chat(self, messages, **kwargs):
        ...         yield ChatResponse(
        ...             message=Message(
        ...                 role=MessageRole.ASSISTANT,
        ...                 chunks=[TextChunk(content="Response")],
        ...             ),
        ...             delta="Response"
        ...         )
        ...     async def achat(self, messages, *, stream=False, **kwargs):
        ...         if stream:
        ...             return await self._astream_chat(messages, **kwargs)
        ...         return ChatResponse(
        ...             message=Message(
        ...                 role=MessageRole.ASSISTANT,
        ...                 chunks=[TextChunk(content="Response")],
        ...             )
        ...         )
        ...     async def _astream_chat(self, messages, **kwargs):
        ...         async def gen():
        ...             yield ChatResponse(
        ...                 message=Message(
        ...                     role=MessageRole.ASSISTANT,
        ...                     chunks=[TextChunk(content="Response")],
        ...                 ),
        ...                 delta="Response"
        ...             )
        ...         return gen()
        ...     @property
        ...     def metadata(self):
        ...         from serapeum.core.llms import Metadata
        ...         return Metadata(is_chat_model=True, is_function_calling_model=True)
        ...     def _prepare_chat_with_tools(self, tools, **kwargs):
        ...         return {"messages": [], "tools": []}
        >>>
        >>> # Now MyLLM has complete() and other completion methods automatically
        >>> llm = MyLLM(model="example")

        ```

    See Also:
        - :mod:`serapeum.core.base.llms.utils`: Decorator-based chat-to-completion adapters
        - :class:`serapeum.core.llms.base.LLM`: Base LLM class
    """

    @overload
    def complete(
        self,
        prompt: str,
        formatted: bool = ...,
        *,
        stream: Literal[False] = ...,
        **kwargs: Any,
    ) -> CompletionResponse: ...

    @overload
    def complete(
        self,
        prompt: str,
        formatted: bool = ...,
        *,
        stream: Literal[True],
        **kwargs: Any,
    ) -> CompletionResponseGen: ...

    def complete(
        self,
        prompt: str,
        formatted: bool = False,
        *,
        stream: bool = False,
        **kwargs: Any,
    ) -> CompletionResponse | CompletionResponseGen:
        """Implement completion by delegating to the chat method.

        Converts the prompt string to a MessageList and calls the chat method,
        then converts the response to the appropriate completion type.

        Args:
            prompt: The prompt string to complete.
            formatted: Whether the prompt is already formatted (unused, for compatibility).
            stream: If True, returns a streaming generator instead of a single response.
            **kwargs: Additional arguments passed to the chat method.

        Returns:
            CompletionResponse when stream=False, CompletionResponseGen when stream=True.

        Examples:
            - Non-streaming completion returns full response text:
                ```python
                >>> from serapeum.core.llms import (
                ...     FunctionCallingLLM, ChatResponse, Message, MessageRole, Metadata,
                ... )
                >>> from serapeum.core.base.llms.types import TextChunk
                >>> from serapeum.core.llms.abstractions.adapters import ChatToCompletion
                >>>
                >>> class DemoLLM(ChatToCompletion, FunctionCallingLLM):
                ...     @property
                ...     def metadata(self):
                ...         return Metadata(is_chat_model=True, is_function_calling_model=True)
                ...     def chat(self, messages, *, stream=False, **kwargs):
                ...         if stream:
                ...             def gen():
                ...                 yield ChatResponse(
                ...                     message=Message(
                ...                         role=MessageRole.ASSISTANT,
                ...                         chunks=[TextChunk(content="Hi")],
                ...                     ),
                ...                     delta="Hi",
                ...                 )
                ...             return gen()
                ...         return ChatResponse(
                ...             message=Message(
                ...                 role=MessageRole.ASSISTANT,
                ...                 chunks=[TextChunk(content="Hi")],
                ...             ),
                ...         )
                ...     async def achat(self, messages, *, stream=False, **kwargs):
                ...         return ChatResponse(
                ...             message=Message(
                ...                 role=MessageRole.ASSISTANT,
                ...                 chunks=[TextChunk(content="Hi")],
                ...             ),
                ...         )
                ...     def _prepare_chat_with_tools(self, tools, **kwargs):
                ...         return {"messages": [], "tools": []}
                >>>
                >>> llm = DemoLLM(model="test")
                >>> response = llm.complete("Hello")
                >>> response.text
                'Hi'

                ```
        """
        messages = MessageList.from_str(prompt)
        if stream:
            chat_response_gen = self.chat(messages, stream=True, **kwargs)  # type: ignore[attr-defined]
            result: CompletionResponse | CompletionResponseGen = (
                ChatResponse.stream_to_completion_response(chat_response_gen)
            )
        else:
            chat_response = self.chat(messages, stream=False, **kwargs)  # type: ignore[attr-defined]
            result = chat_response.to_completion_response()
        return result

    @overload
    async def acomplete(
        self,
        prompt: str,
        formatted: bool = ...,
        *,
        stream: Literal[False] = ...,
        **kwargs: Any,
    ) -> CompletionResponse: ...

    @overload
    async def acomplete(
        self,
        prompt: str,
        formatted: bool = ...,
        *,
        stream: Literal[True],
        **kwargs: Any,
    ) -> CompletionResponseAsyncGen: ...

    async def acomplete(
        self,
        prompt: str,
        formatted: bool = False,
        *,
        stream: bool = False,
        **kwargs: Any,
    ) -> CompletionResponse | CompletionResponseAsyncGen:
        """Implement async completion by delegating to the achat method.

        Converts the prompt string to a MessageList and calls the achat method,
        then converts the response to the appropriate completion type.

        Args:
            prompt: The prompt string to complete.
            formatted: Whether the prompt is already formatted (unused, for compatibility).
            stream: If True, returns an async streaming generator instead of a single response.
            **kwargs: Additional arguments passed to the achat method.

        Returns:
            CompletionResponse when stream=False, CompletionResponseAsyncGen when stream=True.

        Examples:
            - Async non-streaming completion:
                ```python
                >>> import asyncio
                >>> from serapeum.core.llms import (
                ...     FunctionCallingLLM, ChatResponse, Message, MessageRole, Metadata,
                ... )
                >>> from serapeum.core.base.llms.types import TextChunk
                >>> from serapeum.core.llms.abstractions.adapters import ChatToCompletion
                >>>
                >>> class DemoLLM(ChatToCompletion, FunctionCallingLLM):
                ...     @property
                ...     def metadata(self):
                ...         return Metadata(is_chat_model=True, is_function_calling_model=True)
                ...     def chat(self, messages, *, stream=False, **kwargs):
                ...         return ChatResponse(
                ...             message=Message(
                ...                 role=MessageRole.ASSISTANT,
                ...                 chunks=[TextChunk(content="Hi")],
                ...             ),
                ...         )
                ...     async def achat(self, messages, *, stream=False, **kwargs):
                ...         return ChatResponse(
                ...             message=Message(
                ...                 role=MessageRole.ASSISTANT,
                ...                 chunks=[TextChunk(content="Hi")],
                ...             ),
                ...         )
                ...     def _prepare_chat_with_tools(self, tools, **kwargs):
                ...         return {"messages": [], "tools": []}
                >>>
                >>> async def example():
                ...     llm = DemoLLM(model="test")
                ...     response = await llm.acomplete("Hello")
                ...     return response.text
                >>> asyncio.run(example())
                'Hi'

                ```
        """
        messages = MessageList.from_str(prompt)
        if stream:
            chat_response_gen = await self.achat(messages, stream=True, **kwargs)  # type: ignore[attr-defined]
            result: CompletionResponse | CompletionResponseAsyncGen = (
                ChatResponse.astream_to_completion_response(chat_response_gen)
            )
        else:
            chat_response = await self.achat(messages, stream=False, **kwargs)  # type: ignore[attr-defined]
            result = chat_response.to_completion_response()
        return result
