import os
import threading
import time
import re
import discord
import gradio as gr
import torch
import torch.nn as nn
import torch.nn.functional as F
from discord.ext import commands
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer

DISCORD_TOKEN = os.getenv("DISCORD_TOKEN")

# ============================================================
# AI MODEL SETUP (Cosmos T2-Accelerate)
# ============================================================

TOKENIZER_NAME = "Qwen/Qwen2.5-0.5B"
DEFAULT_SYSTEM_PROMPT = "Enable thinking features: INTUITION"

MODELS = {
    "beta2": {
        "full_name": "Cosmos T2-Accelerate Beta 2",
        "repo": "wop/Cosmos-T2-Accelerate-Beta2",
        "short": "Beta 2",
        "desc": "~9.96M params · latest, better quality",
        "checkpoints": {
            "latest": "final.pt",
            "best":   "best.pt",
        },
    },
    "preview": {
        "full_name": "Cosmos T2-Accelerate Preview",
        "repo": "wop/Cosmos-T2-Accelerate-Preview",
        "short": "Preview",
        "desc": "~9.96M params · better quality",
        "checkpoints": {
            "latest": "Cosmos-T2-Accelerate-Preview.pt",
            "best":   "Cosmos-T2-Accelerate-Preview.best.pt",
        },
    },
    "beta": {
        "full_name": "Cosmos T2-Accelerate Beta",
        "repo": "wop/Cosmos-T2-Accelerate-beta",
        "short": "Beta",
        "desc": "~5.03M params · fastest, roughest",
        "checkpoints": {
            "latest": "Cosmos-T2-Accelerate-beta.pt",
            "best":   "Cosmos-T2-Accelerate-beta.best.pt",
        },
    },
}

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE  = torch.float16 if DEVICE == "cuda" else torch.float32

# ---------------- MODEL ARCHITECTURE ----------------

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps
    def forward(self, x):
        rms = x.pow(2).mean(dim=-1, keepdim=True)
        x = x * torch.rsqrt(rms + self.eps)
        return x * self.weight

def rotate_half(x):
    x1 = x[..., ::2]
    x2 = x[..., 1::2]
    return torch.stack((-x2, x1), dim=-1).flatten(-2)

def apply_rope(q, k, cos, sin):
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)

class GQAAttention(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads, rope_base=10000, dropout=0.0):
        super().__init__()
        assert d_model % n_heads == 0
        assert n_heads % n_kv_heads == 0
        self.n_heads    = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim   = d_model // n_heads
        self.dropout    = dropout
        self.q_proj = nn.Linear(d_model, n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(d_model, n_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, rope_cos, rope_sin, past_kv=None, use_cache=False):
        batch, seq_len, _ = x.shape
        q = self.q_proj(x).view(batch, seq_len, self.n_heads,    self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
        q, k = apply_rope(q, k, rope_cos, rope_sin)
        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)
        present_kv = (k, v) if use_cache else None
        if self.n_kv_heads != self.n_heads:
            repeat = self.n_heads // self.n_kv_heads
            k = k.repeat_interleave(repeat, dim=1)
            v = v.repeat_interleave(repeat, dim=1)
        out = F.scaled_dot_product_attention(
            q, k, v, is_causal=(past_kv is None),
            dropout_p=self.dropout if self.training else 0.0
        )
        out = out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
        out = self.o_proj(out)
        return (out, present_kv) if use_cache else out

class SwiGLUMLP(nn.Module):
    def __init__(self, d_model, hidden_dim, dropout=0.0):
        super().__init__()
        self.gate    = nn.Linear(d_model, hidden_dim, bias=False)
        self.up      = nn.Linear(d_model, hidden_dim, bias=False)
        self.down    = nn.Linear(hidden_dim, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)
    def forward(self, x):
        return self.down(self.dropout(F.silu(self.gate(x)) * self.up(x)))

class EngramMemory(nn.Module):
    def __init__(self, d_model, bucket_count, memory_dim, order, pad_id=0, dropout=0.0):
        super().__init__()
        self.bucket_count = bucket_count
        self.order   = order
        self.pad_id  = pad_id
        self.bucket  = nn.Embedding(bucket_count, memory_dim)
        self.query   = nn.Linear(d_model, memory_dim, bias=False)
        self.project = nn.Linear(memory_dim, d_model, bias=False)
        self.gate    = nn.Linear(d_model, d_model, bias=True)
        self.dropout = nn.Dropout(dropout)
        primes = [1, 1315423911, 2654435761, 97531, 433494437]
        self.register_buffer("primes", torch.tensor(primes[:order], dtype=torch.long), persistent=False)

    def hash_tokens(self, idx):
        batch, seq_len = idx.shape
        pad     = torch.full((batch, self.order - 1), self.pad_id, device=idx.device, dtype=idx.dtype)
        history = torch.cat([pad, idx], dim=1)
        hashed  = torch.zeros((batch, seq_len), device=idx.device, dtype=torch.long)
        for offset in range(self.order):
            slice_ = history[:, offset: offset + seq_len].long()
            hashed = (hashed * 1315423911 + slice_ * self.primes[offset]) % self.bucket_count
        return hashed

    def forward(self, x, idx):
        hashed = self.hash_tokens(idx)
        if hashed.size(1) != x.size(1):
            hashed = hashed[:, -x.size(1):]
        query  = torch.tanh(self.query(x))
        memory = self.bucket(hashed) * query
        memory = self.project(memory)
        gate   = torch.sigmoid(self.gate(x))
        return self.dropout(gate * memory)

class Block(nn.Module):
    def __init__(self, d_model, n_heads, n_kv_heads, d_ff, rope_base, dropout=0.0,
                 use_engram=False, engram_bucket_count=64, engram_dim=16, engram_order=3, pad_id=0):
        super().__init__()
        self.norm1  = RMSNorm(d_model)
        self.attn   = GQAAttention(d_model, n_heads, n_kv_heads, rope_base=rope_base, dropout=dropout)
        self.norm2  = RMSNorm(d_model)
        self.engram = EngramMemory(d_model, engram_bucket_count, engram_dim, engram_order,
                                   pad_id=pad_id, dropout=dropout) if use_engram else None
        self.norm3  = RMSNorm(d_model)
        self.mlp    = SwiGLUMLP(d_model, d_ff, dropout=dropout)

    def forward(self, x, idx, rope_cos, rope_sin):
        x = x + self.attn(self.norm1(x), rope_cos, rope_sin)
        if self.engram is not None:
            x = x + self.engram(self.norm2(x), idx)
        return x + self.mlp(self.norm3(x))

    def forward_cached(self, x, idx_context, rope_cos, rope_sin, past_kv=None):
        attn_out, present_kv = self.attn(self.norm1(x), rope_cos, rope_sin,
                                         past_kv=past_kv, use_cache=True)
        x = x + attn_out
        if self.engram is not None:
            x = x + self.engram(self.norm2(x), idx_context)
        x = x + self.mlp(self.norm3(x))
        return x, present_kv

class CosmosT2_Accelerate_LLM(nn.Module):
    def __init__(self, vocab_size, d_model=32, n_layers=6, n_heads=2, n_kv_heads=1,
                 d_ff=256, max_len=1028, rope_base=10000, dropout=0.05, use_engram=True,
                 engram_every=2, engram_bucket_count=64, engram_dim=16, engram_order=3, pad_id=0):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model    = d_model
        self.n_layers   = n_layers
        self.n_heads    = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim   = d_model // n_heads
        self.max_len    = max_len
        self.rope_base  = rope_base
        self.pad_id     = pad_id
        self.tok_emb    = nn.Embedding(vocab_size, d_model)
        self.blocks     = nn.ModuleList()
        for layer_index in range(n_layers):
            block_uses_engram = use_engram and ((layer_index + 1) % engram_every == 0)
            self.blocks.append(Block(
                d_model=d_model, n_heads=n_heads, n_kv_heads=n_kv_heads, d_ff=d_ff,
                rope_base=rope_base, dropout=dropout, use_engram=block_uses_engram,
                engram_bucket_count=engram_bucket_count, engram_dim=engram_dim,
                engram_order=engram_order, pad_id=pad_id,
            ))
        self.norm_f = RMSNorm(d_model)
        self.apply(self._init_weights)

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

    def build_rope(self, seq_len, device, dtype, start_pos=0):
        inv_freq = 1.0 / (self.rope_base ** (
            torch.arange(0, self.head_dim, 2, device=device, dtype=torch.float32) / self.head_dim
        ))
        positions = torch.arange(start_pos, start_pos + seq_len, device=device, dtype=torch.float32)
        freqs = torch.outer(positions, inv_freq)
        cos = freqs.cos().repeat_interleave(2, dim=-1).to(dtype)[None, None, :, :]
        sin = freqs.sin().repeat_interleave(2, dim=-1).to(dtype)[None, None, :, :]
        return cos, sin

    def forward(self, idx, targets=None):
        if idx.size(1) > self.max_len:
            idx = idx[:, -self.max_len:]
        seq_len = idx.size(1)
        rope_cos, rope_sin = self.build_rope(seq_len, idx.device, self.tok_emb.weight.dtype)
        x = self.tok_emb(idx)
        for block in self.blocks:
            x = block(x, idx, rope_cos, rope_sin)
        x = self.norm_f(x)
        logits = F.linear(x, self.tok_emb.weight)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1))
        return logits, loss

    def trim_kv_cache(self, past_kv, max_tokens):
        if past_kv is None:
            return None
        max_tokens = max(0, int(max_tokens))
        trimmed = []
        for k, v in past_kv:
            if max_tokens == 0:
                k = k[:, :, :0, :].contiguous()
                v = v[:, :, :0, :].contiguous()
            elif k.size(2) > max_tokens:
                k = k[:, :, -max_tokens:, :].contiguous()
                v = v[:, :, -max_tokens:, :].contiguous()
            trimmed.append((k, v))
        return trimmed

    @torch.no_grad()
    def forward_cached(self, idx, past_kv=None, cache_pos=0, max_ctx=None, idx_context=None):
        self.eval()
        max_ctx = self.max_len if max_ctx is None else int(max_ctx)
        if past_kv is None:
            idx         = idx[:, -max_ctx:]
            idx_context = idx
            cache_pos   = 0
        else:
            keep_past   = max(0, max_ctx - idx.size(1))
            past_kv     = self.trim_kv_cache(past_kv, keep_past)
            idx_context = idx if idx_context is None else idx_context[:, -max_ctx:]
        seq_len = idx.size(1)
        rope_cos, rope_sin = self.build_rope(
            seq_len, idx.device, self.tok_emb.weight.dtype, start_pos=cache_pos)
        x = self.tok_emb(idx)
        present_kv = []
        for layer_index, block in enumerate(self.blocks):
            layer_past = None if past_kv is None else past_kv[layer_index]
            x, layer_present = block.forward_cached(x, idx_context, rope_cos, rope_sin, past_kv=layer_past)
            present_kv.append(layer_present)
        x = self.norm_f(x)
        logits = F.linear(x, self.tok_emb.weight)
        return logits, present_kv, cache_pos + seq_len

    def sample_next(self, logits, temperature=0.8, top_k=50):
        if logits.dim() == 3:
            logits = logits[:, -1, :]
        if temperature <= 1e-6:
            return torch.argmax(logits, dim=-1, keepdim=True)
        logits = logits / temperature
        if top_k and top_k > 0:
            values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
            logits = logits.masked_fill(logits < values[:, [-1]], float("-inf"))
        probs = F.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1)

    @torch.no_grad()
    def prefill_cache(self, idx, max_ctx=None):
        logits, past_kv, cache_pos = self.forward_cached(idx, past_kv=None, cache_pos=0, max_ctx=max_ctx)
        return logits[:, -1, :], past_kv, cache_pos

    @torch.no_grad()
    def decode_cached(self, idx, past_kv, cache_pos, idx_context, max_ctx=None):
        logits, past_kv, cache_pos = self.forward_cached(
            idx, past_kv=past_kv, cache_pos=cache_pos, max_ctx=max_ctx, idx_context=idx_context)
        return logits[:, -1, :], past_kv, cache_pos

    @torch.no_grad()
    def generate(self, idx, max_new_tokens=128, temperature=0.8, top_k=50,
                 max_ctx=None, stop_ids=None):
        self.eval()
        if stop_ids is None:
            stop_ids = STOP_IDS
        max_ctx = self.max_len if max_ctx is None else int(max_ctx)
        idx = idx[:, -max_ctx:]
        logits, past_kv, cache_pos = self.prefill_cache(idx, max_ctx=max_ctx)
        for step in range(max_new_tokens):
            nxt = self.sample_next(logits, temperature=temperature, top_k=top_k)
            if stop_ids and nxt.numel() == 1 and int(nxt.item()) in stop_ids:
                break
            idx = torch.cat([idx, nxt], dim=1)
            if step + 1 < max_new_tokens:
                logits, past_kv, cache_pos = self.decode_cached(
                    nxt, past_kv, cache_pos, idx[:, -max_ctx:], max_ctx=max_ctx)
        return idx

# ---------------- TOKENIZER + STOP IDS ----------------

tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

def _resolve_stop_ids(tok):
    ids = set()
    for t in ("<|im_end|>", "<|endoftext|>"):
        i = tok.convert_tokens_to_ids(t)
        if isinstance(i, int) and i >= 0 and i != tok.unk_token_id:
            ids.add(i)
    if tok.eos_token_id is not None:
        ids.add(tok.eos_token_id)
    return ids

STOP_IDS = _resolve_stop_ids(tokenizer)

# ---------------- MODEL CACHE ----------------

_MODEL_CACHE = {}
_LOAD_LOCK   = threading.Lock()

def get_model(model_key, ckpt_label):
    key = (model_key, ckpt_label)
    if key in _MODEL_CACHE:
        return _MODEL_CACHE[key]
    with _LOAD_LOCK:
        if key in _MODEL_CACHE:
            return _MODEL_CACHE[key]
        spec     = MODELS[model_key]
        filename = spec["checkpoints"][ckpt_label]
        print(f"Loading {spec['full_name']} [{ckpt_label}] -> {spec['repo']}/{filename}")
        path = hf_hub_download(repo_id=spec["repo"], filename=filename)
        ckpt = torch.load(path, map_location="cpu", weights_only=False)
        cfg  = ckpt.get("config", {})
        resolved = {
            "vocab_size":          cfg.get("vocab_size", len(tokenizer)),
            "d_model":             cfg.get("d_model", 32),
            "n_layers":            cfg.get("n_layers", 4),
            "n_heads":             cfg.get("n_heads", 4),
            "n_kv_heads":          cfg.get("n_kv_heads", 2),
            "d_ff":                cfg.get("d_ff", 256),
            "max_len":             cfg.get("max_len", 1028),
            "rope_base":           cfg.get("rope_base", 10000),
            "dropout":             0.0,
            "use_engram":          cfg.get("use_engram", True),
            "engram_every":        cfg.get("engram_every", 2),
            "engram_bucket_count": cfg.get("engram_bucket_count", cfg.get("engram_buckets", 64)),
            "engram_dim":          cfg.get("engram_dim", 16),
            "engram_order":        cfg.get("engram_order", 3),
            "pad_id":              tokenizer.pad_token_id,
        }
        model = CosmosT2_Accelerate_LLM(**resolved)
        state = ckpt.get("model_state", ckpt)
        state = {k.replace("module.", "", 1): v for k, v in state.items()}
        model.load_state_dict(state, strict=False)
        model = model.to(DEVICE).to(DTYPE).eval()
        n_params = sum(p.numel() for p in model.parameters())
        info = {"params": n_params, "max_len": resolved["max_len"]}
        _MODEL_CACHE[key] = (model, info)
        print(f"  loaded {n_params/1e6:.2f}M params  (ctx {resolved['max_len']})")
        return _MODEL_CACHE[key]

# ---------------- GENERATION HELPERS ----------------

def build_prompt(user_msg, system_msg=DEFAULT_SYSTEM_PROMPT, history=None):
    messages = []
    if system_msg and system_msg.strip():
        messages.append({"role": "system", "content": system_msg.strip()})
    if history:
        for item in history:
            messages.append(item)
    messages.append({"role": "user", "content": user_msg})
    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

def clean_output(text):
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    return text.strip()

def run_model(model_key, ckpt_label, user_msg, max_new=256, temperature=0.7, top_k=40):
    model, info = get_model(model_key, ckpt_label)
    prompt = build_prompt(user_msg)
    ids    = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids.to(DEVICE)
    max_ctx = info["max_len"]
    ids    = ids[:, -max_ctx:]
    out    = model.generate(ids, max_new_tokens=max_new, temperature=temperature,
                            top_k=top_k, max_ctx=max_ctx)
    generated = out[0, ids.size(1):]
    return clean_output(tokenizer.decode(generated.tolist(), skip_special_tokens=True))

# ---------------- PER-USER SESSIONS ----------------

user_sessions = {}

def get_session(user_id):
    if user_id not in user_sessions:
        user_sessions[user_id] = {"model": "beta2", "ckpt": "best"}
    return user_sessions[user_id]

# ---------------- HELP TEXT ----------------

HELP_TEXT = """**Cosmos T2-Accelerate — Commands**

`.ask <message>` — Ask the AI a question
`.model <name>`  — Switch model (`beta2`, `preview`, `beta`)
`.ckpt <label>`  — Switch checkpoint (`best`, `latest`)
`.status`        — Show your current selection
`.models`        — List all available models
`.help`          — Show this message
"""

# ============================================================
# BOT FACTORY — creates a fresh Bot instance each restart
# ============================================================

def make_bot():
    _intents = discord.Intents.default()
    _intents.message_content = True
    _intents.guilds  = True
    _intents.members = True

    b = commands.Bot(command_prefix=".", intents=_intents, help_command=None)

    # ---- events ----

    @b.event
    async def on_ready():
        print(f"Logged in as {b.user}")

    @b.event
    async def on_message(message: discord.Message):
        if message.author == b.user:
            return
        await b.process_commands(message)

    # ---- commands ----

    @b.command(name="help")
    async def cmd_help(ctx):
        await ctx.send(HELP_TEXT)

    @b.command(name="models")
    async def cmd_models(ctx):
        lines = ["**Available models:**"]
        for key, spec in MODELS.items():
            lines.append(f"• `{key}` — {spec['full_name']} · {spec['desc']}")
        lines.append("\nCheckpoints: `best` (lowest val loss) · `latest` (end of training)")
        await ctx.send("\n".join(lines))

    @b.command(name="model")
    async def cmd_model(ctx, *, name: str = None):
        if name is None:
            await ctx.send("Usage: `.model <name>` — Options: `beta2`, `preview`, `beta`")
            return
        name = name.strip().lower()
        if name not in MODELS:
            await ctx.send(f"Unknown model `{name}`. Options: `beta2`, `preview`, `beta`")
            return
        get_session(ctx.author.id)["model"] = name
        spec = MODELS[name]
        await ctx.send(f"✅ Model set to **{spec['full_name']}** ({spec['desc']})")

    @b.command(name="ckpt")
    async def cmd_ckpt(ctx, *, label: str = None):
        if label is None:
            await ctx.send("Usage: `.ckpt <label>` — Options: `best`, `latest`")
            return
        label = label.strip().lower()
        if label not in ("best", "latest"):
            await ctx.send("Unknown checkpoint. Options: `best` or `latest`")
            return
        get_session(ctx.author.id)["ckpt"] = label
        await ctx.send(f"✅ Checkpoint set to **{label}**")

    @b.command(name="status")
    async def cmd_status(ctx):
        sess = get_session(ctx.author.id)
        spec = MODELS[sess["model"]]
        await ctx.send(
            f"**Your current selection:**\n"
            f"• Model: `{sess['model']}` — {spec['full_name']} ({spec['desc']})\n"
            f"• Checkpoint: `{sess['ckpt']}`"
        )

    @b.command(name="ask")
    async def cmd_ask(ctx, *, message: str = None):
        if not message:
            await ctx.send("Usage: `.ask <your message>`")
            return
    
        sess       = get_session(ctx.author.id)
        model_key  = sess["model"]
        ckpt_label = sess["ckpt"]
        spec       = MODELS[model_key]

        thinking_msg = await ctx.send(
            f"⏳ Thinking with **{spec['full_name']}** [{ckpt_label}]…"
        )

        def generate():
            try:
                return run_model(model_key, ckpt_label, message)
            except Exception as e:
                return f"⚠️ Error: {e}"
    
        result = await b.loop.run_in_executor(None, generate)
        await thinking_msg.delete()
    
        if not result:
            result = "*(no output)*"
    
        # Split on 1900 chars but try to break at newlines to avoid
        # cutting mid-sentence or mid-thought
        chunks = []
        while len(result) > 1900:
            split_at = result.rfind("\n", 0, 1900)
            if split_at == -1:
                split_at = 1900
            chunks.append(result[:split_at])
            result = result[split_at:].lstrip("\n")
        chunks.append(result)

        for i, chunk in enumerate(chunks):
            header = f"**[{spec['short']} / {ckpt_label}]** " if i == 0 else f"**[cont. {i+1}]** "
            await ctx.send(f"{header}{chunk}")

    return b

# ============================================================
# BOT RUNNER — recreates the bot object on every crash
# ============================================================

def run_bot():
    if not DISCORD_TOKEN:
        print("DISCORD_TOKEN not set")
        return
    while True:
        b = make_bot()          # fresh Bot + fresh aiohttp session each time
        try:
            b.run(DISCORD_TOKEN)
        except Exception as e:
            print(f"Bot crashed: {e}")
            print("Restarting in 5 seconds…")
            time.sleep(5)

threading.Thread(target=run_bot, daemon=True).start()

# ============================================================
# GRADIO UI
# ============================================================

with gr.Blocks() as demo:
    gr.Markdown("# Cosmos T2-Accelerate — Discord bot is online!")
    gr.Markdown("Use `.ask`, `.model`, `.ckpt`, `.status`, `.models`, `.help` in Discord.")

demo.launch()