"""
tui.py
======
Terminal-based Anomaly Detection Tool for REST API Traffic.

Reads requests from a file or stdin (live stream mode),
processes them through the trained LSTM-Autoencoder,
and displays real-time anomaly alerts in the terminal.

Usage
-----
# File mode — CSIC 2010
python src/tui.py --dataset csic2010 \
    --input data/raw/csic2010/anomalousTrafficTest.txt

# File mode — CIC-IDS2018
python src/tui.py --dataset cicids2018 \
    --input data/raw/cicids2018full/02-22-2018.csv --delay 0.01

# File mode — UNSW-NB15
python src/tui.py --dataset unsw \
    --input "data/raw/unswnb15/Training and Testing Sets/UNSW_NB15_testing-set.csv"

# Live stdin mode — type requests interactively (csic2010 only)
python src/tui.py --dataset csic2010
# Then type: GET /tienda1/index.jsp  (press Enter after each line)

# Pipe stdin from another process
cat requests.txt | python src/tui.py --dataset csic2010

Author : K.A.D.S.D. Kandanaarachchi (2020/ICT/19)
Project: Detecting Anomalous REST API Traffic — IT4216
"""

import argparse
import json
import os
import re
import signal
import sys
import textwrap
import time
from collections import deque
from pathlib import Path
from urllib.parse import parse_qs, unquote, urlparse

import numpy as np
import torch

# ── Add src/ to path so model.py is importable ───────────────────────────────
sys.path.insert(0, str(Path(__file__).parent))
from model import build_model_cicids2018, build_model_csic2010, build_model_unsw

# ═══════════════════════════════════════════════════════════════════════════════
# ANSI COLOURS
# ═══════════════════════════════════════════════════════════════════════════════


class C:
    RESET = "\033[0m"
    BOLD = "\033[1m"
    DIM = "\033[2m"
    RED = "\033[91m"
    GREEN = "\033[92m"
    YELLOW = "\033[93m"
    CYAN = "\033[96m"
    WHITE = "\033[97m"


USE_COLOR = hasattr(sys.stdout, "isatty") and sys.stdout.isatty()


def col(color, text):
    return f"{color}{text}{C.RESET}" if USE_COLOR else text


def bold(text):
    return col(C.BOLD, text)


def dim(text):
    return col(C.DIM, text)


def red(text):
    return col(C.RED, text)


def green(text):
    return col(C.GREEN, text)


def yellow(text):
    return col(C.YELLOW, text)


def cyan(text):
    return col(C.CYAN, text)


# ═══════════════════════════════════════════════════════════════════════════════
# DISPLAY HELPERS
# ═══════════════════════════════════════════════════════════════════════════════

WIDTH = 72


def print_header(
    dataset: str,
    threshold: float,
    model_params: int,
    live: bool = False,
    raw_order: bool = False,
    normal_count: int | None = None,
    attack_count: int | None = None,
):
    # Dataset composition — from the original dataset papers
    DATASET_INFO = {
        "csic2010": {"label": "CSIC 2010", "normal": 72_000, "attack": 25_065},
        "cicids2018": {"label": "CIC-IDS2018", "normal": 2_096_222, "attack": 928},
        "unsw": {"label": "UNSW-NB15", "normal": 93_000, "attack": 164_673},
    }
    info = DATASET_INFO.get(dataset, {})
    n_normal = normal_count if normal_count is not None else info.get("normal", 0)
    n_attack = attack_count if attack_count is not None else info.get("attack", 0)

    mode = "LIVE STDIN" if live else "FILE STREAM"
    print()
    print(col(C.CYAN, "═" * WIDTH))
    print(
        col(
            C.CYAN + C.BOLD,
            "  REST API Anomaly Detection — Live Stream Monitor".center(WIDTH),
        )
    )
    print(col(C.CYAN, "═" * WIDTH))
    print(f"  {bold('Dataset')}   : {info.get('label', dataset.upper())}")
    print(f"  {bold('Mode')}      : {mode}")
    print(f"  {bold('Model')}     : LSTM-Autoencoder  ({model_params:,} parameters)")
    print(
        f"  {bold('Threshold')} : {threshold:.6f}  (95th percentile of normal errors)"
    )
    # Dataset composition
    if n_normal or n_attack:
        total = n_normal + n_attack
        print(col(C.CYAN, "─" * WIDTH))
        print(f"  {bold('Dataset composition')} ({total:,} total records):")
        print(
            f"    {green('Normal')}  : {n_normal:>10,}  ({n_normal / total * 100:.1f}%)"
        )
        print(
            f"    {red('Attack')}  : {n_attack:>10,}  ({n_attack / total * 100:.1f}%)"
        )
    print(col(C.CYAN, "─" * WIDTH))
    print(
        f"  {green('■')} = Normal    {red('■')} = ANOMALY    {yellow('■')} = Uncertain"
    )
    print(col(C.CYAN, "─" * WIDTH))
    if live:
        print(
            col(C.YELLOW, "  Type requests as:  METHOD /path  (e.g. GET /api/users/1)")
        )
        print(col(C.YELLOW, "  Press Ctrl+C to stop and print summary."))
        print(col(C.CYAN, "─" * WIDTH))
    if dataset == "cicids2018" and not raw_order:
        print(
            col(
                C.GREEN,
                "  Sessions are built from Timestamp-sorted flows, matching\n"
                "  preprocessing.py's session construction. Pass --raw-order\n"
                "  to replay strict on-disk row order instead (will not match\n"
                "  evaluate.py's held-out test split either way -- this is a\n"
                "  live-demo illustration, not a re-measurement).",
            )
        )
        print(col(C.CYAN, "─" * WIDTH))
    elif dataset in ("cicids2018", "unsw"):
        print(
            col(
                C.YELLOW,
                "  NOTE: replaying in RAW on-disk row order (--raw-order set,\n"
                "  or no Timestamp column available for this dataset). Flows\n"
                "  adjacent on disk may not be temporally adjacent, so 5-flow\n"
                "  sessions built here may never have occurred as a real\n"
                "  sequence during training -- elevated error here can reflect\n"
                "  that, not a detection failure. Anomaly rates will not match\n"
                "  the formal FPR in evaluation_"
                + dataset
                + ".json; use evaluate.py for the authoritative metric.",
            )
        )
        print(col(C.CYAN, "─" * WIDTH))
    print()


def print_request(idx: int, error: float, threshold: float, display: str):
    margin = threshold * 0.15
    is_alert = error > threshold
    is_unc = not is_alert and error > (threshold - margin)

    if is_alert:
        status = red("  ANOMALY ")
        bar_c = C.RED
    elif is_unc:
        status = yellow("UNCERTAIN ")
        bar_c = C.YELLOW
    else:
        status = green("   NORMAL ")
        bar_c = C.GREEN

    max_bar = 20
    ratio = min(error / (threshold * 3), 1.0)
    filled = int(ratio * max_bar)
    bar = col(bar_c, "█" * filled) + dim("░" * (max_bar - filled))
    label_s = (display[:38] + "…") if len(display) > 39 else display.ljust(39)

    print(f"  [{idx:>5}] {status}  err={error:.5f}  {bar}  {dim(label_s)}")
    return is_alert, is_unc


def print_alert_banner(idx: int, display: str, error: float, threshold: float):
    if error < threshold * 1.5:
        return
    print()
    print(col(C.RED, "  ┌" + "─" * (WIDTH - 4) + "┐"))
    print(
        col(C.RED, "  │")
        + col(C.RED + C.BOLD, f" ANOMALY DETECTED — Session #{idx}".ljust(WIDTH - 4))
        + col(C.RED, "│")
    )
    print(
        col(C.RED, "  │")
        + f"  Error: {error:.6f}  "
        f"(threshold: {threshold:.6f}  ×{error / threshold:.1f})".ljust(WIDTH - 4)
        + col(C.RED, "│")
    )
    trunc = (display[: WIDTH - 10] + "…") if len(display) > WIDTH - 10 else display
    print(col(C.RED, "  │") + f"  {trunc}".ljust(WIDTH - 4) + col(C.RED, "│"))
    print(col(C.RED, "  └" + "─" * (WIDTH - 4) + "┘"))
    print()


def print_summary(
    total: int,
    alerts: int,
    uncertain: int,
    threshold: float,
    elapsed: float,
    top_anomalies: list,
):
    print()
    print(col(C.CYAN, "═" * WIDTH))
    print(col(C.CYAN + C.BOLD, "  SESSION SUMMARY".center(WIDTH)))
    print(col(C.CYAN, "═" * WIDTH))
    print(f"  {bold('Total sessions processed')} : {total:,}")
    if total > 0:
        print(
            f"  {bold('Anomalies detected')}        : "
            f"{red(str(alerts))}  ({alerts / total * 100:.1f}%)"
        )
    print(f"  {bold('Uncertain sessions')}        : {yellow(str(uncertain))}")
    print(
        f"  {bold('Normal sessions')}           : "
        f"{green(str(max(0, total - alerts - uncertain)))}"
    )
    print(f"  {bold('Processing time')}           : {elapsed:.2f}s")
    if total > 0 and elapsed > 0:
        print(
            f"  {bold('Throughput')}               : {total / elapsed:.0f} sessions/sec"
        )

    if top_anomalies:
        print()
        print(col(C.RED, "  Top anomalous sessions:"))
        print(col(C.RED, "  " + "─" * 50))
        for rank, (err, idx, disp) in enumerate(top_anomalies[:5], 1):
            trunc = (disp[:45] + "…") if len(disp) > 46 else disp
            print(f"  {rank}. [{idx:>5}] err={err:.5f}  {dim(trunc)}")
    print(col(C.CYAN, "═" * WIDTH))
    print()


# ═══════════════════════════════════════════════════════════════════════════════
# URL ABSTRACTION  (shared with preprocessing.py)
# ═══════════════════════════════════════════════════════════════════════════════

_SUSPICIOUS = {
    "'",
    '"',
    "<",
    ">",
    ";",
    "--",
    "DROP",
    "SELECT",
    "UNION",
    "INSERT",
    "DELETE",
    "SCRIPT",
    "ALERT",
    "../",
    "%27",
    "%3C",
    "%3E",
}


def _abstract_url(url: str) -> str:
    """
    Must stay byte-for-byte identical to preprocessing.py's _abstract_url().
    The vocabulary in vocab_csic2010_w{window}.json was built from
    preprocessing.py's output — any divergence here causes normal requests
    to be abstracted into templates absent from the vocabulary, mapping
    them to <UNK> and inflating reconstruction error on legitimate
    traffic (false positives).
    """
    if not url:
        return "<EMPTY>"
    try:
        parsed = urlparse(url)

        # Abstract numeric path segments
        abstracted_parts = []
        for part in parsed.path.split("/"):
            if re.match(r"^\d+$", part):
                abstracted_parts.append("{INT}")
            elif re.match(r"^[0-9a-f-]{32,}$", part, re.I):
                abstracted_parts.append("{HASH}")
            else:
                abstracted_parts.append(part)
        abs_path = "/".join(abstracted_parts)

        if not parsed.query:
            return abs_path

        # Abstract query parameter values
        abs_params = []
        for param in parsed.query.split("&"):
            if "=" not in param:
                abs_params.append(param)
                continue
            key, _, value = param.partition("=")
            decoded = unquote(value)
            if re.match(r"^\d+$", decoded):
                abs_params.append(f"{key}={{INT}}")
            elif re.match(r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+$", decoded):
                abs_params.append(f"{key}={{EMAIL}}")
            elif len(decoded) == 0:
                abs_params.append(f"{key}={{EMPTY}}")
            elif any(s in decoded.upper() for s in _SUSPICIOUS):
                abs_params.append(f"{key}={{INJECT}}")
            else:
                abs_params.append(f"{key}={{STR}}")

        return f"{abs_path}?{'&'.join(abs_params)}"

    except Exception:
        return url


# ═══════════════════════════════════════════════════════════════════════════════
# STREAM READERS
# ═══════════════════════════════════════════════════════════════════════════════

_VALID_METHODS = {"GET", "POST", "PUT", "DELETE", "HEAD", "OPTIONS", "PATCH"}


def stream_csic2010(filepath: Path):
    """
    Yield (abstracted_url, display, raw_request) from a CSIC 2010 HTTP
    log file. Requests are separated by blank lines in the raw file.

    raw_request is the UNTRUNCATED "METHOD /path" string (unlike
    `display`, which is truncated to 60 chars for terminal width).
    It exists so that --log runs can record exactly which request
    formed part of a session, for exact-match ground-truth comparison
    against a real attack file via evaluate_stream_log.py — truncated
    strings risk false prefix-matches when two different attack
    requests share the same first 60 characters.
    """
    with open(filepath, "r", encoding="latin-1") as fh:
        content = fh.read()

    for raw in content.strip().split("\n\n"):
        lines = raw.strip().split("\n")
        if not lines or not lines[0].strip():
            continue
        parts = lines[0].strip().split(" ")
        method = parts[0] if parts else ""
        url = parts[1] if len(parts) > 1 else ""
        if method not in _VALID_METHODS:
            continue
        abstracted = _abstract_url(url)
        display = f"{method} {url[:60]}"
        raw_request = f"{method} {url}"
        # line_idx is None here: this reader consumes the raw CSIC
        # dataset files directly, not a mix_stream.py output file, so
        # there is no "<output>_labels.json" sibling to align against.
        yield abstracted, display, raw_request, None


def stream_stdin():
    """
    Read HTTP requests from stdin line by line (live stream mode).

    Accepts lines in the form:
        METHOD /path?param=value
        /path?param=value          (GET assumed)

    Yields (abstracted_url, display, raw_request)
    """
    try:
        for line in sys.stdin:
            line = line.strip()
            if not line or line.startswith("#"):
                continue

            parts = line.split(" ", 1)
            if len(parts) == 2 and parts[0].upper() in _VALID_METHODS:
                method = parts[0].upper()
                url = parts[1]
            else:
                method = "GET"
                url = line

            abstracted = _abstract_url(url)
            display = f"{method} {url[:60]}"
            raw_request = f"{method} {url}"
            yield abstracted, display, raw_request, None

    except (EOFError, KeyboardInterrupt):
        return


def stream_line_file(filepath: Path):
    """
    Read HTTP requests from a one-per-line text file (mix_stream.py output).

    Accepts the format:
        METHOD /path?param=value

    Skips blank lines and lines starting with # (comments).
    Used when --input points at a mixed-stream file rather than a raw CSIC log.

    Yields (abstracted_url, display, raw_request, line_idx).

    line_idx is the 0-based index of this line among DATA lines only
    (comments/blanks excluded) — i.e. it lines up exactly with the
    ground-truth labels array mix_stream.py writes to
    "<output>_labels.json" when run with --save. Use line_idx rather
    than text-matching raw_request against anomalousTrafficTest.txt:
    many attack-file lines are themselves ordinary navigation (images,
    CSS) that also appears verbatim in the normal corpus, so text
    matching alone produces false ground-truth positives.
    """
    with open(filepath, "r", encoding="utf-8", errors="replace") as fh:
        line_idx = 0
        for line in fh:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split(" ", 1)
            if len(parts) == 2 and parts[0].upper() in _VALID_METHODS:
                method = parts[0].upper()
                url = parts[1]
            else:
                # bare path — assume GET
                method = "GET"
                url = line
            abstracted = _abstract_url(url)
            display = f"{method} {url[:60]}"
            raw_request = f"{method} {url}"
            yield abstracted, display, raw_request, line_idx
            line_idx += 1


# ── Feature lists for flow datasets ──────────────────────────────────────────

CICIDS_FEATURES = [
    "Tot Fwd Pkts",
    "Tot Bwd Pkts",
    "TotLen Fwd Pkts",
    "TotLen Bwd Pkts",
    "Pkt Len Mean",
    "Pkt Len Std",
    "Pkt Len Max",
    "Pkt Size Avg",
    "Flow Duration",
    "Flow IAT Mean",
    "Flow IAT Std",
    "Fwd Pkts/s",
    "Bwd Pkts/s",
    "Flow Byts/s",
    "Flow Pkts/s",
    "SYN Flag Cnt",
    "RST Flag Cnt",
    "PSH Flag Cnt",
    "ACK Flag Cnt",
    "Init Fwd Win Byts",
    "Init Bwd Win Byts",
    "Dst Port",
    "Protocol",
]

UNSW_FEATURES = [
    "dur",
    "spkts",
    "dpkts",
    "sbytes",
    "dbytes",
    "rate",
    "sload",
    "dload",
    "smean",
    "dmean",
    "ct_state_ttl",
    "ct_dst_ltm",
    "ct_src_dport_ltm",
    "ct_dst_sport_ltm",
    "ct_dst_src_ltm",
    "proto",
    "service",
    "state",
    "sinpkt",
    "dinpkt",
]


def stream_flows(
    filepath: Path, dataset: str, scaler_path: Path, raw_order: bool = False
):
    """
    Yield (feature_vector, display) from a CIC-IDS2018 or UNSW-NB15 CSV.

    SESSION-ORDER NOTE: by default this function sorts the file by the
    'Timestamp' column (CIC-IDS2018 only -- UNSW-NB15's pre-split CSVs
    have no per-row timestamp) before constructing the 5-flow sliding
    window, mirroring preprocess_cicids2018() in preprocessing.py
    exactly (same dayfirst=True, same errors='coerce').

    This matters because the model was trained on sessions built from
    TEMPORALLY ADJACENT flows. Reading a CSV in raw on-disk row order
    (CICFlowMeter writes flows on completion, not on capture start) can
    group flows that were never actually close together in time into
    one "session" -- producing inflated reconstruction error that
    reflects an artifact of file layout, not a real anomaly. Sorting
    first reconstructs genuine temporal sequences, consistent with how
    the model was trained and how evaluate.py builds its test array.

    Pass raw_order=True to disable sorting and replay strictly in
    on-disk row order instead (e.g. to demonstrate the difference, or
    to simulate a true unsorted live-arrival stream).

    NOTE ON FILL STRATEGY: preprocessing.py fills missing/infinite values
    with the column MEDIAN computed over normal training data before
    fitting the scaler. That per-feature median is not persisted in
    scaler_*.json (only mean/scale are), so it is not recoverable here
    exactly. Filling with the scaler's stored MEAN instead is the
    closest safe approximation: after scaling, (mean - mean) / scale = 0,
    so a missing value contributes neutrally to the reconstruction error
    rather than being pulled toward an arbitrary, scale-dependent offset
    the way fillna(0) does.
    """
    import pandas as pd

    scaler_data = json.load(open(scaler_path))
    mean = np.array(scaler_data["mean"])
    scale = np.array(scaler_data["scale"])
    feat_cols = CICIDS_FEATURES if dataset == "cicids2018" else UNSW_FEATURES
    n_features = len(mean)

    def _clean_and_scale(chunk):
        """Shared cleaning/scaling logic, applied to any chunk/frame."""
        for lc in ("Label", "label"):
            if lc in chunk.columns:
                chunk = chunk[chunk[lc] != lc]

        available = [f for f in feat_cols if f in chunk.columns]
        if not available:
            return None, None

        for col_name in available:
            chunk[col_name] = pd.to_numeric(chunk[col_name], errors="coerce")

        if dataset == "unsw":
            import pandas as _pd

            for cat in ["proto", "service", "state"]:
                if cat in chunk.columns and chunk[cat].dtype == object:
                    chunk[cat] = _pd.Categorical(chunk[cat]).codes.astype(float)

        chunk[available] = chunk[available].replace([np.inf, -np.inf], np.nan)
        fill_values = {
            col: float(mean[idx])
            for idx, col in enumerate(feat_cols)
            if col in available and idx < len(mean)
        }
        chunk[available] = chunk[available].fillna(fill_values)
        chunk[available] = chunk[available].fillna(0)  # safety net only

        X_raw = chunk[available].values.astype(np.float32)
        n_avail = min(X_raw.shape[1], n_features)
        X_scaled = (X_raw[:, :n_avail] - mean[:n_avail]) / (scale[:n_avail] + 1e-9)
        X_scaled = np.nan_to_num(X_scaled, nan=0.0, posinf=0.0, neginf=0.0)

        if X_scaled.shape[1] < n_features:
            pad = np.zeros((X_scaled.shape[0], n_features - X_scaled.shape[1]))
            X_scaled = np.concatenate([X_scaled, pad], axis=1)

        return chunk, X_scaled

    def _yield_rows(chunk, X_scaled):
        for i in range(len(X_scaled)):
            row = chunk.iloc[i]
            if dataset == "cicids2018":
                label = str(row.get("Label", "?"))
                dst = str(row.get("Dst Port", "?"))
                proto = str(row.get("Protocol", "?"))
                display = f"Port={dst} Proto={proto} Label={label}"
            else:
                label = str(row.get("label", row.get("Label", "?")))
                proto = str(row.get("proto", "?"))
                service = str(row.get("service", "?"))
                display = f"Proto={proto} Service={service} Label={label}"
            yield X_scaled[i], display

    can_sort = (not raw_order) and dataset == "cicids2018"

    if can_sort:
        # Load the full file, sort by Timestamp (same as preprocessing.py),
        # then stream from the sorted frame. CIC-IDS2018 day files are a
        # few hundred MB -- comfortably fits in memory on 16GB+ RAM.
        full = pd.read_csv(filepath, low_memory=False)
        for lc in ("Label", "label"):
            if lc in full.columns:
                full = full[full[lc] != lc]
        if "Timestamp" in full.columns:
            full["Timestamp"] = pd.to_datetime(
                full["Timestamp"], dayfirst=True, errors="coerce"
            )

            # A small number of rows in the CIC-IDS2018 day files carry
            # corrupted-but-parseable timestamps (e.g. "10/01/1970"),
            # a known CICFlowMeter capture artifact -- these also carry
            # nonsensical negative Flow Duration values. Sorting WITHOUT
            # filtering these out clusters them all at the very front of
            # the stream (1970 < 2018), so the first 1-2 demo sessions
            # would always be built from these malformed flows instead
            # of real traffic. preprocessing.py never hits this because
            # it sorts the full multi-million-row normal pool, where a
            # handful of misplaced rows don't dominate any single
            # 5-session window the way they do in a 200-session demo
            # slice. Drop rows with a timestamp before the dataset's own
            # capture date (Feb 2018) before sorting/windowing.
            valid_mask = full["Timestamp"] >= "2018-01-01"
            n_dropped = int((~valid_mask).sum())
            if n_dropped:
                print(
                    col(
                        C.DIM,
                        f"  [stream_flows] Dropped {n_dropped} row(s) with "
                        f"corrupted/pre-2018 timestamps before sorting "
                        f"(known CICFlowMeter capture artifact).",
                    )
                )
            full = full[valid_mask]

            full = full.sort_values("Timestamp").reset_index(drop=True)
        else:
            raw_order = True  # no Timestamp column -- fall back silently

        if "Timestamp" in full.columns:
            for start in range(0, len(full), 500):
                chunk = full.iloc[start : start + 500].copy()
                chunk, X_scaled = _clean_and_scale(chunk)
                if chunk is None:
                    continue
                yield from _yield_rows(chunk, X_scaled)
            return

    # raw_order=True, UNSW (no per-row timestamp), or no Timestamp column:
    # stream in on-disk order, chunked, as before.
    for chunk in pd.read_csv(filepath, low_memory=False, chunksize=500):
        chunk, X_scaled = _clean_and_scale(chunk)
        if chunk is None:
            continue
        yield from _yield_rows(chunk, X_scaled)


# ═══════════════════════════════════════════════════════════════════════════════
# MAIN DETECTION ENGINE
# ═══════════════════════════════════════════════════════════════════════════════


def run_detection(
    dataset: str,
    input_path: Path | None,
    data_dir: Path,
    model_dir: Path,
    delay: float = 0.05,
    threshold_override: float = None,
    top_n: int = 5,
    max_rows: int = None,
    raw_order: bool = False,
    window_size: int = 5,
    log_path: Path | None = None,
):
    # ── Validate stdin mode ───────────────────────────────────────────
    if input_path is None and dataset != "csic2010":
        print(red("\n  Stdin mode is only supported for --dataset csic2010."))
        print(red("  Flow datasets (cicids2018, unsw) require a CSV file"))
        print(red("  with pre-extracted network flow features.\n"))
        sys.exit(1)

    # run_id mirrors preprocessing.py/train.py/evaluate.py/active_learning.py:
    # every artifact on disk is now window-suffixed (e.g. "cicids2018_w5")
    # so multiple --window runs can coexist without overwriting each other.
    run_id = f"{dataset}_w{window_size}"

    # ── Load threshold ────────────────────────────────────────────────
    thresh_file = model_dir / f"threshold_{run_id}.json"
    if not thresh_file.exists():
        print(red(f"\n  Threshold file not found: {thresh_file}"))
        print(red(f"  Run: python src/train.py --dataset {dataset} --window {window_size}\n"))
        sys.exit(1)

    thresh_data = json.load(open(thresh_file))
    threshold = threshold_override or thresh_data["threshold"]

    # ── Load model checkpoint ─────────────────────────────────────────
    model_file = model_dir / f"best_{run_id}.pt"
    if not model_file.exists():
        print(red(f"\n  Model file not found: {model_file}\n"))
        sys.exit(1)

    checkpoint = torch.load(model_file, map_location="cpu", weights_only=True)

    if dataset == "csic2010":
        vocab_size = checkpoint["embedding.weight"].shape[0]
        vocab_file = data_dir / f"vocab_{run_id}.json"
        if not vocab_file.exists():
            print(red(f"\n  Vocab file not found: {vocab_file}\n"))
            sys.exit(1)
        vocab_data = json.load(open(vocab_file))
        # Back-compat: older runs saved a bare {token: id} dict; the
        # current format wraps it as {"window_size": N, "vocab": {...}}
        # (see preprocess_csic2010() in preprocessing.py).
        vocab = vocab_data.get("vocab", vocab_data)
        # Pad vocab to match checkpoint size if preprocessing was re-run
        while len(vocab) < vocab_size:
            vocab[f"<PLACEHOLDER_{len(vocab)}>"] = len(vocab)
        model = build_model_csic2010(vocab_size=vocab_size, seq_len=window_size)
        vocab_ref = vocab  # keep reference for inference

    elif dataset == "cicids2018":
        model = build_model_cicids2018(
            n_features=len(CICIDS_FEATURES), seq_len=window_size
        )
        vocab_ref = None

    else:
        model = build_model_unsw(n_features=len(UNSW_FEATURES), seq_len=window_size)
        vocab_ref = None

    model.load_state_dict(checkpoint)
    model.eval()

    total_params = sum(p.numel() for p in model.parameters())
    live_mode = input_path is None
    print_header(
        dataset,
        threshold,
        total_params,
        live=live_mode,
        raw_order=raw_order,
    )

    # ── Set up stream ─────────────────────────────────────────────────
    window: deque = deque(maxlen=window_size)
    # Parallel deque of the untruncated raw "METHOD /path" strings that
    # make up the current window. Kept separate from `window` (which
    # holds abstracted tokens / scaled feature vectors used for
    # inference) so that a --log run can record exactly which raw
    # requests formed each session, for later ground-truth comparison
    # against a real attack file via evaluate_stream_log.py.
    window_raw: deque = deque(maxlen=window_size)
    # Parallel deque of source-file line indices (0-based, data lines
    # only), populated only when reading a mix_stream.py-produced file
    # via stream_line_file. Used to align a session against the exact
    # ground-truth labels mix_stream.py wrote to "<output>_labels.json",
    # rather than text-matching against the full attack corpus.
    window_lineidx: deque = deque(maxlen=window_size)

    # ── Set up log file, if requested ────────────────────────────────
    log_fh = None
    if log_path is not None:
        log_path = Path(log_path)
        log_path.parent.mkdir(parents=True, exist_ok=True)
        log_fh = open(log_path, "w", encoding="utf-8")
        print(dim(f"  Logging session-by-session detections to: {log_path}\n"))

    if live_mode:
        stream = stream_stdin()
    elif dataset == "csic2010":
        # Auto-detect: if the file starts with a '#' comment header it was
        # produced by mix_stream.py (one request per line).
        # Raw CSIC HTTP logs start with an HTTP method like "GET " or "POST ".
        _first_line = input_path.read_bytes()[:6].decode("latin-1", errors="replace")
        if _first_line.startswith("#"):
            stream = stream_line_file(input_path)
        else:
            stream = stream_csic2010(input_path)
    else:
        scaler_key = "cicids2018" if dataset == "cicids2018" else "unsw"
        scaler_path = data_dir / f"scaler_{scaler_key}_w{window_size}.json"
        if not scaler_path.exists():
            print(red(f"\n  Scaler file not found: {scaler_path}\n"))
            sys.exit(1)
        stream = stream_flows(input_path, dataset, scaler_path, raw_order=raw_order)

    # ── Detection loop ────────────────────────────────────────────────
    session_idx = 0
    alert_count = 0
    unc_count = 0
    top_anomalies = []
    margin = threshold * 0.15
    start_time = time.perf_counter()

    interrupted = False

    def _on_interrupt(sig, frame):
        nonlocal interrupted
        interrupted = True
        print(yellow("\n\n  Interrupted — printing summary...\n"))

    signal.signal(signal.SIGINT, _on_interrupt)

    try:
        for item in stream:
            if interrupted:
                break
            if max_rows and session_idx >= max_rows:
                break

            # Unpack stream item
            if dataset == "csic2010":
                abstracted, display, raw_request, line_idx = item
                window.append(abstracted)
                window_raw.append(raw_request)
                window_lineidx.append(line_idx)
            else:
                feat_vec, display = item
                window.append(feat_vec)
                window_raw.append(display)  # no meaningful "raw request" for flow data
                window_lineidx.append(None)

            # Wait until we have a full window
            if len(window) < window_size:
                if live_mode:
                    remaining = window_size - len(window)
                    print(
                        col(
                            C.DIM,
                            f"  Building session window… "
                            f"({remaining} more request(s) needed)",
                        )
                    )
                continue

            session_idx += 1

            # ── Inference ─────────────────────────────────────────────
            with torch.no_grad():
                if dataset == "csic2010":
                    tokens = [vocab_ref.get(u, 1) for u in window]
                    tensor = torch.tensor([tokens], dtype=torch.long)
                else:
                    arr = np.stack(list(window), axis=0)
                    arr = np.nan_to_num(arr, nan=0.0)
                    tensor = torch.tensor(arr[None], dtype=torch.float32)

                error = model.reconstruction_error(tensor).item()

            # ── Classify & display ────────────────────────────────────
            is_alert, is_unc = print_request(session_idx, error, threshold, display)

            if is_alert:
                alert_count += 1
                top_anomalies.append((error, session_idx, display))
                top_anomalies.sort(reverse=True)
                top_anomalies = top_anomalies[:top_n]
                print_alert_banner(session_idx, display, error, threshold)
            elif is_unc:
                unc_count += 1

            # ── Log this session, if requested ──────────────────────────
            # Written as one JSON object per line (JSONL) so a run can be
            # interrupted (Ctrl+C) without corrupting already-written
            # records, and so evaluate_stream_log.py can stream-read it
            # without loading the whole file into memory. `requests` is
            # the up-to-`window_size` untruncated raw requests that made
            # up this session, needed to check ground truth against a
            # real attack file after the fact.
            if log_fh is not None:
                verdict = "anomaly" if is_alert else ("uncertain" if is_unc else "normal")
                record = {
                    "session_idx": session_idx,
                    "error": error,
                    "threshold": threshold,
                    "verdict": verdict,
                    "requests": list(window_raw),
                    "line_indices": list(window_lineidx),
                }
                log_fh.write(json.dumps(record) + "\n")
                log_fh.flush()

            if not live_mode:
                time.sleep(delay)

    except StopIteration:
        pass
    finally:
        # Always close the log file, even on Ctrl+C or an unexpected
        # exception mid-stream, so partial runs still leave a valid,
        # readable JSONL file rather than a truncated/unflushed one.
        if log_fh is not None:
            log_fh.close()

    elapsed = time.perf_counter() - start_time
    top_anomalies.sort(reverse=True)

    print_summary(
        total=session_idx,
        alerts=alert_count,
        uncertain=unc_count,
        threshold=threshold,
        elapsed=elapsed,
        top_anomalies=top_anomalies,
    )

    if log_path is not None:
        print(dim(f"  Session log written → {log_path}"))
        print(
            dim(
                "  Cross-check against ground truth with:\n"
                f"    python evaluate_stream_log.py --log {log_path} "
                "--labels <path/to/demo_stream_labels.json>\n"
                "  (the *_labels.json file mix_stream.py --save writes "
                "alongside its output; this is the recommended, exact "
                "way to check ground truth — the older --attack-file "
                "text-matching mode is still available but over-counts "
                "attacks whose lines also appear as ordinary navigation "
                "in the normal corpus)"
            )
        )
        print()


# ═══════════════════════════════════════════════════════════════════════════════
# CLI
# ═══════════════════════════════════════════════════════════════════════════════


def main():
    parser = argparse.ArgumentParser(
        description="REST API Anomaly Detection — Live Stream Monitor",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=textwrap.dedent("""
        Examples:
          # File mode
          python src/tui.py --dataset csic2010 \\
              --input data/raw/csic2010/anomalousTrafficTest.txt --delay 0.02

          python src/tui.py --dataset cicids2018 \\
              --input data/raw/cicids2018full/02-22-2018.csv --delay 0.01

          python src/tui.py --dataset unsw \\
              --input "data/raw/unswnb15/Training and Testing Sets/UNSW_NB15_testing-set.csv"

          # Live stdin mode (csic2010 only — type requests interactively)
          python src/tui.py --dataset csic2010

          # Pipe from another process
          cat my_requests.txt | python src/tui.py --dataset csic2010
        """),
    )

    parser.add_argument(
        "--dataset",
        required=True,
        choices=["csic2010", "cicids2018", "unsw"],
        help="Which trained model to use",
    )
    parser.add_argument(
        "--input",
        default=None,
        help="Path to input file (HTTP log or CSV). "
        "Omit to read from stdin (csic2010 only).",
    )
    parser.add_argument(
        "--delay",
        type=float,
        default=0.05,
        help="Seconds between sessions in file mode (default 0.05)",
    )
    parser.add_argument(
        "--threshold",
        type=float,
        default=None,
        help="Override the saved anomaly threshold",
    )
    parser.add_argument(
        "--top",
        type=int,
        default=5,
        help="Number of top anomalies to show in summary (default 5)",
    )
    parser.add_argument(
        "--max",
        type=int,
        default=None,
        help="Maximum sessions to process (default: all)",
    )
    parser.add_argument(
        "--base",
        default=".",
        help="Base project directory (default: current dir)",
    )
    parser.add_argument(
        "--raw-order",
        action="store_true",
        help="For flow datasets (cicids2018), replay the CSV in raw "
        "on-disk row order instead of sorting by Timestamp first. "
        "Default is sorted (matches preprocessing.py's session "
        "construction). Use this flag to demonstrate the difference "
        "or to simulate a genuinely unsorted input stream.",
    )
    parser.add_argument(
        "--window",
        type=int,
        default=5,
        help="Sliding window size (must match trained model)",
    )
    parser.add_argument(
        "--log",
        default=None,
        metavar="FILE",
        help="Write a JSONL record per session (index, reconstruction "
        "error, threshold, verdict, and the raw requests in that "
        "session's window) to FILE. Intended for comparing a live-demo "
        "run (e.g. against mix_stream.py output) against a real attack "
        "file afterward with evaluate_stream_log.py, to compute TPR/FPR/"
        "precision/F1 for that specific run.",
    )
    args = parser.parse_args()

    base = Path(args.base)
    data_dir = base / "data" / "processed"
    model_dir = base / "models"

    # Resolve input path
    if args.input is None:
        input_path = None  # stdin mode
    else:
        input_path = Path(args.input)
        if not input_path.exists():
            print(red(f"\n  Input file not found: {input_path}\n"))
            sys.exit(1)

    run_detection(
        dataset=args.dataset,
        input_path=input_path,
        data_dir=data_dir,
        model_dir=model_dir,
        delay=args.delay,
        threshold_override=args.threshold,
        top_n=args.top,
        max_rows=args.max,
        raw_order=args.raw_order,
        window_size=args.window,
        log_path=args.log,
    )


if __name__ == "__main__":
    main()