import gradio as gr
import torch
import sys
import os
from huggingface_hub import hf_hub_download

sys.path.insert(0, os.path.dirname(__file__))

from src.model import Nexus
from src.config import NexusConfig
from tokenizers import Tokenizer

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Loading on {device}")

config = NexusConfig()
model = Nexus(config)

REPO = "JustScriptzz/nexus-smAll-v1"

weights_local = os.path.join(os.path.dirname(__file__), "weights", "nexus_instruct.pt")
if os.path.exists(weights_local):
    weights_path = weights_local
else:
    print("Downloading weights from HuggingFace...")
    weights_path = hf_hub_download(repo_id=REPO, filename="weights/nexus_instruct.pt")

checkpoint = torch.load(weights_path, map_location=device, weights_only=False)
model.load_state_dict(checkpoint["model_state_dict"])
model = model.to(device)
model.eval()
print("Model loaded")

tokenizer_local = os.path.join(os.path.dirname(__file__), "data", "tokenizer.json")
if os.path.exists(tokenizer_local):
    tokenizer_path = tokenizer_local
else:
    tokenizer_path = hf_hub_download(repo_id=REPO, filename="data/tokenizer.json")
tokenizer = Tokenizer.from_file(tokenizer_path)

bos_id = tokenizer.token_to_id("<bos>") or 1
eos_id = tokenizer.token_to_id("<eos>") or 2


def chat(message, history):
    history_ids = [bos_id]
    for user_msg, bot_msg in history:
        if user_msg:
            history_ids.extend(tokenizer.encode(f"User: {user_msg}\nAssistant:").ids)
        if bot_msg:
            history_ids.extend(tokenizer.encode(f" {bot_msg}").ids + [eos_id])

    history_ids.extend(tokenizer.encode(f"User: {message}\nAssistant:").ids)

    input_tensor = torch.tensor([history_ids[-config.max_seq_len:]], dtype=torch.long, device=device)

    with torch.no_grad():
        for _ in range(128):
            seq_len = input_tensor.shape[1]
            if seq_len > config.max_seq_len:
                input_tensor = input_tensor[:, -config.max_seq_len:]

            logits = model(input_tensor, 0)
            logits = logits[:, -1, :] / 0.2

            probs = torch.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)

            input_tensor = torch.cat([input_tensor, next_token], dim=-1)

            if next_token.item() == eos_id:
                break

    response_ids = input_tensor[0].tolist()
    response_ids = response_ids[-(input_tensor.shape[1] - len(history_ids)):]
    
    response = tokenizer.decode(response_ids)
    response = response.split("<eos>")[0].split("User:")[0].replace("Assistant:", "").strip()

    if len(response) < 2:
        response = "[no response]"

    return response


css = """
footer {visibility: hidden}
.message {font-size: 14px}
"""

demo = gr.ChatInterface(
    fn=chat,
    title="Nexus SmAll v1",
    description="A 89.8M parameter transformer trained from scratch. May produce incoherent outputs — it's a tiny model!",
    css=css,
    theme="soft",
)

if __name__ == "__main__":
    demo.launch()