import gradio as gr
import torch, json, os, re
import spaces
from pathlib import Path

HF_TOKEN = os.environ.get("HF_TOKEN")
MODEL_IDS = {
    "PML-12L-O": "K0D3IN/PML-12L-O",
    "PML-22L-O": "K0D3IN/PML-22L-O",
}

COUNTRY_NAMES = [
    "Turkey", "United States", "United Kingdom", "Germany", "France", "Italy",
    "Spain", "Netherlands", "Belgium", "Switzerland", "Austria", "Sweden",
    "Norway", "Denmark", "Finland", "Poland", "Czech Republic", "Hungary",
    "Romania", "Bulgaria", "Greece", "Croatia", "Serbia",
    "Brazil", "Argentina", "Chile", "Colombia", "Mexico", "Peru",
    "India", "China", "Japan", "South Korea", "Singapore", "Indonesia",
    "Philippines", "Thailand", "Vietnam", "Russia", "Ukraine",
    "Australia", "New Zealand", "South Africa", "Egypt", "Morocco", "Nigeria",
    "Israel", "Saudi Arabia", "UAE", "Canada", "Ireland", "Portugal",
]

COUNTRY_CODE = {n: c for n, c in zip(COUNTRY_NAMES, [
    "TR","US","GB","DE","FR","IT","ES","NL","BE","CH","AT","SE","NO","DK","FI",
    "PL","CZ","HU","RO","BG","GR","HR","RS","BR","AR","CL","CO","MX","PE",
    "IN","CN","JP","KR","SG","ID","PH","TH","VN","RU","UA","AU","NZ","ZA",
    "EG","MA","NG","IL","SA","AE","CA","IE","PT",
])}

MODELS = {}

def load_model(name):
    if name not in MODELS:
        from model_v5 import PasswordLLaMA
        from tokenizers import Tokenizer
        import safetensors.torch
        
        device = "cpu"
        repo_id = MODEL_IDS[name]
        from huggingface_hub import hf_hub_download
        
        config_path = hf_hub_download(repo_id=repo_id, filename="config.json", token=HF_TOKEN)
        with open(config_path) as f:
            cfg = json.load(f)
        
        model = PasswordLLaMA(
            vocab_size=cfg.get("vocab_size", 8192),
            n_layer=cfg.get("n_layer", 22),
            n_embd=cfg.get("n_embd", 384),
            n_head=cfg.get("n_head", 6),
            max_seq_len=cfg.get("max_seq_len", 48),
        )
        
        weights_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", token=HF_TOKEN)
        state = safetensors.torch.load_file(weights_path, device=device)
        model.load_state_dict(state, strict=True)
        model.eval()
        
        toker_path = hf_hub_download(repo_id=repo_id, filename="tokenizer.json", token=HF_TOKEN)
        toker = Tokenizer.from_file(toker_path)
        
        MODELS[name] = (model, toker)
    return MODELS[name]

@spaces.GPU
def generate(country, username, pw_len, model_name, num_pw):
    if not username or not username.strip():
        return "Please enter a username"
    cc = COUNTRY_CODE.get(country, "US")
    username = username.strip()
    pw_len = int(pw_len) if pw_len and pw_len > 0 else 0
    
    model, toker = load_model(model_name)
    
    if pw_len > 0:
        prompt = f"[USER:{username}][COUNTRY:{cc}][LEN:{pw_len}]:"
    else:
        prompt = f"[USER:{username}][COUNTRY:{cc}]:"
    
    prefix_ids = toker.encode(prompt).ids
    device = next(model.parameters()).device
    pad_id = toker.token_to_id("<PAD>")
    eos_id = toker.token_to_id("<EOS>")
    
    results = []
    seen = set()
    bs = min(num_pw, 64)
    
    with torch.no_grad():
        while len(results) < num_pw:
            ids = torch.full((bs, 48), pad_id, dtype=torch.long, device=device)
            ids[:, :len(prefix_ids)] = torch.tensor(prefix_ids, device=device)
            cur_len = len(prefix_ids)
            finished = torch.zeros(bs, dtype=torch.bool, device=device)
            max_new = 16 if pw_len == 0 else pw_len + 4
            
            for _ in range(max_new):
                if finished.all():
                    break
                logits = model(ids[:, :cur_len])
                nxt = logits[torch.arange(bs), -1, :] / 0.8
                vals, _ = torch.topk(nxt, 50)
                nxt[nxt < vals[:, -1:]] = float("-inf")
                probs = torch.softmax(nxt, dim=-1)
                nids = torch.multinomial(probs, 1).squeeze(-1)
                finished |= (nids == eos_id)
                nids[finished] = pad_id
                ids[:, cur_len] = nids
                cur_len += 1
                if pw_len > 0 and cur_len - len(prefix_ids) >= pw_len:
                    break
            
            for i in range(bs):
                pw = toker.decode(ids[i].tolist())
                for t in ["<BOS>", "<EOS>", "<PAD>", "<UNK>"]:
                    pw = pw.replace(t, "")
                pw = re.sub(r"\[[A-Z]+:[^\]]*\]", "", pw)
                pw = re.sub(r"^\s*:\s*", "", pw)  # artık tag ayracı
                pw = pw.strip()
                if pw and pw not in seen and len(pw) >= (pw_len if pw_len > 0 else 4):
                    seen.add(pw)
                    results.append(pw)
    
    return "\n".join(results[:num_pw])

with gr.Blocks(title="PML Password Generator") as demo:
    gr.Markdown("# 🔐 PML Password Generator\nGenerate culturally-aware passwords with USER+COUNTRY conditioning.")
    
    with gr.Row():
        with gr.Column(scale=1):
            country = gr.Dropdown(choices=COUNTRY_NAMES, value="Turkey", label="Country")
            username = gr.Textbox(label="Username", placeholder="john_doe")
            length = gr.Slider(0, 30, 0, step=1, label="Length (0 = auto)")
            model_sel = gr.Dropdown(choices=["PML-12L-O", "PML-22L-O"], value="PML-22L-O", label="Model")
            count = gr.Slider(1, 50, 10, step=1, label="How many?")
            btn = gr.Button("🎲 Generate", variant="primary")
        with gr.Column(scale=2):
            out = gr.Textbox(lines=20, label="Passwords", placeholder="Generated passwords...")
    
    btn.click(fn=generate, inputs=[country, username, length, model_sel, count], outputs=out)
    gr.Markdown("---\n### Models\n[PML-12L-O](https://huggingface.co/K0D3IN/PML-12L-O) • [PML-22L-O](https://huggingface.co/K0D3IN/PML-22L-O) • [PML-6L](https://huggingface.co/K0D3IN/PML-6L)\n\n### Support Open Source AI Research\n**Monero (XMR):** `83iqXtvVu28ZiL9bsATMerSgbFFiD1J1jc96CcxJLEnAW3KBmBKedWnUAeLvLvEA9aBiUBpHQJs1iNHYtkTLZbNUEymobSS`\n\n**Bitcoin (BTC):** `bc1qmnlvpukcgl0hsr7nje0x8555mhtxjt80wtmlxm`")

if __name__ == "__main__":
    demo.queue()
    demo.launch(server_name="0.0.0.0", theme=gr.themes.Soft())