import os
import sys
import torch
import torch.nn.functional as F

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from model.model import ModelConfig, Transformer

def hinglish_interactive_chat():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print("="*60)
    print("  Romanized Hindi (Hinglish) Foundational Model CLI")
    print("="*60)
    
    # Load 125M Model
    config = ModelConfig.get_125m(vocab_size=16384)
    model = Transformer(config).to(device)
    
    checkpoint_path = "/home/adminuser/foundational_model/checkpoints/model_125m_final.pt"
    if os.path.exists(checkpoint_path):
        model.load_state_dict(torch.load(checkpoint_path, map_location=device))
        print(f"Loaded trained weights from {checkpoint_path}")
    else:
        print("Notice: No trained checkpoint found yet. Running in demo mode with randomly initialized model.")

    model.eval()
    print("\nType your prompt in Romanized Hindi (e.g., 'kya kar rahe ho?')")
    print("Type 'exit' to quit.\n")

    while True:
        try:
            prompt = input("User (Hinglish) > ")
            if prompt.strip().lower() == "exit":
                break
            if not prompt.strip():
                continue

            # Tokenize & Generate
            print("Model (Hinglish) > ", end="", flush=True)
            # Dummy generation representation
            fake_response = "main ek foundational AI model hoon aur aapki madad karne ke liye taiyaar hoon!"
            print(fake_response + "\n")

        except KeyboardInterrupt:
            break

if __name__ == "__main__":
    hinglish_interactive_chat()