from config import (
    DEVICE,
    CHECKPOINT_DIR,
)

from data.tokenizer import Tokenizer
from model.builder_model import BuilderModel
from utils.checkpoint import latest_checkpoint, load_checkpoint


def main():
    if DEVICE == "cuda":
        print("Cuda Mode: true")
    else:
        print("Cuda Mode: false")

    print("Loading tokenizer...")

    tokenizer = Tokenizer()
    tokenizer.load("data")

    print("Loading model...")

    model = BuilderModel(tokenizer).to(DEVICE)

    checkpoint = latest_checkpoint(CHECKPOINT_DIR)

    if checkpoint is None:
        print("No checkpoint found!")
        return

    print(f"Loading checkpoint: {checkpoint}")

    load_checkpoint(
        checkpoint,
        model,
        optimizer=None,
        map_location=DEVICE,
    )

    model.eval()

    print("\nMinecraft AI Builder")
    print("Type 'exit' to quit.\n")

    while True:

        prompt = input("> ")

        if prompt.lower() in ("exit", "quit"):
            break

        blocks, token_count = model.generate(prompt)

        print("\nGenerated Structure:\n")

        if len(blocks) == 0:
            print("(No blocks generated)")
            continue

        for block in blocks:
            print(
                f"{block['block']:20}"
                f" x={block['x']:4}"
                f" y={block['y']:4}"
                f" z={block['z']:4}"
            )

        print(f"\nTotal blocks: {len(blocks)}")
        print(f"Generated tokens: {token_count}\n")


if __name__ == "__main__":
    main()