from typing import List, Optional
import weakref
import torch
from llama_cpp import Llama, LlamaState, llama_token_to_str  # type: ignore

from langdash.llm import LLM
from langdash.llm_session import LLMGenerationSessionForRawText
import langdash.sampling as sampling
from ._tokenizer.bytes_dict_tokenizer import BytesDictTokenizer
from ._mixins.tensor_based_infer_mixin import TensorBasedInferWithSessionMixin


class LlamaWrapper:
  model: Llama
  tokenizer: BytesDictTokenizer
  last_called_session: Optional[weakref.ref]
  eos_token: int
  bos_token: int

  def __init__(self, *args, **kwargs):
    self.model = Llama(*args, **kwargs)
    mapping = [
      llama_token_to_str(self.model.ctx, tokid)
      for tokid in range(self.model.n_vocab())
    ]
    self.tokenizer = BytesDictTokenizer(
      lambda text, **_k: self.model.
      tokenize(text.encode("utf-8"), add_bos=False),
      lambda tokens, **_k: self.model.detokenize(tokens).decode("utf-8"),
      mapping
    )
    self.last_called_session = None
    self.eos_token = Llama.token_eos()
    self.bos_token = Llama.token_bos()

  def eval(self, tokens: List[int]) -> torch.Tensor:
    self.model.eval(tokens)
    return torch.from_numpy(self.model._scores[-1, :])

  def enter_session(self, session: "LlamaCppSession"):
    if self.last_called_session is None:
      self.last_called_session = weakref.ref(session)
      return
    last_called_session = self.last_called_session()
    if session == last_called_session:
      return
    elif last_called_session is not None:
      last_called_session._logits = self.load_logits_from_llama()
      last_called_session._state = self.model.save_state()
    if session._state is not None:
      self.model.load_state(session._state)
    self.last_called_session = weakref.ref(session)

  def load_logits_from_llama(self) -> torch.Tensor:
    return torch.Tensor(self.model.eval_logits[-1])

  def clone_state(self, session: "LlamaCppSession") -> LlamaState:
    self.enter_session(session)
    return self.model.save_state()

  def set_state(self, session: "LlamaCppSession", state: Optional[LlamaState]):
    self.enter_session(session)
    self.model.reset()
    if state is not None:
      self.model.load_state(state)


class LlamaCppSession(
  TensorBasedInferWithSessionMixin,
  LLMGenerationSessionForRawText["LlamaCppModel", LlamaState, torch.Tensor]
):
  """
  Session for llama.cpp model.
  """

  def __init__(self, *args, **kwargs):
    super().__init__(*args, **kwargs)

    def load_model(llm: LlamaCppModel):
      return LlamaWrapper(model_path=llm._model_path, **llm._model_kwargs)

    self._model = self._ld._get_model_internal(self._llm, load_model)
    self._logits = None
    self._state = None
    self._next_token = None

  def _eval(self, token: int):
    return self._model.eval([token])

  def _eval_mult(self, tokens: List[int]):
    return self._model.eval(tokens)

  def _next_token_logits_raw(self):
    self._model.enter_session(self)
    if self._next_token is None:
      if self._logits is None:
        raise ValueError("cannot predict next probability for empty input")
      logits = self._logits
    else:
      logits = self._eval(self._next_token[0])
    return logits

  def next_token_logits(self) -> List[float]:
    return self._next_token_logits_raw().tolist()

  def next_token_probs(self) -> List[float]:
    return sampling.logits_to_probs(self._next_token_logits_raw()).tolist()

  def set_state(self, state: Optional[LlamaState]):
    self._model.set_state(self, state)
    if state is None:
      self._logits = None
    else:
      self._logits = self._model.load_logits_from_llama()

  def clone_state(self) -> LlamaState:
    return self._model.clone_state(self)

  def tokenize(self, text: str, add_special_tokens: bool = False) -> List[int]:
    return self._model.tokenizer.encode(
      text, add_special_tokens=add_special_tokens
    )

  def decode(self, tokens: List[int]) -> str:
    return self._model.tokenizer.decode(tokens)

  def _on_first_inject(self):
    self._model.model.reset()
    self._eval(self._model.bos_token)

  # Wrapper for public functions to flush the old session states

  def inject(self, *a, **k):
    self._model.enter_session(self)
    return LLMGenerationSessionForRawText.inject(self, *a, **k)

  def flush_token(self, *a, **k):
    self._model.enter_session(self)
    return LLMGenerationSessionForRawText.inject(self, *a, **k)


class LlamaCppModel(LLM[LlamaCppSession]):
  """
  llama.cpp model.
  """

  Session = LlamaCppSession

  def __init__(self, model_path: str, **model_kwargs):
    self._model_path = model_path
    self._model_kwargs = model_kwargs

  def session(self, **kwargs):
    return LlamaCppSession(self, **kwargs)
