import os

import torch
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput

from .bard_core import LM, KVCache, softcap
from .bard_game import DEFAULT_OPENING, StoryEngine, Tale
from .bard_steering import LENGTHS
from .bard_tokenizer import FORMAT_FILE, VOCAB_FILE, BardTokenizer
from .configuration_bard import BardConfig


class BardForStory(PreTrainedModel):
    """bard-small: a small language model that tells interactive tales. Use `begin` and `act` to play one:

        bard = AutoModelForCausalLM.from_pretrained("<repo>", trust_remote_code=True)
        tale = bard.begin("You wake up in a dungeon.")
        print(tale.text, tale.choices)
        tale = bard.act(tale, tale.choices[0])      # or any words of your own
    """

    config_class = BardConfig
    main_input_name = "input_ids"

    def __init__(self, config: BardConfig):
        super().__init__(config)
        self.lm = LM(config.vocab_size, config.d_model, config.n_heads, config.n_layers, config.window,
                     config.global_every, parallel=config.parallel, rope=config.rope)
        self._tokenizer = None
        self._engine = None
        self.post_init()

    def _init_weights(self, module):
        pass

    @property
    def tokenizer(self) -> BardTokenizer:
        if self._tokenizer is None:
            where = getattr(self.config, "_name_or_path", "") or ""
            if os.path.isdir(where):
                self._tokenizer = BardTokenizer.from_dir(where)
            else:
                from huggingface_hub import hf_hub_download
                self._tokenizer = BardTokenizer(hf_hub_download(where, VOCAB_FILE), hf_hub_download(where, FORMAT_FILE))
        return self._tokenizer

    def forward(self, input_ids, number_values=None, **kwargs):
        if number_values is None:
            number_values = torch.zeros_like(input_ids)
        cache = KVCache(self.lm.body)
        h = self.lm.hidden(input_ids, number_values, cache=cache)
        return CausalLMOutput(logits=softcap(torch.nn.functional.linear(h, self.lm.head.weight)))

    @property
    def engine(self) -> StoryEngine:
        if self._engine is None:
            self._engine = StoryEngine(self.lm, self.tokenizer, self.config.vars_weight, self.config.context)
        return self._engine

    def begin(self, opening: str = DEFAULT_OPENING, length: str = "medium", state: dict | None = None,
              on_text=None, seed: int | None = None) -> Tale:
        """Begin a tale. `opening` is how it starts (the model goes on from it); `length` is "short", "medium" or
        "long"; `state` the variables the player carries (default: 100 money). `on_text(chunk)` is called with the
        narration as it is written. Returns a Tale: .text (the passage), .choices, .state, .ended."""
        if length not in LENGTHS:
            raise ValueError(f"length must be one of {list(LENGTHS)}")
        tale = self.engine.new_tale(opening, length, {"money": 100} if state is None else state)
        return self.engine.turn(tale, None, on_text=on_text, seed=seed)

    def act(self, tale: Tale, move: str, on_text=None, seed: int | None = None) -> Tale:
        """Tell the next turn: `move` is one of tale.choices, or the player's own words. The tale is updated in place
        and returned."""
        if tale.ended:
            raise ValueError("this tale has ended")
        return self.engine.turn(tale, move, on_text=on_text, seed=seed)

    def play(self, opening: str = DEFAULT_OPENING, length: str = "medium"):
        """Play a tale in the terminal: pick a choice by its number or type your own move. /short, /medium and /long
        change how much is told, /quit leaves."""
        tale = self.begin(opening, length, on_text=print_chunk)
        while True:
            show_choices(tale)
            if tale.ended:
                print("\nThe end.")
                return tale
            move = self.ask(tale)
            if move is None:
                return tale
            if move:
                print()
                tale = self.act(tale, move, on_text=print_chunk)

    @staticmethod
    def ask(tale):
        """The player's next move: a choice's text, their own words, "" to ask again, or None to leave."""
        try:
            answer = input("\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            return None
        if answer == "/quit":
            return None
        if answer in ("/short", "/medium", "/long"):
            tale.length = answer[1:]
            return ""
        if answer.isdigit() and 1 <= int(answer) <= len(tale.choices):
            return tale.choices[int(answer) - 1]
        return answer


def print_chunk(chunk):
    print(chunk, end="", flush=True)


def show_choices(tale):
    print()
    for number, choice in enumerate(tale.choices, 1):
        print(f"  {number}. {choice}")
    for name, old, new in tale.changes:
        print(f"  [{name}: {old} -> {new}]")