import os
import re
import sys
import glob
import math

import torch
import torch.nn as nn
import torch.nn.functional as F

NUM_THREADS = os.cpu_count() or 4
torch.set_num_threads(NUM_THREADS)
try:
    torch.set_num_interop_threads(max(1, NUM_THREADS // 2))
except RuntimeError:
    pass

DEVICE = torch.device("cpu")
print(f"🧵 Number of CPU threads : {NUM_THREADS}")
TOKEN_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE)

def tokenize(text: str):
    return TOKEN_RE.findall(text.lower())

class Vocab:
    PAD, UNK, BOS, EOS = "<pad>", "<unk>", "<bos>", "<eos>"

    def __init__(self):
        self.stoi = {}
        self.itos = []

    def encode(self, text, add_bos=False, add_eos=False):
        ids = [self.stoi.get(t, self.stoi[Vocab.UNK]) for t in tokenize(text)]
        if add_bos:
            ids = [self.stoi[Vocab.BOS]] + ids
        if add_eos:
            ids = ids + [self.stoi[Vocab.EOS]]
        return ids

    def decode(self, ids):
        toks = [self.itos[i] for i in ids if 0 <= i < len(self.itos)]
        toks = [t for t in toks if t != Vocab.PAD and t != Vocab.BOS]
        out = []
        for t in toks:
            if t == Vocab.EOS:
                break
            out.append(t)
        text = " ".join(out)
        text = re.sub(r"\s+([.,!?;:])", r"\1", text)
        return text

    def __len__(self):
        return len(self.itos)

class CausalSelfAttention(nn.Module):
    def __init__(self, d_model, n_head, dropout):
        super().__init__()
        assert d_model % n_head == 0, "d_model must be evenly divisible by n_head"
        self.n_head = n_head
        self.head_dim = d_model // n_head
        self.qkv = nn.Linear(d_model, 3 * d_model)
        self.proj = nn.Linear(d_model, d_model)
        self.attn_drop = nn.Dropout(dropout)
        self.resid_drop = nn.Dropout(dropout)

    def forward(self, x, attn_mask):
        B, T, C = x.shape
        qkv = self.qkv(x)
        q, k, v = qkv.split(C, dim=2)
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        att = att.masked_fill(attn_mask, float("-inf"))
        att = F.softmax(att, dim=-1)
        att = self.attn_drop(att)
        out = att @ v
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_drop(self.proj(out))

class TransformerBlock(nn.Module):
    def __init__(self, d_model, n_head, d_ff, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_head, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.GELU(),
            nn.Linear(d_ff, d_model),
            nn.Dropout(dropout),
        )

    def forward(self, x, attn_mask):
        x = x + self.attn(self.ln1(x), attn_mask)
        x = x + self.mlp(self.ln2(x))
        return x

class OSW1Model(nn.Module):
    def __init__(self, vocab_size, cfg: dict, pad_id: int):
        super().__init__()
        self.cfg = cfg
        self.pad_id = pad_id
        self.block_size = cfg["block_size"]

        self.tok_emb = nn.Embedding(vocab_size, cfg["d_model"])
        self.pos_emb = nn.Embedding(cfg["block_size"], cfg["d_model"])
        self.drop = nn.Dropout(cfg["dropout"])
        self.blocks = nn.ModuleList([
            TransformerBlock(cfg["d_model"], cfg["n_head"], cfg["d_ff"], cfg["dropout"])
            for _ in range(cfg["n_layer"])
        ])
        self.ln_f = nn.LayerNorm(cfg["d_model"])
        self.head = nn.Linear(cfg["d_model"], vocab_size, bias=False)
        self.head.weight = self.tok_emb.weight  # weight tying

    def forward(self, idx):
        B, T = idx.shape
        pos = torch.arange(T, device=idx.device).unsqueeze(0)
        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))

        mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=idx.device), diagonal=1)
        for block in self.blocks:
            x = block(x, mask)
        x = self.ln_f(x)
        return self.head(x)

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=0.85, top_k=40, eos_id=None):
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits = self(idx_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-5)
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float("-inf")
            probs = F.softmax(logits, dim=-1)
            next_id = torch.multinomial(probs, num_samples=1)
            idx = torch.cat([idx, next_id], dim=1)
            if eos_id is not None and next_id.item() == eos_id:
                break
        return idx

def find_checkpoint():
    candidates = glob.glob("opensoftware_world_osw1_*.pth")
    if not candidates:
        return None
    candidates.sort(key=os.path.getmtime, reverse=True)
    return candidates[0]

def load_checkpoint(path: str):
    print(f"📦 Loading: {path}")
    ckpt = torch.load(path, map_location="cpu")

    cfg = ckpt["config"]
    vocab = Vocab()
    vocab.stoi = ckpt["vocab_stoi"]
    vocab.itos = ckpt["vocab_itos"]
    pad_id = ckpt["pad_id"]

    model = OSW1Model(len(vocab), cfg, pad_id=pad_id).to(DEVICE)
    model.load_state_dict(ckpt["model_state_dict"])
    model.eval()

    param_count = ckpt.get("param_count", sum(p.numel() for p in model.parameters()))
    training_time = ckpt.get("training_time_sec", None)
    final_loss = ckpt.get("final_loss", None)

    print("\n" + "=" * 64)
    print("🧠  OpenSoftware-World OSW1 — LOADED MODEL INFORMATION")
    print("=" * 64)
    print(f"  File                 : {path}")
    print(f"  Vocab size           : {len(vocab):,}")
    print(f"  Number of parameters : {param_count:,}")
    print(f"  d_model / n_layer    : {cfg['d_model']} / {cfg['n_layer']}")
    print(f"  n_head / d_ff        : {cfg['n_head']} / {cfg['d_ff']}")
    print(f"  Context window       : {cfg['block_size']}")
    if training_time is not None:
        print(f"  Training time        : {training_time/60:.2f} minutes")
    if final_loss is not None:
        print(f"  Final training loss  : {final_loss:.4f}")
    print("=" * 64 + "\n")

    return model, vocab, cfg

def chat_loop(model: OSW1Model, vocab: Vocab):
    print("=" * 64)
    print("💬 OSW1 ready! You can start chatting. Type 'exit' to quit.")
    print("=" * 64)

    eos_id = vocab.stoi[Vocab.EOS]
    bos_id = vocab.stoi[Vocab.BOS]

    while True:
        try:
            user_in = input("\nYou: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n👋 Goodbye!")
            break

        if user_in.lower() in ("exit", "quit"):
            print("👋 Goodbye!")
            break
        if not user_in:
            continue

        ids = [bos_id] + vocab.encode(user_in)
        x = torch.tensor([ids], dtype=torch.long)
        out = model.generate(x, max_new_tokens=60, temperature=0.85, top_k=40, eos_id=eos_id)
        answer_ids = out[0, len(ids):].tolist()
        answer = vocab.decode(answer_ids)
        print(f"OSW1: {answer if answer else '(...silence...)'}")

def main():
    if len(sys.argv) > 1:
        ckpt_path = sys.argv[1]
        if not os.path.isfile(ckpt_path):
            print(f"❌ File not found: {ckpt_path}")
            sys.exit(1)
    else:
        ckpt_path = find_checkpoint()
        if ckpt_path is None:
            print(
                "❌ No checkpoint files found in the directory.\n"
                "   Please train a model using 'python train_osw1.py' or\n"
                "   specify a checkpoint file using 'python model_init.py <file_path>'."
            )
            sys.exit(1)

    model, vocab, cfg = load_checkpoint(ckpt_path)
    chat_loop(model, vocab)


if __name__ == "__main__":
    main()