# file: train_and_generate.py
import os
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from architecture import TransformerBlock
from dataset import build_or_load_tokenizer, AutoregressiveLogDataset


class LogSentryLM(nn.Module):
    """Decoder-only (GPT-style) LM. Defaults are deliberately tiny (~2M params)
    so it trains in under a minute on a laptop GPU."""

    def __init__(self, vocab_size, d_model=128, n_heads=4, n_layers=3, max_seq_len=64):
        super().__init__()
        # max_seq_len is the context window. It's fixed up front because positions
        # are LEARNED embeddings (one vector per slot), not computed on the fly.
        self.max_seq_len = max_seq_len
        self.token_embeddings = nn.Embedding(vocab_size, d_model)
        self.position_embeddings = nn.Embedding(max_seq_len, d_model)  # attention has no order sense on its own

        self.blocks = nn.ModuleList([
            TransformerBlock(d_model=d_model, n_heads=n_heads) for _ in range(n_layers)
        ])
        self.ln_final = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)  # project back to vocab logits

    def forward(self, idx):
        b, s = idx.shape
        positions = torch.arange(0, s, device=idx.device).unsqueeze(0)
        # token identity + position — the model needs both
        x = self.token_embeddings(idx) + self.position_embeddings(positions)

        for block in self.blocks:
            x = block(x, is_causal=True)

        return self.lm_head(self.ln_final(x))

    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=0.7):
        """Autoregressive sampling — predict one token, append, repeat.

        temperature < 1 sharpens toward the likely token (safe/repetitive),
        > 1 flattens it (more diverse, more mistakes).
        """
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.max_seq_len:]  # keep only what fits the context window
            logits = self(idx_cond)
            logits = logits[:, -1, :] / temperature  # only the last position predicts the next token
            probs = torch.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)  # sample rather than argmax for variety
            idx = torch.cat((idx, next_token), dim=1)
        return idx


@torch.no_grad()
def score_line(model, tokenizer, text, device, max_seq_len=128):
    """Anomaly score for one line = average next-token loss.

    The model learned normal logs, so a line it predicts well scores LOW and a
    surprising line scores HIGH. Returns (loss, perplexity); None if too short.
    """
    model.eval()
    ids = tokenizer.encode(text).ids
    if len(ids) < 2:
        return None  # need at least one (context -> target) pair
    ids = ids[: max_seq_len + 1]
    x = torch.tensor([ids[:-1]], dtype=torch.long, device=device)
    y = torch.tensor([ids[1:]], dtype=torch.long, device=device)
    logits = model(x)
    loss = F.cross_entropy(logits.view(-1, logits.size(-1)), y.view(-1))
    return loss.item(), math.exp(loss.item())  # perplexity is the intuitive scale


def fit_anomaly_threshold(model, tokenizer, lines, device, max_seq_len=128, k=2.0):
    """Derive the normal-vs-anomalous cutoff from the corpus itself.

    Most lines are normal, so threshold = mean + k*std of the scores. Lower k =
    more sensitive (more flags/false positives); higher k = stricter.
    """
    scored = []
    for line in lines:
        result = score_line(model, tokenizer, line, device, max_seq_len)
        if result is not None:
            loss, ppl = result
            scored.append((loss, ppl, line))

    losses = [s[0] for s in scored]
    mean = sum(losses) / len(losses)
    var = sum((l - mean) ** 2 for l in losses) / len(losses)
    std = math.sqrt(var)
    threshold = mean + k * std
    scored.sort(key=lambda s: s[0], reverse=True)  # most anomalous first
    return threshold, mean, std, scored


def generate_mock_logs(file_path):
    """Fallback synthetic corpus. Lines are duplicated so the tiny model has an
    obvious pattern to latch onto if the real logs are missing."""
    os.makedirs(os.path.dirname(file_path), exist_ok=True)
    mock_data = """[ERROR] nginx failed. Fix: run systemctl restart nginx
[INFO] database connection secure.
[ERROR] out of memory. Fix: upgrade ram capacity
[ERROR] nginx failed. Fix: run systemctl restart nginx
[INFO] disk space clean.
[ERROR] out of memory. Fix: upgrade ram capacity
""" * 50
    with open(file_path, "w") as f:
        f.write(mock_data)


if __name__ == "__main__":
    # Prefer GPU: cuda -> apple mps -> cpu
    if torch.cuda.is_available():
        device = "cuda"
    elif torch.backends.mps.is_available():
        device = "mps"
    else:
        device = "cpu"
    print(f"[SETUP] Using device: {device}")

    LOG_FILE = "data/raw_logs.txt"
    TOKENIZER_FILE = "tokenizer/log_tokenizer.json"

    # --- Data ---
    if not os.path.exists(LOG_FILE):
        print("[DATA] raw_logs.txt not found — generating synthetic mock logs as fallback.")
        generate_mock_logs(LOG_FILE)
    tokenizer = build_or_load_tokenizer(LOG_FILE, TOKENIZER_FILE, vocab_size=10000)  # ceiling; small corpus fills ~5.5k

    max_len = 128
    dataset = AutoregressiveLogDataset(LOG_FILE, tokenizer, max_seq_len=max_len, stride=64)
    dataloader = DataLoader(dataset, batch_size=64, shuffle=True)  # bigger batch = better GPU use
    print(f"[DATA] {len(dataset)} windows | {len(dataloader)} batches/epoch")

    # --- Model ---
    model = LogSentryLM(vocab_size=tokenizer.get_vocab_size(), max_seq_len=max_len).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)  # 1e-3 is fine at this size; drop to 3e-4 if it diverges
    loss_criterion = nn.CrossEntropyLoss()

    # --- Train ---
    EPOCHS = 5
    n_batches = len(dataloader)
    print("[TRAIN] Starting pre-training pass...", flush=True)
    model.train()
    for epoch in range(EPOCHS):
        total_loss = 0
        for step, (x, y) in enumerate(dataloader, start=1):
            x, y = x.to(device), y.to(device)
            optimizer.zero_grad()
            logits = model(x)
            # flatten (B, T, vocab) -> (B*T, vocab) so loss covers every position
            loss = loss_criterion(logits.view(-1, logits.size(-1)), y.view(-1))
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

            # flush so progress shows even when piped to a file (stdout buffering)
            if step % 20 == 0 or step == n_batches:
                print(f"[TRAIN] Epoch {epoch+1}/{EPOCHS} | batch {step}/{n_batches} "
                      f"| loss {loss.item():.4f}", flush=True)

        print(f"[TRAIN] Epoch {epoch+1}/{EPOCHS} done | avg CE loss "
              f"{total_loss/n_batches:.4f}", flush=True)

    # --- Generate (quick sanity check) ---
    prompt_text = "Invalid user"  # a phrase that actually appears in the corpus
    prompt_tokens = tokenizer.encode(prompt_text).ids
    input_tensor = torch.tensor([prompt_tokens], dtype=torch.long, device=device)

    generated_ids = model.generate(input_tensor, max_new_tokens=6, temperature=0.5)
    decoded_output = tokenizer.decode(generated_ids[0].tolist())

    print(f"\n[GEN] Seed  : {prompt_text}")
    print(f"[GEN] Output: {decoded_output}")

    # --- Anomaly detection ---
    with open(LOG_FILE, "r") as f:
        corpus_lines = [ln.strip() for ln in f if ln.strip()]

    threshold, mean, std, scored = fit_anomaly_threshold(
        model, tokenizer, corpus_lines, device, max_seq_len=max_len, k=2.0
    )
    print(f"\n[ANOMALY] Score stats: mean={mean:.3f} std={std:.3f} | "
          f"threshold (mean+2*std)={threshold:.3f}", flush=True)

    print("[ANOMALY] Top 5 most anomalous lines in the corpus:")
    for loss, ppl, line in scored[:5]:
        flag = "ANOMALY" if loss > threshold else "normal "
        print(f"  [{flag}] loss={loss:.3f} ppl={ppl:8.1f} | {line[:90]}")

    # A couple of real lines + one injected off-distribution line (should score huge)
    print("[ANOMALY] Predicting on new lines:")
    test_lines = [
        "Dec 10 06:55:46 LabSZ sshd[24200]: Invalid user webmaster from 173.234.31.186",
        "Failed password for root from 112.95.230.3 port 49204 ssh2",
        "kjshdf!! TOTALLY RANDOM gibberish $$$ 999 not-a-real-log ~~~~",
    ]
    for line in test_lines:
        result = score_line(model, tokenizer, line, device, max_seq_len=max_len)
        if result is None:
            print(f"  [skipped: too short] {line}")
            continue
        loss, ppl = result
        flag = "ANOMALY" if loss > threshold else "normal "
        print(f"  [{flag}] loss={loss:.3f} ppl={ppl:8.1f} | {line[:90]}")

    # --- Save weights + threshold so detect_anomalies.py can run standalone ---
    MODEL_FILE = "model/logsentry_lm.pt"
    os.makedirs(os.path.dirname(MODEL_FILE), exist_ok=True)
    torch.save({
        "model_state": model.state_dict(),
        "vocab_size": tokenizer.get_vocab_size(),
        "max_seq_len": max_len,
        "threshold": threshold,
        "score_mean": mean,
        "score_std": std,
    }, MODEL_FILE)
    print(f"\n[SAVE] Saved model + anomaly threshold to {MODEL_FILE}")