# app.py — CodVa-2 SFT backend (serves custom HTML UI)

import os
import math
import time
import json
import threading

os.environ["OMP_NUM_THREADS"] = "2"
os.environ["MKL_NUM_THREADS"] = "2"
os.environ["OPENBLAS_NUM_THREADS"] = "1"

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download, login, HfApi
from flask import Flask, request, Response, send_file, jsonify
import re

torch.set_num_threads(2)
torch.set_num_interop_threads(1)

HF_TOKEN      = os.environ.get("HF_TOKEN", "")
HF_DATA_REPO  = "Bc-AI/MiniCoder-Instruction-dataset"
HF_MODEL_REPO = os.environ.get("MODEL_REPO", "hugging-science/CodVa-2-SFT")

if HF_TOKEN:
    login(token=HF_TOKEN)

# ─────────────────────────────────────────────────────────────────────────────
#  CONFIG
# ─────────────────────────────────────────────────────────────────────────────
@dataclass
class Config:
    vocab_size:       int   = 50304
    d_model:          int   = 896
    n_layers:         int   = 18
    n_heads:          int   = 14
    n_kv_heads:       int   = 2
    max_len:          int   = 2048
    rope_theta:       float = 500_000.0
    window_size:      int   = 512
    pattern_mult:     float = 2.5
    reason_mult:      float = 0.75
    reason_depth:     int   = 2
    gate_hidden:      int   = 64
    gate_init:        float = 0.0
    diff_lambda_init: float = 0.8
    tie_embeddings:   bool  = True

    @property
    def head_dim(self):
        return self.d_model // self.n_heads

    @property
    def pattern_dim(self):
        return ((int(self.d_model * self.pattern_mult) + 255) // 256) * 256

    @property
    def reason_dim(self):
        return ((int(self.d_model * self.reason_mult) + 63) // 64) * 64

# ─────────────────────────────────────────────────────────────────────────────
#  MODEL
# ─────────────────────────────────────────────────────────────────────────────
class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.w   = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        x32 = x.float()
        return (x32 * torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps) * self.w).to(x.dtype)

def precompute_rope(head_dim, max_len, theta, device):
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
    pos = torch.arange(max_len, device=device, dtype=torch.float32)
    freqs = torch.outer(pos, inv_freq)
    return freqs.cos(), freqs.sin()

def apply_rope(x, cos, sin):
    B, H, L, D = x.shape
    h = D // 2
    c = cos[:L].unsqueeze(0).unsqueeze(0)
    s = sin[:L].unsqueeze(0).unsqueeze(0)
    return torch.cat([x[..., :h] * c - x[..., h:] * s, x[..., h:] * c + x[..., :h] * s], dim=-1)

def repeat_kv(x, n_rep):
    if n_rep == 1: return x
    B, H, L, D = x.shape
    return x.unsqueeze(2).expand(B, H, n_rep, L, D).reshape(B, H * n_rep, L, D)

class DifferentialAttention(nn.Module):
    def __init__(self, cfg, layer_idx, local=False):
        super().__init__()
        self.n_pairs    = cfg.n_heads // 2
        self.n_kv_pairs = max(1, cfg.n_kv_heads // 2)
        self.n_rep      = self.n_pairs // self.n_kv_pairs
        self.head_dim   = cfg.head_dim
        self.local      = local
        self.window     = cfg.window_size
        d               = cfg.d_model
        self.wq       = nn.Linear(d, 2 * self.n_pairs    * self.head_dim, bias=False)
        self.wk       = nn.Linear(d, 2 * self.n_kv_pairs * self.head_dim, bias=False)
        self.wv       = nn.Linear(d,     self.n_kv_pairs * self.head_dim, bias=False)
        self.wo       = nn.Linear(    self.n_pairs * self.head_dim, d,    bias=False)
        self.q_norm   = RMSNorm(self.head_dim)
        self.k_norm   = RMSNorm(self.head_dim)
        self.lambda1  = nn.Parameter(torch.tensor(0.0))
        self.lambda2  = nn.Parameter(torch.tensor(0.0))
        self.out_norm = RMSNorm(self.head_dim)

    def _window_mask(self, L, device):
        idx  = torch.arange(L, device=device)
        dist = idx.unsqueeze(0) - idx.unsqueeze(1)
        mask = (dist > 0) | (dist < -self.window)
        return mask.float().masked_fill(mask, float('-inf'))

    def forward(self, x, cos, sin):
        B, L, _ = x.shape
        hd = self.head_dim
        q_all  = self.wq(x).view(B, L, self.n_pairs,    2, hd).transpose(1, 2)
        k_all  = self.wk(x).view(B, L, self.n_kv_pairs, 2, hd).transpose(1, 2)
        v      = self.wv(x).view(B, L, self.n_kv_pairs,    hd).transpose(1, 2)
        q1, q2 = q_all[..., 0, :], q_all[..., 1, :]
        k1, k2 = k_all[..., 0, :], k_all[..., 1, :]
        q1 = self.q_norm(q1); q2 = self.q_norm(q2)
        k1 = self.k_norm(k1); k2 = self.k_norm(k2)
        q1 = apply_rope(q1, cos, sin); q2 = apply_rope(q2, cos, sin)
        k1 = apply_rope(k1, cos, sin); k2 = apply_rope(k2, cos, sin)
        k1 = repeat_kv(k1, self.n_rep); k2 = repeat_kv(k2, self.n_rep)
        v  = repeat_kv(v,  self.n_rep)
        scale = 1.0 / math.sqrt(hd)
        mask  = (self._window_mask(L, x.device) if self.local
                 else torch.zeros(L, L, device=x.device).masked_fill(
                     ~torch.ones(L, L, device=x.device, dtype=torch.bool).tril(), float('-inf')))
        a1   = F.softmax(torch.matmul(q1, k1.transpose(-2, -1)) * scale + mask, dim=-1)
        a2   = F.softmax(torch.matmul(q2, k2.transpose(-2, -1)) * scale + mask, dim=-1)
        lam  = torch.exp(self.lambda1) - torch.exp(self.lambda2) + 0.5
        out  = torch.matmul(a1 - lam * a2, v)
        out  = self.out_norm(out).transpose(1, 2).contiguous().view(B, L, -1)
        return self.wo(out)

class DualPathFFN(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        d, pd, rd     = cfg.d_model, cfg.pattern_dim, cfg.reason_dim
        self.pat_gate = nn.Linear(d,  pd, bias=False)
        self.pat_up   = nn.Linear(d,  pd, bias=False)
        self.pat_down = nn.Linear(pd, d,  bias=False)
        layers = [nn.Linear(d, rd, bias=False), nn.SiLU()]
        for _ in range(cfg.reason_depth - 1):
            layers += [nn.Linear(rd, rd, bias=False), nn.SiLU()]
        layers.append(nn.Linear(rd, d, bias=False))
        self.reason = nn.Sequential(*layers)
        self.merge  = nn.Parameter(torch.zeros(d))
    def forward(self, x):
        pat = self.pat_down(F.silu(self.pat_gate(x)) * self.pat_up(x))
        w   = torch.sigmoid(self.merge)
        return w * pat + (1.0 - w) * self.reason(x)

class TokenImportanceGate(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.gate_hidden, bias=True),
            nn.SiLU(),
            nn.Linear(cfg.gate_hidden, 1, bias=True),
        )
    def forward(self, x):
        return x * (0.5 + torch.sigmoid(self.net(x)))

class Block(nn.Module):
    def __init__(self, cfg, layer_idx):
        super().__init__()
        self.norm1 = RMSNorm(cfg.d_model)
        self.norm2 = RMSNorm(cfg.d_model)
        self.attn  = DifferentialAttention(cfg, layer_idx, local=(layer_idx % 2 == 0))
        self.ffn   = DualPathFFN(cfg)
    def forward(self, x, cos, sin):
        x = x + self.attn(self.norm1(x), cos, sin)
        x = x + self.ffn(self.norm2(x))
        return x

class CodVa2(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg        = cfg
        self.embed      = nn.Embedding(cfg.vocab_size, cfg.d_model)
        self.importance = TokenImportanceGate(cfg)
        self.blocks     = nn.ModuleList([Block(cfg, i) for i in range(cfg.n_layers)])
        self.final_norm = RMSNorm(cfg.d_model)
        self.register_buffer("rope_cos", torch.zeros(cfg.max_len, cfg.head_dim // 2))
        self.register_buffer("rope_sin", torch.zeros(cfg.max_len, cfg.head_dim // 2))
        self._rope_ready = False

    def _init_rope(self, device):
        c, s = precompute_rope(self.cfg.head_dim, self.cfg.max_len, self.cfg.rope_theta, device)
        self.rope_cos.copy_(c); self.rope_sin.copy_(s)
        self._rope_ready = True

    def forward(self, tokens):
        B, L   = tokens.shape
        device = tokens.device
        if not self._rope_ready:
            self._init_rope(device)
        cos = self.rope_cos[:L]
        sin = self.rope_sin[:L]
        x   = self.importance(self.embed(tokens))
        for block in self.blocks:
            x = block(x, cos, sin)
        return F.linear(self.final_norm(x), self.embed.weight)

# ─────────────────────────────────────────────────────────────────────────────
#  LOAD MODEL + TOKENIZER
# ─────────────────────────────────────────────────────────────────────────────
print("[init] loading tokenizer...")
tok_path  = hf_hub_download(
    repo_id=HF_DATA_REPO,
    filename="codva2_sft_tokenizer.json",
    repo_type="dataset",
    token=HF_TOKEN or None,
)
tokenizer = Tokenizer.from_file(tok_path)

IM_START_ID = tokenizer.token_to_id("<|im_start|>")
IM_END_ID   = tokenizer.token_to_id("<|im_end|>")
EOS_ID      = tokenizer.token_to_id("<|endoftext|>") or 50256
STOP_IDS    = {IM_END_ID, EOS_ID}

MAX_CTX     = 2048   # total context budget for ring display

print("[init] loading model...")
cfg            = Config()
cfg.vocab_size = (tokenizer.get_vocab_size() + 63) // 64 * 64
model          = CodVa2(cfg)

api   = HfApi()
files = list(api.list_repo_files(repo_id=HF_MODEL_REPO, repo_type="model", token=HF_TOKEN or None))
ckpts = sorted(
    [f for f in files if "codva2_sft" in f and f.endswith(".pt")],
    key=lambda x: int(x.split("step")[-1].split(".")[0]) if "step" in x else 0
)

if ckpts:
    print(f"[init] loading: {ckpts[-1]}")
    wpath = hf_hub_download(HF_MODEL_REPO, ckpts[-1], repo_type="model", token=HF_TOKEN or None)
    ckpt  = torch.load(wpath, map_location="cpu", weights_only=False)
    model.load_state_dict(ckpt["model"], strict=True)
else:
    print("[init] WARNING: no SFT weights found")

model.eval()
print(f"[init] ready | {sum(p.numel() for p in model.parameters())/1e6:.1f}M params")

# ─────────────────────────────────────────────────────────────────────────────
#  GENERATION STATE (for stop button)
# ─────────────────────────────────────────────────────────────────────────────
gen_lock    = threading.Lock()
stop_signal = threading.Event()

SYSTEM_PROMPT = "You are CodVa-2, a coding-specialized AI assistant. You write clean, correct, well-commented code. When solving problems, think step by step."

def format_prompt(user_msg: str) -> str:
    return (
        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
        f"<|im_start|>user\n{user_msg}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )

# ─────────────────────────────────────────────────────────────────────────────
#  FLASK APP
# ─────────────────────────────────────────────────────────────────────────────
app = Flask(__name__)

@app.route("/")
def index():
    return send_file("index.html")

@app.route("/stop", methods=["POST"])
def stop():
    stop_signal.set()
    return jsonify({"ok": True})

@app.route("/generate", methods=["POST"])
def generate():
    data        = request.json
    prompt      = (data.get("prompt") or "").strip()
    max_new     = int(data.get("max_new", 256))
    temperature = float(data.get("temperature", 0.3))
    top_p       = float(data.get("top_p", 0.95))
    top_k       = int(data.get("top_k", 50))

    if not prompt:
        return jsonify({"error": "empty prompt"}), 400

    stop_signal.clear()

    def stream():
        full_prompt = format_prompt(prompt)
        enc = tokenizer.encode(full_prompt)
        ids = enc.ids
        if len(ids) > MAX_CTX - max_new:
            ids = ids[-(MAX_CTX - max_new):]

        prompt_len = len(ids)
        x          = torch.tensor([ids], dtype=torch.long)
        generated  = []
        t0         = time.time()

        # Send initial context info
        yield f"data: {json.dumps({'type': 'ctx', 'prompt_tokens': prompt_len, 'max_ctx': MAX_CTX, 'max_new': max_new})}\n\n"

        with torch.no_grad():
            for step in range(max_new):
                if stop_signal.is_set():
                    break

                x_in   = x[:, -MAX_CTX:] if x.size(1) > MAX_CTX else x
                logits = model(x_in)[0, -1, :]

                # Sample
                logits = logits / max(temperature, 1e-5)
                if top_k > 0:
                    k = min(top_k, logits.size(-1))
                    topk_vals, _ = torch.topk(logits, k)
                    logits[logits < topk_vals[-1]] = float('-inf')
                probs = F.softmax(logits, dim=-1)
                sorted_p, sorted_idx = probs.sort(descending=True)
                cumsum = sorted_p.cumsum(0)
                sorted_p[(cumsum - sorted_p) > top_p] = 0.0
                sorted_p = sorted_p / sorted_p.sum().clamp(min=1e-9)
                next_tok = sorted_idx[torch.multinomial(sorted_p, 1)].item()

                generated.append(next_tok)
                x = torch.cat([x, torch.tensor([[next_tok]])], dim=1)

                # Decode full generated so far
                raw = tokenizer.decode(generated, skip_special_tokens=False)

                elapsed = time.time() - t0
                tps     = len(generated) / max(elapsed, 1e-3)
                total_tokens = prompt_len + len(generated)

                payload = {
                    "type":         "token",
                    "raw":          raw,
                    "tok_id":       next_tok,
                    "step":         len(generated),
                    "max_new":      max_new,
                    "tps":          round(tps, 1),
                    "elapsed":      round(elapsed, 1),
                    "total_tokens": total_tokens,
                    "prompt_tokens": prompt_len,
                    "max_ctx":      MAX_CTX,
                }
                yield f"data: {json.dumps(payload)}\n\n"

                if next_tok in STOP_IDS:
                    break

        elapsed = time.time() - t0
        tps     = len(generated) / max(elapsed, 1e-3)
        yield f"data: {json.dumps({'type': 'done', 'tps': round(tps,1), 'elapsed': round(elapsed,1), 'tokens': len(generated)})}\n\n"

    return Response(stream(), mimetype="text/event-stream",
                    headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, threaded=True)