import random

from functools import partial
from argparse import ArgumentParser

import torch

from torch.cuda import is_available as cuda_is_available

from colored import fore_rgb, style

from src.nope_gpt.model import NoPEGPT
from nope_gpt.memory import BufferWindowMemory

DEFAULT_SYSTEM_MESSAGE = (
    "You're a helpful AI assistant named NoPE GPT. "
    "Your job is to chat and answer questions. "
    "If you don't know the answer to a question, say that you don't know. "
)

WHITE = (255, 255, 255)


def main():
    parser = ArgumentParser(description="Chat with the instruction-tuned model.")

    parser.add_argument(
        "--checkpoint_path", default="./checkpoints/checkpoint.pt", type=str
    )
    parser.add_argument("--max_tokens", default=2000, type=int)
    parser.add_argument("--colorize_tokens", action="store_true")
    parser.add_argument("--context_length", default=4096, type=int)
    parser.add_argument("--temperature", default=0.7, type=float)
    parser.add_argument("--top_k", default=500, type=int)
    parser.add_argument("--top_p", default=0.9, type=float)
    parser.add_argument("--repeat_penalty", default=0.1, type=float)
    parser.add_argument("--repeat_window", default=50, type=int)
    parser.add_argument("--max_message_history", default=4, type=int)
    parser.add_argument("--device", default="cpu", type=str)
    parser.add_argument("--seed", default=None, type=int)

    args = parser.parse_args()

    if "cuda" in args.device and not cuda_is_available():
        raise RuntimeError("Cuda is not available.")

    if args.seed is not None:
        torch.manual_seed(args.seed)
        random.seed(args.seed)

    checkpoint = torch.load(
        args.checkpoint_path, map_location=args.device, weights_only=False
    )

    tokenizer = checkpoint["tokenizer"]

    model = NoPEGPT(**checkpoint["model_args"])

    model.resize_token_embeddings(tokenizer.vocabulary_size)

    model.add_lora_adapters(**checkpoint["lora_args"])

    model.load_state_dict(checkpoint["model"])

    model.merge_lora_adapters()

    model.to(args.device)

    print("Model checkpoint loaded")

    model.eval()

    system_message = input("Enter a system message: ")

    if not system_message:
        system_message = DEFAULT_SYSTEM_MESSAGE

    system_message = {
        "role": "system",
        "content": system_message,
    }

    memory = BufferWindowMemory(args.max_message_history)

    generate = partial(
        model.generate,
        max_tokens=args.max_tokens,
        context_length=args.context_length,
        temperature=args.temperature,
        top_k=args.top_k,
        top_p=args.top_p,
        repeat_penalty=args.repeat_penalty,
        repeat_window=args.repeat_window,
    )

    while True:
        prompt = input("Enter a prompt: ")

        user_message = {
            "role": "user",
            "content": prompt,
        }

        memory.add_message(user_message)

        messages = [system_message] + memory.get_history()

        tokens = tokenizer.tokenize_prompt(messages)

        prompt = torch.tensor(tokens, dtype=torch.int64, device=args.device)

        response = ""

        for token, probability in generate(prompt):
            token, probability = token.item(), probability.item()

            if token in tokenizer.stop_tokens:
                break

            if args.colorize_tokens:
                intensity = int(probability * 255)

                r, g, b = 255 - intensity, 0, intensity
            else:
                r, g, b = WHITE

            color = fore_rgb(r, g, b)

            out = tokenizer.decode_single_token(token)

            print(f"{color}{out}{style("reset")}", end="", flush=True)

            response += out

        print("\n")

        if "y" not in input("Go again? (yes|no): ").lower():
            break

        assistant_message = {
            "role": "assistant",
            "content": response,
        }

        memory.add_message(assistant_message)


if __name__ == "__main__":
    main()
