# ==============================================================================
# 🚀 ViuTranslate — Interactive Translation CLI
# ==============================================================================
# Run: python inference.py
# Type any sentence in English or Hindi to get real-time neural translation.
# ==============================================================================

import os
import sys
import torch
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download

# UTF-8 encoding
if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")

cur_dir = os.path.dirname(os.path.abspath(__file__)) if "__file__" in locals() else os.getcwd()
if cur_dir not in sys.path:
    sys.path.insert(0, cur_dir)

from model import ViuAI
from config import ViuAIConfig

REPO_ID = "ViuAI/ViuTranslate"
EOT_ID = 64002

def load_model():
    print("=" * 75)
    print(f"🚀 Initializing ViuTranslate-500M from {REPO_ID}...")
    print("=" * 75)
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"• Hardware Device: {device.type.upper()}")
    
    # Checkpoint candidates
    ckpt_candidates = [
        "viutranslate_final.pt",
        "checkpoints/viutranslate_final.pt",
        os.path.join(cur_dir, "viutranslate_final.pt")
    ]
    ckpt_path = None
    for c in ckpt_candidates:
        if os.path.exists(c):
            ckpt_path = c
            break
            
    if ckpt_path is None:
        print(f"📥 Downloading viutranslate_final.pt from {REPO_ID}...")
        ckpt_path = hf_hub_download(repo_id=REPO_ID, filename="viutranslate_final.pt")
        
    tok_candidates = [
        "tokenizer.json",
        os.path.join(cur_dir, "tokenizer.json")
    ]
    tok_path = None
    for tc in tok_candidates:
        if os.path.exists(tc):
            tok_path = tc
            break
            
    if tok_path is None:
        print(f"📥 Downloading tokenizer.json from {REPO_ID}...")
        tok_path = hf_hub_download(repo_id=REPO_ID, filename="tokenizer.json")
        
    tokenizer = Tokenizer.from_file(tok_path)
    
    cfg = ViuAIConfig(vocab_size=64003, context_length=2048)
    model = ViuAI(cfg).to(device)
    
    state = torch.load(ckpt_path, map_location=device, weights_only=False)
    weights = state.get("model_state_dict", state)
    model.load_state_dict(weights, strict=False)
    model.eval()
    
    print("✅ ViuTranslate Engine loaded and ready for inference!\n")
    return model, tokenizer, device

@torch.no_grad()
def translate(model, tokenizer, device, text: str, mode: str = "direct") -> str:
    text = text.strip()
    if mode == "direct":
        prompt = f"<|user|>\n{text}<|endofturn|>\n<|assistant|>\n"
    elif mode == "to_hi":
        prompt = f"<|user|>\nTranslate to Hindi: '{text}'<|endofturn|>\n<|assistant|>\n"
    elif mode == "to_en":
        prompt = f"<|user|>\nTranslate to English: '{text}'<|endofturn|>\n<|assistant|>\n"
    else:
        prompt = f"<|user|>\n{text}<|endofturn|>\n<|assistant|>\n"

    input_ids = torch.tensor([tokenizer.encode(prompt).ids], dtype=torch.long, device=device)
    prompt_len = input_ids.shape[1]
    
    out = model.generate(
        input_ids,
        max_new_tokens=150,
        temperature=0.2,
        top_p=0.9,
        repetition_penalty=1.15,
        eos_token_id=EOT_ID
    )
    gen_tokens = out[0][prompt_len:].tolist()
    if EOT_ID in gen_tokens:
        gen_tokens = gen_tokens[:gen_tokens.index(EOT_ID)]
        
    return tokenizer.decode(gen_tokens).strip()

def interactive_loop():
    model, tokenizer, device = load_model()
    print("💡 Enter text to translate (Google Translate style). Type 'exit' or 'quit' to stop.\n")
    
    while True:
        try:
            inp = input("📝 [Input]: ").strip()
            if not inp:
                continue
            if inp.lower() in ["exit", "quit", "q"]:
                print("👋 Exiting ViuTranslate.")
                break
                
            out = translate(model, tokenizer, device, inp, mode="direct")
            print(f"🌐 [ViuTranslate]: {out}\n")
        except (KeyboardInterrupt, EOFError):
            print("\n👋 Exiting ViuTranslate.")
            break

if __name__ == "__main__":
    interactive_loop()