"""Minimal load+generate helper for the AetherStory model on Hugging Face.

Users can copy this file next to ``model.safetensors`` / ``config.json`` /
``tokenizer.json`` and run::

    from modeling_aetherstory import StoryTeller
    t = StoryTeller.from_dir(".")
    print(t("Once upon a time"))
"""

import json
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file


# --- model definition (mirrors src/model.py) -------------------------------
class _Attn(nn.Module):
    def __init__(s, d, h, drop):
        super().__init__()
        s.h, s.hd, s.scale = h, d // h, (d // h) ** -0.5
        s.qkv = nn.Linear(d, 3 * d, bias=False)
        s.proj = nn.Linear(d, d, bias=False)
        s.drop = nn.Dropout(drop)
    def forward(s, x):
        B, T, C = x.shape
        qkv = s.qkv(x).reshape(B, T, 3, s.h, s.hd).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        att = (q @ k.transpose(-2, -1)) * s.scale
        att = att.masked_fill(~torch.tril(torch.ones(T, T, device=x.device, dtype=torch.bool)), float("-inf"))
        att = s.drop(F.softmax(att, dim=-1))
        y = (att @ v).transpose(1, 2).contiguous().reshape(B, T, C)
        return s.drop(s.proj(y))


class _FFN(nn.Module):
    def __init__(s, d, f, drop):
        super().__init__()
        s.fc1, s.fc2, s.drop = nn.Linear(d, f, bias=False), nn.Linear(f, d, bias=False), nn.Dropout(drop)
    def forward(s, x):
        return s.drop(s.fc2(F.gelu(s.fc1(x))))


class _Block(nn.Module):
    def __init__(s, d, h, f, drop):
        super().__init__()
        s.ln1, s.attn = nn.LayerNorm(d), _Attn(d, h, drop)
        s.ln2, s.ffn = nn.LayerNorm(d), _FFN(d, f, drop)
    def forward(s, x):
        x = x + s.attn(s.ln1(x)); x = x + s.ffn(s.ln2(x)); return x


class AetherStoryModel(nn.Module):
    def __init__(s, cfg):
        super().__init__()
        s.cfg = cfg
        s.tok_emb = nn.Embedding(cfg["vocab_size"], cfg["d_model"])
        s.pos_emb = nn.Parameter(torch.zeros(1, cfg["max_seq_len"], cfg["d_model"]))
        s.blocks = nn.ModuleList([_Block(cfg["d_model"], cfg["n_heads"], cfg["ffn_dim"], cfg["dropout"]) for _ in range(cfg["n_layers"])])
        s.ln_f = nn.LayerNorm(cfg["d_model"])
        if cfg.get("tie_embeddings", True):
            s.head_bias = nn.Parameter(torch.zeros(cfg["vocab_size"])); s.lm_head = None
        else:
            s.head_bias = None; s.lm_head = nn.Linear(cfg["d_model"], cfg["vocab_size"], bias=False)
    def forward(s, idx, targets=None):
        B, T = idx.shape
        x = s.tok_emb(idx) + s.pos_emb[:, :T, :]
        for b in s.blocks: x = b(x)
        x = s.ln_f(x)
        logits = s.lm_head(x) if s.lm_head is not None else (x @ s.tok_emb.weight.t() + s.head_bias)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.reshape(-1, s.cfg["vocab_size"]), targets.reshape(-1), ignore_index=s.cfg.get("pad_token_id", 0))
        return logits, loss
    @torch.no_grad()
    def generate(s, idx, max_new, temperature=0.9, top_k=40, eos_token_id=None):
        s.eval()
        for _ in range(max_new):
            ic = idx if idx.size(1) <= s.cfg["max_seq_len"] else idx[:, -s.cfg["max_seq_len"]:]
            logits, _ = s(ic)
            logits = logits[:, -1, :] / max(temperature, 1e-5)
            if top_k and top_k > 0:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = float("-inf")
            nxt = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
            idx = torch.cat([idx, nxt], dim=1)
            if eos_token_id is not None and (nxt == eos_token_id).all(): break
        return idx


class StoryTeller:
    @classmethod
    def from_dir(cls, d="."):
        d = Path(d)
        cfg = json.loads((d / "config.json").read_text(encoding="utf-8"))
        tok_payload = json.loads((d / "tokenizer.json").read_text(encoding="utf-8"))
        model = AetherStoryModel(cfg)
        state = load_file(str(d / "model.safetensors"))
        model.load_state_dict({k: v for k, v in state.items()})
        model.eval()
        w2i = {}
        for word, idx in tok_payload["special_tokens"].items(): w2i[word] = idx
        for idx, word in enumerate(tok_payload["vocab"]):
            if idx not in w2i: w2i[word] = idx
        i2w = {v: k for k, v in w2i.items()}
        inst = cls(); inst.model = model; inst.w2i = w2i; inst.i2w = i2w
        inst.bos = tok_payload["special_tokens"]["<bos>"]; inst.eos = tok_payload["special_tokens"]["<eos>"]
        return inst
    def _encode(self, text):
        import re
        ids = [self.bos]
        for t in re.findall(r"\w+|[^\w\s]|\s+", text.lower()):
            if t.strip() and t in self.w2i: ids.append(self.w2i[t])
        return ids
    def _decode(self, ids):
        out = []
        for i in ids:
            if i == self.bos or i == 0: continue
            if i == self.eos: break
            w = self.i2w.get(i, "")
            if w.startswith("<") and w.endswith(">"): continue
            out.append(w)
        t = " ".join(out)
        import re
        return re.sub(r"\s+([,.;:!?\'\"()])", r"\1", t).strip()
    def __call__(self, prompt, max_tokens=80, temperature=0.9, top_k=40, seed=None):
        if seed is not None: torch.manual_seed(seed)
        ids = self._encode(prompt)
        idx = torch.tensor([ids], dtype=torch.long)
        out = self.model.generate(idx, max_tokens, temperature, top_k, self.eos)
        return self._decode(out[0].tolist())