import os
import math
import json
import time
import threading
import torch
import torch.nn as nn
from torch.nn import functional as F

training_state = {
    "running": False,
    "epoch": 0,
    "total_epochs": 0,
    "step": 0,
    "total_steps": 0,
    "loss": None,
    "best_loss": None,
    "loss_history": [],
    "log": [],
    "done": False,
    "error": None,
    "model_name": None,
}

_stop_flag = threading.Event()


class MultiHeadSelfAttention(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        assert n_embd % n_head == 0
        self.n_head = n_head
        self.n_embd = n_embd
        self.head_size = n_embd // n_head
        self.c_attn = nn.Linear(n_embd, 3 * n_embd, bias=False)
        self.c_proj = nn.Linear(n_embd, n_embd, bias=False)
        self.attn_drop = nn.Dropout(dropout)
        self.resid_drop = nn.Dropout(dropout)
        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)))

    def forward(self, x):
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        q = q.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_size).transpose(1, 2)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(self.head_size))
        att = att.masked_fill(self.mask[:T, :T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_drop(att)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_drop(self.c_proj(y))


class FeedForward(nn.Module):
    def __init__(self, n_embd, dropout):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_embd, 4 * n_embd),
            nn.GELU(),
            nn.Linear(4 * n_embd, n_embd),
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)


class Block(nn.Module):
    def __init__(self, n_embd, n_head, block_size, dropout):
        super().__init__()
        self.ln1 = nn.LayerNorm(n_embd)
        self.attn = MultiHeadSelfAttention(n_embd, n_head, block_size, dropout)
        self.ln2 = nn.LayerNorm(n_embd)
        self.ff = FeedForward(n_embd, dropout)

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


class MiniGPT(nn.Module):
    def __init__(self, vocab_size, block_size, n_embd, n_layer, n_head, dropout):
        super().__init__()
        self.block_size = block_size
        self.tok_emb = nn.Embedding(vocab_size, n_embd)
        self.pos_emb = nn.Embedding(block_size, n_embd)
        self.drop = nn.Dropout(dropout)
        self.blocks = nn.Sequential(*[
            Block(n_embd, n_head, block_size, dropout) for _ in range(n_layer)
        ])
        self.ln_f = nn.LayerNorm(n_embd)
        self.head = nn.Linear(n_embd, vocab_size, bias=False)
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
        if isinstance(module, nn.Linear) and module.bias is not None:
            nn.init.zeros_(module.bias)

    def forward(self, idx, targets=None):
        B, T = idx.size()
        pos = torch.arange(T, device=idx.device)
        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
        x = self.blocks(x)
        x = self.ln_f(x)
        logits = self.head(x)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=0.8, top_k=40):
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]
            logits, _ = self(idx_cond)
            logits = logits[:, -1, :] / temperature
            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)
            idx_next = torch.multinomial(probs, num_samples=1)
            idx = torch.cat((idx, idx_next), dim=1)
        return idx


def get_batch(data, block_size, batch_size, device):
    ix = torch.randint(len(data) - block_size, (batch_size,))
    x = torch.stack([data[i:i+block_size] for i in ix])
    y = torch.stack([data[i+1:i+block_size+1] for i in ix])
    return x.to(device), y.to(device)


def start_training(text, config, model_name):
    global training_state, _stop_flag
    _stop_flag.clear()
    training_state = {
        "running": True,
        "epoch": 0,
        "total_epochs": config["epochs"],
        "step": 0,
        "total_steps": 0,
        "loss": None,
        "best_loss": None,
        "loss_history": [],
        "log": [],
        "done": False,
        "error": None,
        "model_name": model_name,
    }

    def run():
        try:
            device = "cpu"
            block_size = config.get("block_size", 128)
            batch_size = config.get("batch_size", 32)
            n_embd = config.get("n_embd", 256)
            n_layer = config.get("n_layer", 4)
            n_head = config.get("n_head", 4)
            dropout = config.get("dropout", 0.1)
            lr = config.get("lr", 3e-4)
            epochs = config.get("epochs", 5)

            chars = sorted(set(text))
            vocab_size = len(chars)
            stoi = {c: i for i, c in enumerate(chars)}
            itos = {i: c for i, c in enumerate(chars)}

            data = torch.tensor([stoi[c] for c in text], dtype=torch.long)
            steps_per_epoch = max(1, len(data) // (block_size * batch_size))
            total_steps = steps_per_epoch * epochs
            training_state["total_steps"] = total_steps

            model = MiniGPT(vocab_size, block_size, n_embd, n_layer, n_head, dropout).to(device)
            param_count = sum(p.numel() for p in model.parameters())
            training_state["log"].append(f"Model ready: {param_count:,} parameters | Vocab: {vocab_size} chars")

            optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
            scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps)

            best_loss = float('inf')
            global_step = 0

            for epoch in range(1, epochs + 1):
                if _stop_flag.is_set():
                    training_state["log"].append("Training stopped by user.")
                    break

                training_state["epoch"] = epoch
                epoch_loss = 0.0

                for step in range(steps_per_epoch):
                    if _stop_flag.is_set():
                        break
                    xb, yb = get_batch(data, block_size, batch_size, device)
                    _, loss = model(xb, yb)
                    optimizer.zero_grad()
                    loss.backward()
                    nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                    optimizer.step()
                    scheduler.step()
                    global_step += 1
                    epoch_loss += loss.item()
                    training_state["step"] = global_step
                    training_state["loss"] = round(loss.item(), 4)

                avg_loss = epoch_loss / steps_per_epoch
                training_state["log"].append(f"Epoch {epoch}/{epochs} — avg loss: {avg_loss:.4f}")
                training_state["loss_history"].append({"epoch": epoch, "loss": round(avg_loss, 4)})

                if avg_loss < best_loss:
                    best_loss = avg_loss
                    training_state["best_loss"] = round(best_loss, 4)

            save_dir = os.path.join(os.path.dirname(__file__), "models", "trained")
            os.makedirs(save_dir, exist_ok=True)
            save_path = os.path.join(save_dir, f"{model_name}.pt")

            torch.save({
                "model_state": model.state_dict(),
                "config": {
                    "vocab_size": vocab_size,
                    "block_size": block_size,
                    "n_embd": n_embd,
                    "n_layer": n_layer,
                    "n_head": n_head,
                    "dropout": dropout,
                },
                "stoi": stoi,
                "itos": itos,
                "model_name": model_name,
            }, save_path)

            training_state["log"].append(f"Model saved: models/trained/{model_name}.pt")
            training_state["log"].append(f"Best loss: {best_loss:.4f}")
            training_state["done"] = True
            training_state["running"] = False

        except Exception as e:
            training_state["error"] = str(e)
            training_state["running"] = False
            training_state["done"] = True

    t = threading.Thread(target=run, daemon=True)
    t.start()


def stop_training():
    _stop_flag.set()


def get_trained_models():
    save_dir = os.path.join(os.path.dirname(__file__), "models", "trained")
    if not os.path.exists(save_dir):
        return []
    results = []
    for f in os.listdir(save_dir):
        if f.endswith(".pt"):
            results.append({"name": f, "type": "pt"})
        elif os.path.isdir(os.path.join(save_dir, f)) and os.path.exists(os.path.join(save_dir, f, "config.json")):
            results.append({"name": f, "type": "hf"})
    return results


# ── MiniGPT inference ─────────────────────────────────────────────────────────

_inference_cache = {}


def load_trained_model(model_name):
    if model_name in _inference_cache:
        return _inference_cache[model_name]
    save_dir = os.path.join(os.path.dirname(__file__), "models", "trained")
    path = os.path.join(save_dir, model_name)
    checkpoint = torch.load(path, map_location="cpu", weights_only=False)
    cfg = checkpoint["config"]
    model = MiniGPT(
        vocab_size=cfg["vocab_size"],
        block_size=cfg["block_size"],
        n_embd=cfg["n_embd"],
        n_layer=cfg["n_layer"],
        n_head=cfg["n_head"],
        dropout=0.0,
    )
    model.load_state_dict(checkpoint["model_state"])
    model.eval()
    _inference_cache[model_name] = {
        "model": model,
        "stoi": checkpoint["stoi"],
        "itos": checkpoint["itos"],
        "block_size": cfg["block_size"],
    }
    return _inference_cache[model_name]


def generate_text(model_name, prompt, max_tokens=200, temperature=0.8, top_k=40, think_mode=False):
    mc = load_trained_model(model_name)
    model = mc["model"]
    stoi = mc["stoi"]
    itos = mc["itos"]
    block_size = mc["block_size"]

    formatted = f"User: {prompt}\n<think>" if think_mode else prompt
    encoded = [stoi.get(c, 0) for c in formatted]
    idx = torch.tensor([encoded], dtype=torch.long)
    out = model.generate(idx, max_new_tokens=max_tokens, temperature=temperature, top_k=top_k)
    tokens = out[0].tolist()
    full_output = "".join(itos.get(i, "") for i in tokens)
    generated = full_output[len(formatted):]

    if think_mode:
        if "</think>" in generated:
            parts = generated.split("</think>", 1)
            return {"think_block": parts[0].strip(), "response": parts[1].strip(), "thinking": True}
        return {"think_block": generated.strip(), "response": generated.strip(), "thinking": True}
    return {"response": generated, "thinking": False}


# ── ExocoreV1 inference ───────────────────────────────────────────────────────

_exocore_cache = {}


def is_exocore_pt(model_name):
    save_dir = os.path.join(os.path.dirname(__file__), "models", "trained")
    meta_path = os.path.join(save_dir, model_name + ".meta")
    if os.path.exists(meta_path):
        try:
            with open(meta_path, "r") as f:
                meta = json.load(f)
            return meta.get("exocore_type") in ("qwen3", "exocoreV1")
        except Exception:
            pass
    path = os.path.join(save_dir, model_name)
    if os.path.exists(path) and path.endswith(".pt"):
        try:
            ckpt = torch.load(path, map_location="cpu", weights_only=False)
            return ckpt.get("exocore_type") in ("qwen3", "exocoreV1")
        except Exception:
            pass
    return False


def load_exocore_model(model_name):
    import gc
    if model_name in _exocore_cache:
        return _exocore_cache[model_name]
    from exocore_model import ExocoreLM
    save_dir = os.path.join(os.path.dirname(__file__), "models", "trained")
    path = os.path.join(save_dir, model_name)
    ckpt = torch.load(path, map_location="cpu", weights_only=False)
    cfg = ckpt["config"]
    tok_json = ckpt["tokenizer_json"]
    model = ExocoreLM(cfg)
    model.load_state_dict(ckpt["model_state"], strict=True)
    del ckpt
    gc.collect()
    model.eval()
    from tokenizers import Tokenizer
    tok = Tokenizer.from_str(tok_json)
    _exocore_cache[model_name] = {"model": model, "tokenizer": tok, "config": cfg, "name": "ExocoreV1"}
    return _exocore_cache[model_name]


def generate_exocore_stream(model_name, messages, max_tokens=512, temperature=0.7,
                             think_mode=False, deep_think=False, search_context=None):
    from exocore_model import build_chat_prompt, IM_END, EOT

    mc = load_exocore_model(model_name)
    tok = mc["tokenizer"]
    model = mc["model"]
    cfg = mc["config"]
    max_pos = min(cfg.get("max_position_embeddings", 8192), 4096)

    def sample_next(ids):
        ctx = ids[:, -min(ids.shape[1], max_pos):]
        with torch.no_grad():
            logits, _ = model(ctx)
        logits = logits[:, -1, :].float() / max(temperature, 1e-5)
        v, _ = torch.topk(logits, min(50, logits.size(-1)))
        logits[logits < v[:, [-1]]] = float("-inf")
        probs = torch.softmax(logits, dim=-1)
        return torch.multinomial(probs, 1)

    def encode(prompt):
        enc = tok.encode(prompt)
        return torch.tensor([enc.ids], dtype=torch.long)

    def collect_tokens(ids, max_t, stop_ids=(IM_END, EOT)):
        out_ids = []
        for _ in range(max_t):
            nxt = sample_next(ids)
            tid = nxt.item()
            ids = torch.cat([ids, nxt], dim=1)
            if tid in stop_ids:
                break
            out_ids.append(tid)
        return tok.decode(out_ids), ids

    if deep_think:
        think1_prompt = build_chat_prompt(messages, think=True, search_context=search_context)
        think1_raw, _ = collect_tokens(encode(think1_prompt), min(max_tokens * 2, 1024))
        think1 = think1_raw.replace("<think>", "").replace("</think>", "").strip()

        think2_prompt = build_chat_prompt(messages, think=True, prior_thinking=think1, search_context=search_context)
        think2_raw, _ = collect_tokens(encode(think2_prompt), min(max_tokens, 512))
        think2 = think2_raw.replace("<think>", "").replace("</think>", "").strip()

        combined = f"[Pass 1]\n{think1}\n\n[Pass 2]\n{think2}"
        yield ("think_block", combined)

        answer_prompt = build_chat_prompt(messages, think=False, prior_thinking=combined, search_context=search_context)
        ids = encode(answer_prompt)
        for _ in range(max_tokens):
            nxt = sample_next(ids)
            tid = nxt.item()
            ids = torch.cat([ids, nxt], dim=1)
            if tid in (IM_END, EOT):
                break
            yield ("token", tok.decode([tid]))

    else:
        prompt = build_chat_prompt(messages, think=think_mode, search_context=search_context)
        ids = encode(prompt)
        buf = ""
        in_think = False
        think_done = not think_mode
        think_buf = ""
        think_sent = False

        for _ in range(max_tokens):
            nxt = sample_next(ids)
            tid = nxt.item()
            ids = torch.cat([ids, nxt], dim=1)
            if tid in (IM_END, EOT):
                break

            buf += tok.decode([tid])

            if not think_done:
                if not in_think and "<think>" in buf:
                    buf = buf.split("<think>", 1)[1]
                    in_think = True
                if in_think:
                    if "</think>" in buf:
                        ttext, buf = buf.split("</think>", 1)
                        think_buf += ttext
                        think_done = True
                        in_think = False
                        if not think_sent:
                            yield ("think_block", think_buf.strip())
                            think_sent = True
                    else:
                        think_buf += buf
                        buf = ""
                    continue

            if not in_think and buf:
                yield ("token", buf)
                buf = ""

        if buf and not in_think:
            yield ("token", buf)