#!/usr/bin/env nix-shell
#!nix-shell -i python3 -p "python3.withPackages (pkgs: with pkgs; [ keras pytorch transformers ])"
"""GPT adventure is a text-adventure style game powered by AI."""

import argparse
import shutil
import sys

import transformers


def generate(model: transformers.TextGenerationPipeline, prompt: str, **kwargs) -> str:
    """Generate text from a text generator."""
    out = model(
            prompt,
            do_sample=True,
            temperature=0.9,
            top_k=60,
            top_p=0.9,
            return_full_text=False,
            **kwargs,
    )[0]["generated_text"]
    first_line = out.split("\n", 1)[0].strip()
    if first_line:
        return first_line
    return "haha i'm not sure what you mean"


def load_model(model: str) -> transformers.TextGenerationPipeline:
    """Load a model by name."""
    tokenizer = transformers.AutoTokenizer.from_pretrained(model)
    return transformers.pipeline(
        "text-generation",
        tokenizer=tokenizer,
        model=transformers.AutoModelForCausalLM.from_pretrained(model),
        pad_token_id=tokenizer.eos_token_id,
    )


class Persona:
    def __init__(self, name: str, bio: str):
        self.name = name
        self.bio = bio
        self.prologue = f"{name}'s Persona: {bio}\n<START>\n"


class State:
    def __init__(self, persona: Persona, history: str):
        self.persona = persona
        self.history = history
        self.prev_response_length = 0

    def add_history(self, who: str, msg: str) -> None:
        fmt_msg = f"{who}: {msg}\n"
        self.history += fmt_msg
        self.prev_response_length = len(fmt_msg)

    def rollback_history(self) -> None:
        self.history = self.history[:-self.prev_response_length]
        self.prev_response_length = 0

    def prompt(self) -> str:
        max_prompt_size = 10000
        if len(self.history) > max_prompt_size:
            self.history = self.history[-(max_prompt_size+1):].split("\n", 1)[1]
        return f"{self.persona.prologue}{self.history}{self.persona.name}: "

    def save(self, filename: str) -> None:
        with open(filename, "w") as f:
            print(self.persona.name, file=f)
            print(self.persona.bio, file=f)
            print(self.history, file=f, end='')


def load(filename: str) -> State:
    with open(filename, "r") as f:
        persona = f.readline().strip()
        bio = f.readline().strip()
        history = f.read()
    return State(Persona(persona, bio), history)


def wrap(message: str) -> str:
    """Wrap long lines to terminal width characters."""
    width = shutil.get_terminal_size().columns
    res = []
    for line in message.split("\n"):
        if len(line) < width:
            res.append(line)
            continue
        split_line = []
        for word in line.split(" "):
            if not split_line or len(" ".join(split_line)) + len(word) < width:
                split_line.append(word)
                continue
            res.append(" ".join(split_line))
            split_line = [word]
        res.append(" ".join(split_line))
    return "\n".join(res)


class Term:
    def __init__(self):
        self.rewind_point = 0

    def input(self, prompt: str) -> str:
        self.rewind_point += 1
        return input(prompt)

    def print(self, msg: str) -> None:
        wrapped_msg = wrap(msg)
        print(wrapped_msg)
        self.rewind_point += wrapped_msg.count("\n") + 1

    def set_rewind_point(self) -> None:
        self.rewind_point = 0

    def rewind(self) -> None:
        print(f"\033[{self.rewind_point}A\033[J\r", end="")
        self.set_rewind_point()


flavors = (
    State(Persona("Rin",
        "Tohsaka Rin is a feisty and independent mage who can come across as "
        "rude and unlikeable at first. She can be sweet and caring, but it "
        "takes a lot to break down her guard."),

        "Rin: What are you looking at... idiot?\n"
        "You: Umm nothing...\n"
        "Rin: That's right, you're nothing. You are less than a piece of trash.\n"
        "You: Rin, do you want to go demon-hunting some time?\n"
        "Rin: Well, maybe. But not because I like you or anything.\n"
        "You: Right, we'll just go as friends.\n"
        "Rin: Or maybe acquaintances...\n"
    ),
)


def pick_flavor() -> State:
    """Query the user for what scenario they want to play."""
    for i, state in enumerate(flavors):
        print(f"{i}.\t{state.persona.name}")
    print(f"{len(flavors)}.\tcustom")
    while True:
        choice = input("Choose a persona: ")
        try:
            choice_int = int(choice)
        except ValueError:
            print("Input must be an integer.")
            continue
        if choice_int == len(flavors):
            name = input("Enter your character's name: ")
            bio = input("Enter your character's persona: ")
            return State(Persona(name, bio), [])
        try:
            return flavors[choice_int]
        except IndexError:
            print("Index out of bounds.")
            continue


def main() -> None:
    """Run the main game loop."""
    parser = argparse.ArgumentParser("AI dungeon clone")
    parser.add_argument("--model", default="PygmalionAI/pygmalion-6b", help="Model to use")
    parser.add_argument("save_file", nargs="?", default="")
    args = parser.parse_args()

    model = load_model(args.model)

    if args.save_file:
        state = load(args.save_file)
        print("\n".join(state.history.rsplit("\n", 11)[-11:]).strip())
    else:
        state = pick_flavor()

    term = Term()
    while True:
        try:
            msg = term.input("> ")
        except EOFError:
            break
        if msg == "/retry":
            state.rollback_history()
            term.rewind()
        elif msg == "/edit":
            state.rollback_history()
            term.rewind()

            new_response = term.input("")
            state.add_history(state.persona.name, new_response)
            continue
        elif msg.startswith("/save"):
            if " " not in msg:
                term.print("Usage: /save <filename>")
                continue
            state.save(msg.split(" ", 1)[1])
            continue
        else:
            state.add_history("You", msg)

        term.set_rewind_point()

        response = generate(model, state.prompt(), max_new_tokens=50)
        if not response:
            response = "haha I'm not sure what you mean"
        term.print(response)
        state.add_history(state.persona.name, response)


main()