"""A conversation held in the model's KV cache, with audio prefilled while the user is still speaking.

The cache always ends right after a complete chat-template segment, so each step only appends:

    begin_turn():  [<bos>system…<|im_end|>\\n]  <|im_start|>user\\n
    feed_audio():  audio tokens, encoded and prefilled as the mic chunks arrive
    generate():    <|im_end|>\\n<|im_start|>assistant\\n<|mixed_start|>  → generated text + speech frames

The model's own generated tokens (including its speech frames) stay in the cache, so the next turn starts from
there instead of re-encoding the whole conversation.
"""

import copy
import time

import numpy as np
import soxr
import torch

MAX_CONTEXT = 16_000  # tokens; past this the conversation is reset at the next turn


class StreamingSession:
    def __init__(self, engine):
        self.engine = engine
        self.model = engine.model
        self.tok = engine.tokenizer
        self.fe = engine.processor.feature_extractor
        self.im_end = self.tok.convert_tokens_to_ids("<|im_end|>")
        self.reset()

    # ---- state -------------------------------------------------------------------------------------------------

    def reset(self):
        self.cache = None
        self.mask = None
        self.key = None  # (mode, system prompt): a change starts a fresh context
        self.needs_close = False  # last generation stopped without <|im_end|>
        self.turn_snapshot = None
        self._end_audio()

    def _end_audio(self):
        self.resampler = None
        self.feature_cache = None
        self.encoder_cache = None
        self.audio_in_s = 0.0
        self.audio_tokens = 0

    @property
    def length(self) -> int:
        return 0 if self.cache is None else self.cache.get_seq_length()

    def snapshot(self):
        return copy.deepcopy((self.cache, self.mask, self.needs_close))

    def restore(self, snap):
        self.cache, self.mask, self.needs_close = snap
        self._end_audio()

    # ---- appending -----------------------------------------------------------------------------------------------

    def _append_embeds(self, embeds: torch.Tensor):
        n = embeds.shape[1]
        if self.cache is None:
            from transformers import DynamicCache

            self.cache = DynamicCache(config=self.model.config.text_config)
            self.mask = torch.ones((1, 0), dtype=torch.long, device=self.model.device)
        self.mask = torch.cat((self.mask, self.mask.new_ones((1, n))), dim=-1)
        position_ids = self.mask.cumsum(-1).sub(1)[:, -n:]
        self.model.model.language_model(
            inputs_embeds=embeds,
            past_key_values=self.cache,
            attention_mask=self.mask,
            position_ids=position_ids,
            use_cache=True,
        )

    def append_text(self, text: str):
        ids = self.tok(text, add_special_tokens=False, return_tensors="pt").input_ids.to(self.model.device)
        self._append_embeds(self.model.get_input_embeddings()(ids))

    def close_prefix(self) -> str:
        """Text that closes the previous assistant turn, if generation left it open."""
        return "<|im_end|>\n" if self.needs_close else "\n" if self.cache is not None else ""

    # ---- a turn --------------------------------------------------------------------------------------------------

    @torch.inference_mode()
    def begin_turn(self, key, system_block: str, keep_history: bool) -> dict:
        """Open a user turn. `system_block` is the rendered `<bos><|im_start|>system…<|im_end|>\\n` prefix."""
        info = {}
        if not keep_history or key != self.key or self.length > MAX_CONTEXT:
            if self.cache is not None:
                info["context_reset"] = "history off" if not keep_history else "settings changed" if key != self.key else "context full"
            self.reset()
            self.key = key
        self._end_audio()
        t = time.perf_counter()
        self.turn_snapshot = self.snapshot()
        info["snapshot_ms"] = round((time.perf_counter() - t) * 1000, 1)
        prefix = self.close_prefix() if self.cache is not None else system_block
        self.append_text(prefix + "<|im_start|>user\n")
        self.needs_close = False
        info["context_tokens"] = self.length
        return info

    @torch.inference_mode()
    def feed_audio(self, pcm: np.ndarray, rate: int, final: bool = False):
        """Resample, extract mel frames and prefill whatever audio tokens they complete."""
        if self.resampler is None:
            self.resampler = soxr.ResampleStream(rate, self.fe.sampling_rate, 1, dtype="float32")
        pcm16 = self.resampler.resample_chunk(np.asarray(pcm, dtype=np.float32), last=final)
        self.audio_in_s += len(pcm16) / self.fe.sampling_rate
        features, self.feature_cache = self.fe.stream(pcm16, self.feature_cache, is_final=final)
        mel = features["input_features"].to(self.model.device)
        if not mel.shape[-1] and not final:
            return
        out = self.model.prefill_audio(
            mel, past_key_values=self.cache, encoder_cache=self.encoder_cache, attention_mask=self.mask, is_final=final
        )
        self.audio_tokens += out.attention_mask.shape[1] - self.mask.shape[1]
        self.cache, self.mask, self.encoder_cache = out.past_key_values, out.attention_mask, out.encoder_cache

    @torch.inference_mode()
    def generate(self, suffix: str, **kwargs):
        """Append `suffix` and generate from the cache; the cache then covers the whole reply."""
        ids = self.tok(suffix, add_special_tokens=False, return_tensors="pt").input_ids.to(self.model.device)
        mask = torch.cat((self.mask, self.mask.new_ones((1, ids.shape[1]))), dim=-1)
        try:
            out = self.model.generate(
                input_ids=ids, attention_mask=mask, past_key_values=self.cache, return_audio_eos=True,
                return_dict_in_generate=True, **kwargs,
            )
        finally:
            # the cache was advanced in place, whether generation finished or was interrupted
            self.mask = torch.ones((1, self.length), dtype=torch.long, device=self.model.device)
        self.needs_close = int(out.sequences[0, -1]) != self.im_end
        return out