import os
import uuid
import time
import multiprocessing
from typing import List, Optional
from collections import deque

from . import llama_cpp


class Llama:
    """High-level Python wrapper for a llama.cpp model."""

    def __init__(
        self,
        model_path: str,
        # NOTE: The following parameters are likely to change in the future.
        n_ctx: int = 512,
        n_parts: int = -1,
        seed: int = 1337,
        f16_kv: bool = False,
        logits_all: bool = False,
        vocab_only: bool = False,
        use_mlock: bool = False,
        embedding: bool = False,
        n_threads: Optional[int] = None,
    ) -> "Llama":
        """Load a llama.cpp model from `model_path`.

        Args:
            model_path: Path to the model.
            n_ctx: Maximum context size.
            n_parts: Number of parts to split the model into. If -1, the number of parts is automatically determined.
            seed: Random seed. 0 for random.
            f16_kv: Use half-precision for key/value cache.
            logits_all: Return logits for all tokens, not just the last token.
            vocab_only: Only load the vocabulary no weights.
            use_mlock: Force the system to keep the model in RAM.
            embedding: Embedding mode only.
            n_threads: Number of threads to use. If None, the number of threads is automatically determined.

        Raises:
            ValueError: If the model path does not exist.

        Returns:
            A Llama instance.
        """
        self.model_path = model_path

        self.params = llama_cpp.llama_context_default_params()
        self.params.n_ctx = n_ctx
        self.params.n_parts = n_parts
        self.params.seed = seed
        self.params.f16_kv = f16_kv
        self.params.logits_all = logits_all
        self.params.vocab_only = vocab_only
        self.params.use_mlock = use_mlock
        self.params.embedding = embedding

        self.last_n = 64
        self.max_chunk_size = n_ctx

        self.n_threads = n_threads or multiprocessing.cpu_count()

        if not os.path.exists(model_path):
            raise ValueError(f"Model path does not exist: {model_path}")

        self.ctx = llama_cpp.llama_init_from_file(
            self.model_path.encode("utf-8"), self.params
        )

    def tokenize(self, text: bytes) -> List[int]:
        """Tokenize a string.

        Args:
            text: The utf-8 encoded string to tokenize.

        Returns:
            A list of tokens.
        """
        n_ctx = llama_cpp.llama_n_ctx(self.ctx)
        tokens = (llama_cpp.llama_token * n_ctx)()
        n_tokens = llama_cpp.llama_tokenize(
            self.ctx,
            text,
            tokens,
            n_ctx,
            True,
        )
        if n_tokens < 0:
            raise RuntimeError(f'Failed to tokenize: text="{text}" n_tokens={n_tokens}')
        return list(tokens[:n_tokens])

    def detokenize(self, tokens: List[int]) -> bytes:
        """Detokenize a list of tokens.

        Args:
            tokens: The list of tokens to detokenize.

        Returns:
            The detokenized string.
        """
        output = b""
        for token in tokens:
            output += llama_cpp.llama_token_to_str(self.ctx, token)
        return output

    def embed(self, text: str):
        """Embed a string.

        Args:
            text: The utf-8 encoded string to embed.

        Returns:
            A list of embeddings.
        """
        tokens = self.tokenize(text.encode("utf-8"))
        self._eval(tokens, 0)
        embeddings = llama_cpp.llama_get_embeddings(self.ctx)
        return embeddings[:llama_cpp.llama_n_embd(self.ctx)]

    def _eval(self, tokens: List[int], n_past):
        rc = llama_cpp.llama_eval(
            self.ctx,
            (llama_cpp.llama_token * len(tokens))(*tokens),
            len(tokens),
            n_past,
            self.n_threads,
        )
        if rc != 0:
            raise RuntimeError(f"Failed to evaluate: {rc}")

    def _sample(self, last_n_tokens, top_p, top_k, temp, repeat_penalty):
        return llama_cpp.llama_sample_top_p_top_k(
            self.ctx,
            (llama_cpp.llama_token * len(last_n_tokens))(*last_n_tokens),
            len(last_n_tokens),
            top_k=top_k,
            top_p=top_p,
            temp=temp,
            repeat_penalty=repeat_penalty,
        )

    def _generate(self, past_tokens, max_tokens, top_p, top_k, temp, repeat_penalty):
        last_n_tokens = deque([0] * self.last_n, maxlen=self.last_n)
        last_n_tokens.extend(past_tokens)
        for i in range(max_tokens):
            token = self._sample(
                last_n_tokens,
                top_p=top_p,
                top_k=top_k,
                temp=temp,
                repeat_penalty=repeat_penalty,
            )
            yield token
            self._eval([token], len(past_tokens) + i)

    def _call(
        self,
        prompt: str,
        suffix: Optional[str] = None,
        max_tokens: int = 16,
        temperature: float = 0.8,
        top_p: float = 0.95,
        logprobs: Optional[int] = None,
        echo: bool = False,
        stop: List[str] = [],
        repeat_penalty: float = 1.1,
        top_k: int = 40,
        stream: bool = False,
    ):
        completion_id = f"cmpl-{str(uuid.uuid4())}"
        created = int(time.time())
        completion_tokens = []
        prompt_tokens = self.tokenize(prompt.encode("utf-8"))

        if len(prompt_tokens) + max_tokens > llama_cpp.llama_n_ctx(self.ctx):
            raise ValueError(
                f"Requested tokens exceed context window of {llama_cpp.llama_n_ctx(self.ctx)}"
            )

        # Process prompt in chunks to avoid running out of memory
        for i in range(0, len(prompt_tokens), self.max_chunk_size):
            chunk = prompt_tokens[i : min(len(prompt_tokens), i + self.max_chunk_size)]
            self._eval(chunk, n_past=i)

        if stop is not None:
            stop = [s.encode("utf-8") for s in stop]

        finish_reason = None
        for token in self._generate(
            prompt_tokens, max_tokens, top_p, top_k, temperature, repeat_penalty
        ):
            if token == llama_cpp.llama_token_eos():
                finish_reason = "stop"
                break
            completion_tokens.append(token)

            text = self.detokenize(completion_tokens)
            any_stop = [s for s in stop if s in text]
            if len(any_stop) > 0:
                first_stop = any_stop[0]
                text = text[: text.index(first_stop)]
                finish_reason = "stop"
                break

            if stream:
                start = len(self.detokenize(completion_tokens[:-1]))
                longest = 0
                for s in stop:
                    for i in range(len(s), 0, -1):
                        if s[-i:] == text[-i:]:
                            if i > longest:
                                longest = i
                            break
                yield {
                    "id": completion_id,
                    "object": "text_completion",
                    "created": created,
                    "model": self.model_path,
                    "choices": [
                        {
                            "text": text[start : len(text) - longest].decode("utf-8"),
                            "index": 0,
                            "logprobs": None,
                            "finish_reason": None,
                        }
                    ],
                }

        if finish_reason is None:
            finish_reason = "length"

        if stream:
            if finish_reason == "stop":
                start = len(self.detokenize(completion_tokens[:-1]))
                text = text[start:].decode("utf-8")
            else:
                text = ""
            yield {
                "id": completion_id,
                "object": "text_completion",
                "created": created,
                "model": self.model_path,
                "choices": [
                    {
                        "text": text,
                        "index": 0,
                        "logprobs": None,
                        "finish_reason": finish_reason,
                    }
                ],
            }
            return

        text = text.decode("utf-8")

        if echo:
            text = prompt + text

        if suffix is not None:
            text = text + suffix

        if logprobs is not None:
            logprobs = llama_cpp.llama_get_logits(
                self.ctx,
            )[:logprobs]

        yield {
            "id": completion_id,
            "object": "text_completion",
            "created": created,
            "model": self.model_path,
            "choices": [
                {
                    "text": text,
                    "index": 0,
                    "logprobs": logprobs,
                    "finish_reason": finish_reason,
                }
            ],
            "usage": {
                "prompt_tokens": len(prompt_tokens),
                "completion_tokens": len(completion_tokens),
                "total_tokens": len(prompt_tokens) + len(completion_tokens),
            },
        }

    def __call__(
        self,
        prompt: str,
        suffix: Optional[str] = None,
        max_tokens: int = 16,
        temperature: float = 0.8,
        top_p: float = 0.95,
        logprobs: Optional[int] = None,
        echo: bool = False,
        stop: List[str] = [],
        repeat_penalty: float = 1.1,
        top_k: int = 40,
        stream: bool = False,
    ):
        """Generate text from a prompt.

        Args:
            prompt: The prompt to generate text from.
            suffix: A suffix to append to the generated text. If None, no suffix is appended.
            max_tokens: The maximum number of tokens to generate.
            temperature: The temperature to use for sampling.
            top_p: The top-p value to use for sampling.
            logprobs: The number of logprobs to return. If None, no logprobs are returned.
            echo: Whether to echo the prompt.
            stop: A list of strings to stop generation when encountered.
            repeat_penalty: The penalty to apply to repeated tokens.
            top_k: The top-k value to use for sampling.
            stream: Whether to stream the results.

        Raises:
            ValueError: If the requested tokens exceed the context window.
            RuntimeError: If the prompt fails to tokenize or the model fails to evaluate the prompt.

        Returns:
            Response object containing the generated text.
        """
        call = self._call(
            prompt=prompt,
            suffix=suffix,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            logprobs=logprobs,
            echo=echo,
            stop=stop,
            repeat_penalty=repeat_penalty,
            top_k=top_k,
            stream=stream,
        )
        if stream:
            return call
        return next(call)

    def __del__(self):
        llama_cpp.llama_free(self.ctx)