import torch
import torch.nn as nn
import torch.nn.functional as F
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
import os, pickle, sys

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_PATH = os.environ.get("MODEL_PATH", "best_generator.pt")
VOCAB_PATH = os.environ.get("VOCAB_PATH", "vocab.pt")
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "7"))

# ---------------------------------------------------------------------------
# Model architecture
# ---------------------------------------------------------------------------

class RotaryPositionalEmbedding(nn.Module):
    def __init__(self, dim=32):
        super().__init__()
        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim).float() / dim))
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, x):
        seq_len = x.shape[1]
        t = torch.arange(seq_len, device=x.device).float()
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        return freqs[:seq_len]


class MultiHeadAttentionWithRoPE(nn.Module):
    def __init__(self, d_model=256, n_heads=8, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        self.W_q = nn.Linear(d_model, d_model, bias=True)
        self.W_k = nn.Linear(d_model, d_model, bias=True)
        self.W_v = nn.Linear(d_model, d_model, bias=True)
        self.W_o = nn.Linear(d_model, d_model, bias=True)

        self.rope = RotaryPositionalEmbedding(dim=self.d_head)
        self.dropout = nn.Dropout(dropout)

    def _rotate_half(self, x):
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)

    def _apply_rope(self, x, freqs):
        return x * freqs.cos() + self._rotate_half(x) * freqs.sin()

    def forward(self, x, mask=None):
        B, T, C = x.shape
        Q = self.W_q(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        K = self.W_k(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        V = self.W_v(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)

        freqs = self.rope(x)
        Q = self._apply_rope(Q, freqs)
        K = self._apply_rope(K, freqs)

        attn = Q @ K.transpose(-2, -1) / (self.d_head ** 0.5)
        if mask is not None:
            attn = attn.masked_fill(mask == 0, float("-inf"))
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)

        out = (attn @ V).transpose(1, 2).contiguous().view(B, T, C)
        return self.W_o(out)


class FeedForward(nn.Module):
    def __init__(self, d_model=256, d_ff=1024, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff, bias=True)
        self.linear2 = nn.Linear(d_ff, d_model, bias=True)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.linear2(self.dropout(F.gelu(self.linear1(x))))


class DecoderBlockWithRoPE(nn.Module):
    def __init__(self, d_model=256, n_heads=8, d_ff=1024, dropout=0.1):
        super().__init__()
        self.self_attention = MultiHeadAttentionWithRoPE(d_model, n_heads, dropout)
        self.feed_forward = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        x = x + self.dropout(self.self_attention(self.norm1(x), mask))
        x = x + self.dropout(self.feed_forward(self.norm2(x)))
        return x


class GPTModelWithRoPE(nn.Module):
    def __init__(self, vocab_size=42962, d_model=256, n_heads=8,
                 n_layers=3, d_ff=1024, dropout=0.1):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.embedding_dropout = nn.Dropout(dropout)
        self.decoder_blocks = nn.ModuleList([
            DecoderBlockWithRoPE(d_model, n_heads, d_ff, dropout)
            for _ in range(n_layers)
        ])
        self.final_norm = nn.LayerNorm(d_model)
        self.output_projection = nn.Linear(d_model, vocab_size, bias=True)

    def forward(self, x, mask=None):
        x = self.embedding_dropout(self.token_embedding(x))
        for block in self.decoder_blocks:
            x = block(x, mask)
        x = self.final_norm(x)
        logits = self.output_projection(x)
        return logits

    @torch.no_grad()
    def generate(self, input_ids, max_new_tokens=7, temperature=1.0,
                 top_k=50, top_p=0.9, eos_token_id=None):
        self.eval()
        B, T = input_ids.shape
        # Causal mask for current sequence length
        mask = torch.tril(torch.ones(1, 1, T, T, device=input_ids.device))

        for _ in range(max_new_tokens):
            logits = self(input_ids, mask)
            logits = logits[:, -1, :]

            if temperature > 0:
                logits = logits / temperature

            # Replace nan/inf with -inf to avoid crashes
            logits = torch.nan_to_num(logits, nan=float("-inf"), posinf=float("-inf"), neginf=float("-inf"))

            if top_k > 0:
                values, _ = torch.topk(logits, top_k, dim=-1)
                logits[logits < values[:, -1:]] = float("-inf")

            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
                sorted_indices_to_remove = cumulative_probs > top_p
                sorted_logits[sorted_indices_to_remove] = float("-inf")
                logits = sorted_logits.scatter(1, sorted_indices, sorted_logits)

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)

            # Update causal mask for new token
            input_ids = torch.cat([input_ids, next_token], dim=-1)
            new_T = input_ids.shape[1]
            mask = torch.tril(torch.ones(1, 1, new_T, new_T, device=input_ids.device))

            if eos_token_id is not None and next_token.item() == eos_token_id:
                break

        return input_ids


# ---------------------------------------------------------------------------
# Vocab (in case no vocab file exists)
# ---------------------------------------------------------------------------

class DummyVocab:
    """Minimal vocab that passes through token IDs as text."""
    word2idx = {'<BOS>': 0, '<EOS>': 1, '<PAD>': 2, '<UNK>': 3}
    idx2word = {0: '<BOS>', 1: '<EOS>', 2: '<PAD>', 3: '<UNK>'}

    def __init__(self):
        for i in range(4, 42962):
            self.word2idx[f'<TOK_{i}>'] = i
            self.idx2word[i] = f'<TOK_{i}>'

    def text_to_indices_for_generator(self, text):
        return [self.word2idx.get(w, 3) for w in text.strip().split()]

    def indices_to_text(self, indices):
        return ' '.join(self.idx2word.get(i, '<UNK>') for i in indices)


# ---------------------------------------------------------------------------
# Generation function
# ---------------------------------------------------------------------------

@torch.no_grad()
def generate_clickbait(model, vocab, prompt="", max_new_tokens=7,
                       temperature=1.0, top_k=50, top_p=0.9):
    model.eval()

    if prompt:
        tokens = vocab.text_to_indices_for_generator(prompt)
        input_ids = torch.LongTensor([[vocab.word2idx['<BOS>']] + tokens]).to(DEVICE)
    else:
        input_ids = torch.LongTensor([[vocab.word2idx['<BOS>']]]).to(DEVICE)

    generated = model.generate(
        input_ids,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        eos_token_id=vocab.word2idx.get('<EOS>', 1)
    )

    generated_text = vocab.indices_to_text(generated[0].tolist())

    # Clean up special tokens for display
    for tok in ['<BOS>', '<EOS>', '<PAD>', '<UNK>']:
        generated_text = generated_text.replace(tok, '')
    generated_text = generated_text.strip()

    return generated_text


# ---------------------------------------------------------------------------
# API
# ---------------------------------------------------------------------------

app = FastAPI(title="ClickBait Generator API")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class GenerateRequest(BaseModel):
    prompt: str = ""
    temperature: float = 1.0
    top_k: int = 50
    num_samples: int = 1


class GenerateResponse(BaseModel):
    texts: list[str]


model = None
vocab = None
using_dummy_vocab = False


@app.on_event("startup")
def load_artifacts():
    global model, vocab, using_dummy_vocab
    if not os.path.exists(MODEL_PATH):
        print(f"[WARN] {MODEL_PATH} not found. Model loading skipped.")
        return

    try:
        checkpoint = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=True)

        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
            state_dict = checkpoint['model_state_dict']
        else:
            state_dict = checkpoint

        vocab_size = state_dict['token_embedding.weight'].shape[0]
        model = GPTModelWithRoPE(vocab_size=vocab_size).to(DEVICE)

        # Remove causal_mask (we compute it dynamically)
        sd_clean = {k: v for k, v in state_dict.items() if k != 'causal_mask'}
        missing, unexpected = model.load_state_dict(sd_clean, strict=False)
        if missing:
            print(f"[WARN] Missing keys: {missing}")
        if unexpected:
            print(f"[WARN] Unexpected keys: {unexpected}")

        model.eval()

        # Load vocab
        if os.path.exists(VOCAB_PATH):
            vocab = pickle.load(open(VOCAB_PATH, "rb"))
            using_dummy_vocab = False
            print(f"[OK] Vocab loaded: {len(vocab.word2idx)} tokens")
        else:
            print(f"[WARN] {VOCAB_PATH} not found, using fallback dummy vocab")
            vocab = DummyVocab()
            using_dummy_vocab = True

        print(f"[OK] Model loaded on {DEVICE} | params: {sum(p.numel() for p in model.parameters()):,}")

    except Exception as e:
        print(f"[ERROR] Failed to load model: {e}")
        import traceback
        traceback.print_exc()


@app.get("/api/health")
def health():
    return {
        "status": "ok",
        "model_loaded": model is not None,
        "vocab_loaded": vocab is not None,
        "using_dummy_vocab": using_dummy_vocab,
        "device": str(DEVICE),
        "vocab_size": len(vocab.word2idx) if vocab else 0,
        "note": "vocab.pt not found — output shows token IDs" if using_dummy_vocab else "ready",
    }


@app.post("/api/generate", response_model=GenerateResponse)
def generate(req: GenerateRequest):
    if model is None:
        raise HTTPException(status_code=503, detail="Model not loaded")

    # Clamp / validate
    temperature = max(0.1, min(3.0, req.temperature))
    top_k = max(1, min(200, req.top_k))
    count = max(1, min(20, req.num_samples))

    texts = []
    for _ in range(count):
        text = generate_clickbait(
            model, vocab or DummyVocab(),
            prompt=req.prompt,
            temperature=temperature,
            top_k=top_k,
            max_new_tokens=MAX_NEW_TOKENS,
        )
        texts.append(text)

    return GenerateResponse(texts=texts)


# ---------------------------------------------------------------------------
# Static files + SPA fallback
# ---------------------------------------------------------------------------

STATIC_DIR = os.path.dirname(os.path.abspath(__file__))

@app.get("/")
def serve_index():
    return FileResponse(os.path.join(STATIC_DIR, "index.html"))

@app.get("/{path:path}")
def serve_static(path: str):
    file_path = os.path.join(STATIC_DIR, path)
    if os.path.isfile(file_path):
        return FileResponse(file_path)
    return FileResponse(os.path.join(STATIC_DIR, "index.html"))


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 8000))
    uvicorn.run(app, host="0.0.0.0", port=port)