import asyncio
from time import sleep

try:
    from jupyter_ai_magics.providers import BaseProvider
    from langchain_community.llms import FakeListLLM
except ImportError:

    class BaseProvider:
        pass

    class FakeListLLM:
        pass


FIBONACCI = """\
def fib(n):
    if n < 2: return n
    return fib(n - 1) + fib(n - 2)
"""


class FibonacciCompletionProvider(BaseProvider, FakeListLLM):  # type: ignore[misc, valid-type]

    id = "my_provider"
    name = "My Provider"
    model_id_key = "model"
    models = ["model_a"]

    def __init__(self, **kwargs):
        kwargs["responses"] = ["This fake response will not be used for completion"]
        kwargs["model_id"] = "model_a"
        super().__init__(**kwargs)

    async def generate_inline_completions(self, request):
        raise ValueError("IPython only supports streaming models.")

    async def stream_inline_completions(self, request):
        from jupyter_ai.completions.models import (
            InlineCompletionList,
            InlineCompletionReply,
        )

        assert request.number > 0
        token = f"t{request.number}s0"
        last_line = request.prefix.splitlines()[-1]

        if not FIBONACCI.startswith(last_line):
            return

        yield InlineCompletionReply(
            list=InlineCompletionList(
                items=[
                    {"insertText": "", "isIncomplete": True, "token": token},
                ]
            ),
            reply_to=request.number,
        )

        async for reply in self._stream(
            FIBONACCI[len(last_line) :],
            request.number,
            token,
        ):
            yield reply

    async def _stream(self, sentence, request_number, token, start_with=""):
        from jupyter_ai.completions.models import InlineCompletionStreamChunk

        suggestion = start_with

        for fragment in sentence.split(" "):
            await asyncio.sleep(0.05)
            if suggestion:
                suggestion += " "
            suggestion += fragment
            yield InlineCompletionStreamChunk(
                type="stream",
                response={"insertText": suggestion, "token": token},
                reply_to=request_number,
                done=False,
            )

        # finally, send a message confirming that we are done
        yield InlineCompletionStreamChunk(
            type="stream",
            response={"insertText": suggestion, "token": token},
            reply_to=request_number,
            done=True,
        )


class SlowStartingCompletionProvider(BaseProvider, FakeListLLM):  # type: ignore[misc, valid-type]
    id = "slow_provider"
    name = "Slow Provider"
    model_id_key = "model"
    models = ["model_a"]

    def __init__(self, **kwargs):
        kwargs["responses"] = ["This fake response will be used for completion"]
        kwargs["model_id"] = "model_a"
        sleep(10)
        super().__init__(**kwargs)