#!/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."""

from typing import List, Optional
import argparse

import transformers


class Model:
    def __init__(self, name: str):
        self.name = name
        self.tokenizer = transformers.AutoTokenizer.from_pretrained(name)
        self.model = transformers.pipeline(
            "text-generation",
            tokenizer=self.tokenizer,
            model=transformers.AutoModelForCausalLM.from_pretrained(name),
            pad_token_id=self.tokenizer.eos_token_id,
        )

    def generate(self, prompt: str) -> str:
        """Generate text from a text generator."""
        out = self.model(
                prompt,
                do_sample=True,
                temperature=0.9,
                top_k=60,
                top_p=0.9,
                return_full_text=False,
                max_new_tokens=50,
        )[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 token_length(self, msg: str) -> int:
        return len(self.tokenizer(msg)["input_ids"])


class State:
    def __init__(self, model: Model, character: str, bio: str, history: List[str], filename: Optional[str]):
        self.model = model
        self.character = character
        self.bio = bio
        self.history = history
        self.filename = filename
        self.prologue = f"{character}'s Persona: {bio}\n<START>\n"

    def add_history(self, who: str, msg: str) -> None:
        self.history.append(f"{who}: {msg}")

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

    def prompt(self) -> str:
        return self.prologue + "\n".join(self.history) + f"\n{self.character}: "

    def trim_history(self):
        while self.history and self.model.token_length(self.prompt()) > 1024:
            self.history = self.history[1:]

    def print_context(self) -> None:
        for elem in self.history[-10:]:
            if elem.startswith("You: "):
                print("> " + elem.removeprefix("You: "))
            else:
                print(elem.removeprefix(f"{self.character}: "))

    def autosave(self) -> None:
        if self.filename:
            self.save(self.filename)

    def save(self, filename: str) -> None:
        self.filename = filename
        with open(filename, "w") as f:
            print(self.model.name, file=f)
            print(self.character, file=f)
            print(self.bio, file=f)
            for elem in self.history:
                print(elem, file=f)


def load(filename: str) -> State:
    with open(filename, "r") as f:
        model = f.readline().strip()
        character = f.readline().strip()
        bio = f.readline().strip()
        history = f.read().strip().split("\n")
    return State(Model(model), character, bio, history, filename)


def clear() -> None:
    print(f"\033[2J", flush=True)


def new_character(model: Model) -> State:
    print("Creating a new character.")
    print("You may have to heavily edit (/edit) the AI's first few responses")
    print("in order for it to learn the appropriate tone.")
    name = input("Enter your character's name: ")
    bio = input("Enter a few sentences describing your character: ")
    return State(model, name, bio, [], None)


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

    if args.save_file:
        state = load(args.save_file)
        state.print_context()
    else:
        state = new_character(Model(args.model))

    while True:
        state.autosave()

        try:
            msg = input("> ")
        except EOFError:
            break
        if msg == "/help":
            commands = (
                    "/help",
                    "/retry",
                    "/edit",
                    "/undo",
                    "/save",
            )
            print("Commands:")
            for cmd in commands:
                print(f"\t{cmd}")
        elif msg == "/retry":
            state.rollback_history()
            clear()
            state.print_context()
        elif msg == "/edit":
            state.rollback_history()
            clear()
            state.print_context()

            new_response = input("")
            state.add_history(state.character, new_response)
            continue
        elif msg == "/undo":
            state.rollback_history()
            state.rollback_history()
            clear()
            state.print_context()
            continue
        elif msg.startswith("/save"):
            if " " not in msg:
                print("Usage: /save <filename>")
                continue
            state.save(msg.split(" ", 1)[1])
            continue
        else:
            state.add_history("You", msg)

        state.trim_history()

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


main()