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

# Ensure UTF-8 console output on Windows
if sys.platform == "win32":
    try:
        sys.stdout.reconfigure(encoding="utf-8")
        sys.stdin.reconfigure(encoding="utf-8")
    except Exception:
        pass

REPO_ID = "ViuAI/ViuAI-500M"
EOT_ID = 64002

def main():
    print("=" * 70)
    print("🚀  ViuAI Sarus-500M -- 1-Click Interactive Terminal Chat")
    print("=" * 70)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"[*] Compute Device: {device.upper()}")

    # 1. Download or locate local architecture & tokenizer
    base_dir = os.path.dirname(os.path.abspath(__file__))
    code_dir = os.path.join(base_dir, "code")
    
    if os.path.exists(os.path.join(code_dir, "model.py")):
        sys.path.insert(0, code_dir)
    else:
        print("[*] Downloading model architecture from Hugging Face...")
        p_model = hf_hub_download(repo_id=REPO_ID, filename="code/model.py")
        p_cfg = hf_hub_download(repo_id=REPO_ID, filename="code/config.py")
        sys.path.insert(0, os.path.dirname(p_model))

    from model import ViuAI
    from config import ViuAIConfig

    tok_local = os.path.join(base_dir, "tokenizer", "tokenizer.json")
    if not os.path.exists(tok_local):
        tok_local = os.path.join(base_dir, "tokenizer.json")
    if not os.path.exists(tok_local):
        print("[*] Downloading tokenizer...")
        tok_local = hf_hub_download(repo_id=REPO_ID, filename="tokenizer.json")

    tokenizer = Tokenizer.from_file(tok_local)
    cfg = ViuAIConfig(vocab_size=tokenizer.get_vocab_size(), context_length=2048)
    model = ViuAI(cfg).to(device)

    # 2. Download or locate DPO 5k checkpoint
    ckpt_local = os.path.join(base_dir, "dpo_checkpoints", "dpo_5k", "dpo_5k_final.pt")
    if not os.path.exists(ckpt_local):
        print("[*] Downloading DPO-5K Final Weights (~944 MB)...")
        ckpt_local = hf_hub_download(repo_id=REPO_ID, filename="dpo_checkpoints/dpo_5k/dpo_5k_final.pt")

    print(f"[*] Loading weights from {os.path.basename(ckpt_local)}...")
    state = torch.load(ckpt_local, map_location=device, weights_only=True)
    raw_sd = state.get("model", state.get("model_state_dict", state))
    clean_sd = {k.replace("_orig_mod.", "").replace("module.", ""): v for k, v in raw_sd.items()}
    model.load_state_dict(clean_sd, strict=False)
    model.eval()

    print("\n[OK] Model Ready! Type 'exit' to quit.\n" + "-" * 70)

    while True:
        try:
            prompt = input("\n👤 USER >> ").strip()
            if not prompt:
                continue
            if prompt.lower() in ["exit", "quit", "q"]:
                print("\nSession ended.")
                break

            formatted = f"<|user|>\n{prompt}<|endofturn|>\n<|assistant|>\n"
            ids = tokenizer.encode(formatted).ids
            idx = torch.tensor([ids], dtype=torch.long, device=device)

            print("🤖 SARUS >> ", end="", flush=True)

            generated = []
            prev_text = ""

            with torch.no_grad():
                for _ in range(300):
                    idx_cond = idx if idx.size(1) <= cfg.context_length else idx[:, -cfg.context_length:]
                    logits, _ = model(idx_cond)
                    logits = logits[:, -1, :]

                    # Repetition penalty
                    if generated:
                        gen_t = torch.tensor(list(set(generated)), dtype=torch.long, device=device)
                        score = logits[:, gen_t]
                        logits[:, gen_t] = torch.where(score < 0, score * 1.15, score / 1.15)

                    # Top-K & Top-P Sampling
                    sampling_logits = logits / 0.15
                    top_k_val = min(50, sampling_logits.size(-1))
                    kth_vals, _ = torch.topk(sampling_logits, top_k_val)
                    sampling_logits[sampling_logits < kth_vals[:, [-1]]] = float("-inf")

                    sorted_logits, sorted_indices = torch.sort(sampling_logits, descending=True)
                    cum_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
                    remove_mask = cum_probs > 0.90
                    remove_mask[..., 1:] = remove_mask[..., :-1].clone()
                    remove_mask[..., 0] = False
                    indices_to_remove = torch.zeros_like(sampling_logits, dtype=torch.bool).scatter_(1, sorted_indices, remove_mask)
                    sampling_logits[indices_to_remove] = float("-inf")

                    probs = torch.softmax(sampling_logits, dim=-1)
                    idx_next = torch.multinomial(probs, num_samples=1)
                    token_id = idx_next.item()

                    if token_id == EOT_ID:
                        break

                    generated.append(token_id)
                    idx = torch.cat([idx, idx_next], dim=1)

                    curr_text = tokenizer.decode(generated)
                    chunk = curr_text[len(prev_text):]
                    print(chunk, end="", flush=True)
                    prev_text = curr_text

            print()

        except KeyboardInterrupt:
            print("\nSession interrupted.")
            break

if __name__ == "__main__":
    main()