"""Tales: the state of one, and the engine that tells it turn by turn."""
import re
import time
from dataclasses import dataclass, field

import torch

from .bard_rules import ANNOUNCE, apply_changes, clean_input, narration_so_far, parse_turn, state_text
from .bard_sampling import generate
from .bard_steering import LengthControl

DEFAULT_OPENING = "You wake up in a dungeon."
SAMPLING = dict(temperature=0.8, top_k=50, top_p=0.95, rep_penalty=1.1, rep_window=128)
TURN = re.compile(r"<\|start\|>(state|story|player)<\|content\|>(.*?)<\|end\|>", re.S)


@dataclass
class Tale:
    history: str
    state: dict = field(default_factory=dict)
    title: str = ""
    length: str = "medium"
    text: str = ""
    choices: list = field(default_factory=list)
    ended: bool = False
    turns: int = 0
    pending: str | None = None
    changes: list = field(default_factory=list)
    created: float = field(default_factory=time.time)

    def to_dict(self):
        return dict(self.__dict__)

    @classmethod
    def from_dict(cls, data):
        return cls(**data)

    def transcript(self):
        entries = []
        for role, content in TURN.findall(self.history):
            if role == "player":
                entries.append(("player", content))
            elif role == "story":
                narration = content.split("<|vars|>")[0].split("<|choices|>")[0]
                entries.append(("story", narration.strip()))
        return entries


class StoryEngine:
    def __init__(self, lm, tok, vars_weight, context):
        self.lm = lm
        self.tok = tok
        self.vars_weight = vars_weight
        self.context = context

    def new_tale(self, opening, length, state):
        opening = clean_input(opening) or DEFAULT_OPENING
        state = dict(state)
        history = f"<|start|>state<|content|>{state_text(state)}<|end|>"
        return Tale(history=history, state=state, title=opening, length=length, pending=opening)

    def turn(self, tale, move, on_text=None, seed=None, max_new=450, max_parts=20, **sampling):
        if tale.pending is None:
            move = clean_input(move or "")
            if not move:
                raise ValueError("a move is needed: pick one of the choices, or write what you do")
        prompt, opening = self.prompt(tale, move)
        if opening and on_text:
            on_text(opening)

        raw = self.write(tale.length, prompt, opening, on_text, seed, max_new, max_parts, sampling)
        return self.finish(tale, prompt, opening, raw)

    def prompt(self, tale, move):
        if tale.pending is not None:
            return tale.history + "<|start|>story<|content|>" + tale.pending, tale.pending
        player = f"<|start|>player<|content|>{move}<|end|>"
        return tale.history + player + "<|start|>story<|content|>", ""

    @torch.no_grad()
    def write(self, length, prompt, opening, on_text, seed, max_new, max_parts, sampling):
        if seed is None:
            seed = int(time.time() * 1000) % 2**31
        rng = torch.Generator(device=self.lm.head.weight.device).manual_seed(seed)
        control = LengthControl(self.tok, length, self.vars_weight)
        settings = {**SAMPLING, **sampling}

        raw, shown = "", 0
        for _ in range(max_parts):
            control.start_part(raw)
            tokens = []
            sampled = generate(self.lm, self.tok, prompt + raw, max_new, self.context, rng,
                               adjust_logits=control, **settings)
            for token, number in sampled:
                tokens.append((token, number))
                if on_text:
                    shown = self.stream(on_text, raw, tokens, shown)
            raw += self.decode(tokens)
            if len(tokens) < max_new:
                break
        return raw

    def stream(self, on_text, raw, tokens, shown):
        visible = narration_so_far(raw + self.decode(tokens)).rstrip("�")
        if len(visible) > shown:
            on_text(visible[shown:])
            shown = len(visible)
        return shown

    def decode(self, tokens):
        ids, vals = zip(*tokens) if tokens else ((), ())
        return self.tok.decode(ids, vals)

    def finish(self, tale, prompt, opening, raw):
        text, changes, choices, _ = parse_turn(opening + raw)
        before = dict(tale.state)
        apply_changes(tale.state, changes, text)

        tale.history = prompt + raw + "<|end|>" + f"<|start|>state<|content|>{state_text(tale.state)}<|end|>"
        tale.text = text
        tale.choices = choices
        tale.ended = not choices
        tale.turns += 1
        tale.pending = None
        tale.changes = [(name, before.get(name), value) for name, value in tale.state.items()
                        if name in ANNOUNCE and before.get(name) != value]
        return tale