server.py

#!/usr/bin/env python3
"""ROGLAG demo server.

Stdlib-only HTTP server for a small security-log RAG demo:
- loads synthetic JSONL logs
- exposes dashboard/search/Q&A APIs
- uses a local retrieval + rule explanation fallback
- optionally calls OpenRouter when OPENROUTER_API_KEY is set
"""
from __future__ import annotations

import json
import heapq
import hashlib
import csv
import io
import math
import os
import re
import sqlite3
import sys
import time
import urllib.error
import urllib.request
from array import array
from collections import Counter, defaultdict
from dataclasses import dataclass
from datetime import datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any
from urllib.parse import parse_qs, urlparse

try:
    from app.generator import generate as gen_records, generate_live, to_csv
    GENERATOR_OK = True
except Exception:  # generator is optional; core demo still runs without it
    GENERATOR_OK = False


ROOT = Path(__file__).resolve().parent
WEB_DIR = ROOT / "web"

# Severity fallback when an ingested row has no severity column.
CATEGORY_DEFAULT_SEVERITY = {
    "benign": "info", "brute_force": "high", "port_scan": "medium",
    "web_attack": "high", "priv_esc": "critical", "data_exfil": "critical",
    "c2_beacon": "critical",
}


def resolve_project_path(env_name: str, default: Path) -> Path:
    value = os.getenv(env_name)
    if not value:
        return default
    path = Path(value)
    return path if path.is_absolute() else ROOT / path


def load_dotenv(path: Path = ROOT / ".env") -> None:
    if not path.exists():
        return
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip().strip('"').strip("'")
        if key and key not in os.environ:
            os.environ[key] = value


load_dotenv()

DATA_FILE = resolve_project_path("LOGRAG_DATA_FILE", ROOT / "Data" / "files" / "security_logs.jsonl")
VECTOR_FILE = resolve_project_path("LOGRAG_VECTOR_FILE", ROOT / "Data" / "files" / "lograg_vectors.sqlite")

DEFAULT_MODEL = os.getenv("OPENROUTER_MODEL", "openrouter/free")
DEFAULT_EMBED_MODEL = os.getenv("OPENROUTER_EMBED_MODEL", "nvidia/llama-nemotron-embed-vl-1b-v2:free")

# Qdrant (production vector RAG). Set QDRANT_URL to enable; empty = use local fallback.
QDRANT_URL = os.getenv("QDRANT_URL", "").rstrip("/")
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY", "")
QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "roglag_logs")

CATEGORY_META = {
    "benign": {
        "label": "ปกติ",
        "mitre": "-",
        "risk": 10,
        "runbook": "บันทึกไว้เป็น baseline และใช้เทียบพฤติกรรมผิดปกติในช่วงเวลาเดียวกัน",
    },
    "brute_force": {
        "label": "Brute force",
        "mitre": "T1110",
        "risk": 78,
        "runbook": "ตรวจ IP ต้นทาง, lock account ชั่วคราว, เพิ่ม MFA, block IP ถ้ามี failed login ถี่ผิดปกติ",
    },
    "port_scan": {
        "label": "Port scan",
        "mitre": "T1046",
        "risk": 62,
        "runbook": "ตรวจ firewall drop burst, จำกัด exposure ของ port สำคัญ, เพิ่ม deny rule สำหรับ source ที่ scan ซ้ำ",
    },
    "web_attack": {
        "label": "Web attack",
        "mitre": "T1190",
        "risk": 82,
        "runbook": "ตรวจ request payload, WAF rule, access log รอบเวลาเดียวกัน และ patch endpoint ที่รับ parameter เสี่ยง",
    },
    "priv_esc": {
        "label": "Privilege escalation",
        "mitre": "T1068/T1548",
        "risk": 94,
        "runbook": "ตรวจ sudo command, revoke session, review sudoers, rotate credential และดู command history",
    },
    "data_exfil": {
        "label": "Data exfiltration",
        "mitre": "T1041",
        "risk": 96,
        "runbook": "ตัด connection ปลายทาง, ตรวจปริมาณ bytes, หา process เจ้าของ traffic และ preserve evidence",
    },
    "c2_beacon": {
        "label": "C2 beacon",
        "mitre": "T1071",
        "risk": 95,
        "runbook": "isolate host, ตรวจ interval คงที่, block C2 destination และเก็บ memory/network artifact",
    },
}

TOKEN_RE = re.compile(r"[a-zA-Z0-9_.:/-]+|[\u0E00-\u0E7F]+")
SQLI_RE = re.compile(r"(union\s+select|drop\s+table|or\s+'?1'?='?1|sleep\()", re.I)
XSS_RE = re.compile(r"(<script|onerror=|document\.cookie)", re.I)

QUERY_EXPANSIONS = {
    "brute_force": ["brute_force", "failed", "password", "login_failed", "ssh"],
    "port_scan": ["port_scan", "drop", "syn", "dpt", "scan"],
    "web_attack": ["web_attack", "sqli", "xss", "union", "select", "script", "payload"],
    "priv_esc": ["priv_esc", "sudo", "root", "shadow", "passwd"],
    "data_exfil": ["data_exfil", "outbound", "bytes", "exfil", "dst"],
    "c2_beacon": ["c2_beacon", "c2", "beacon", "conn", "interval"],
}

QUERY_HINTS = {
    "brute": "brute_force",
    "bruteforce": "brute_force",
    "failed login": "brute_force",
    "เดารหัส": "brute_force",
    "ล็อกอิน": "brute_force",
    "scan": "port_scan",
    "port": "port_scan",
    "สแกน": "port_scan",
    "sqli": "web_attack",
    "sql": "web_attack",
    "xss": "web_attack",
    "payload": "web_attack",
    "เว็บ": "web_attack",
    "sudo": "priv_esc",
    "root": "priv_esc",
    "privilege": "priv_esc",
    "ยกระดับ": "priv_esc",
    "exfil": "data_exfil",
    "outbound": "data_exfil",
    "ข้อมูลรั่ว": "data_exfil",
    "ส่งข้อมูล": "data_exfil",
    "c2": "c2_beacon",
    "beacon": "c2_beacon",
    "interval": "c2_beacon",
    "คงที่": "c2_beacon",
}


def tokenize(text: str) -> list[str]:
    return [t.lower() for t in TOKEN_RE.findall(text or "") if len(t) > 1]


def expand_query(query: str) -> tuple[list[str], set[str]]:
    lowered = query.lower()
    categories = {cat for cat in CATEGORY_META if cat in lowered}
    for hint, category in QUERY_HINTS.items():
        if hint in lowered:
            categories.add(category)

    terms = tokenize(query.replace("_", " "))
    for category in categories:
        terms.extend(QUERY_EXPANSIONS.get(category, []))
    return terms, categories


def parse_csv_text(text: str) -> list[dict[str, Any]]:
    reader = csv.DictReader(io.StringIO(text or ""))
    rows: list[dict[str, Any]] = []
    for raw in reader:
        row = {(k or "").strip(): (v.strip() if isinstance(v, str) else v) for k, v in raw.items()}
        row = {k: (v if v != "" else None) for k, v in row.items()}
        if row.get("message") or row.get("source"):
            rows.append(row)
    return rows


def json_response(handler: BaseHTTPRequestHandler, payload: Any, status: int = 200) -> None:
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    handler.send_response(status)
    handler.send_header("Content-Type", "application/json; charset=utf-8")
    handler.send_header("Content-Length", str(len(body)))
    handler.end_headers()
    handler.wfile.write(body)


def read_json_body(handler: BaseHTTPRequestHandler) -> dict[str, Any]:
    length = int(handler.headers.get("Content-Length") or 0)
    if length <= 0:
        return {}
    raw = handler.rfile.read(length)
    return json.loads(raw.decode("utf-8"))


def detect_rule(row: dict[str, Any]) -> dict[str, Any]:
    message = row.get("message", "")
    category = "benign"
    reasons: list[str] = []

    if "Failed password" in message:
        category = "brute_force"
        reasons.append("พบ failed SSH login")
    if "DROP" in message and "FLAGS=SYN" in message:
        category = "port_scan"
        reasons.append("พบ SYN หลายปลายทาง/หลาย port")
    if SQLI_RE.search(message) or XSS_RE.search(message):
        category = "web_attack"
        reasons.append("พบ SQLi/XSS payload ใน request")
    if "sudo:" in message and "USER=root" in message:
        category = "priv_esc"
        reasons.append("พบ sudo ไปยัง root")
    if "OUTBOUND" in message and "bytes=" in message:
        category = "data_exfil"
        reasons.append("พบ outbound transfer ขนาดใหญ่")
    if "interval=" in message and "CONN" in message:
        category = "c2_beacon"
        reasons.append("พบ connection interval คงที่แบบ beacon")

    meta = CATEGORY_META[category]
    expected = row.get("category")
    return {
        "category": category,
        "label": meta["label"],
        "risk_score": meta["risk"],
        "severity": row.get("severity", "info"),
        "mitre": meta["mitre"],
        "reasons": reasons or ["ไม่พบ pattern เสี่ยงตาม rule พื้นฐาน"],
        "matches_label": expected == category,
    }


@dataclass
class SearchHit:
    score: float
    row: dict[str, Any]
    rule: dict[str, Any]
    retrieval: str = "keyword"


def normalize_vector(values: list[float]) -> array:
    norm = math.sqrt(sum(float(v) * float(v) for v in values)) or 1.0
    return array("f", (float(v) / norm for v in values))


def dot_product(left: array, right: array) -> float:
    return sum(a * b for a, b in zip(left, right))


def call_openrouter_embeddings(inputs: list[str], model: str = DEFAULT_EMBED_MODEL) -> list[list[float]]:
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        raise RuntimeError("OPENROUTER_API_KEY is not set")

    payload = {"model": model, "input": inputs}
    req = urllib.request.Request(
        "https://openrouter.ai/api/v1/embeddings",
        data=json.dumps(payload).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "http://localhost:8765",
            "X-Title": "ROGLAG Demo",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            result = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"OpenRouter embedding error {exc.code}: {detail}") from exc

    data = sorted(result["data"], key=lambda item: item["index"])
    return [item["embedding"] for item in data]


class VectorIndex:
    def __init__(self, path: Path, rows: list[dict[str, Any]]):
        self.path = path
        self.rows = rows
        self.ready = path.exists()
        self.model = ""
        self.dimensions = 0
        self.count = 0
        self.error = None
        if self.ready:
            try:
                self._load_meta()
            except Exception as exc:
                self.ready = False
                self.error = str(exc)

    def _load_meta(self) -> None:
        with sqlite3.connect(self.path) as conn:
            meta = dict(conn.execute("select key, value from meta").fetchall())
            self.model = meta.get("model", "")
            self.dimensions = int(meta.get("dimensions", "0"))
            self.count = int(conn.execute("select count(*) from vectors").fetchone()[0])

    def search(self, query: str, limit: int = 8) -> list[SearchHit]:
        if not self.ready:
            return []
        model = self.model or DEFAULT_EMBED_MODEL
        query_vector = normalize_vector(call_openrouter_embeddings([query], model=model)[0])
        scored: list[tuple[float, int]] = []
        with sqlite3.connect(self.path) as conn:
            cursor = conn.execute("select row_id, embedding from vectors")
            for row_id, blob in cursor:
                values = array("f")
                values.frombytes(blob)
                score = dot_product(query_vector, values)
                scored.append((score, int(row_id)))
        hits = heapq.nlargest(limit, scored, key=lambda item: item[0])
        return [
            SearchHit(
                score=round(score, 4),
                row=self.rows[row_id],
                rule=detect_rule(self.rows[row_id]),
                retrieval="vector",
            )
            for score, row_id in hits
            if 0 <= row_id < len(self.rows)
        ]


def qdrant_search(query: str, limit: int = 8) -> list[SearchHit]:
    """Query the real Qdrant vector DB: embed the query then ANN search."""
    if not QDRANT_URL:
        return []
    vector = call_openrouter_embeddings([query])[0]
    headers = {"Content-Type": "application/json"}
    if QDRANT_API_KEY:
        headers["api-key"] = QDRANT_API_KEY
    req = urllib.request.Request(
        f"{QDRANT_URL}/collections/{QDRANT_COLLECTION}/points/search",
        data=json.dumps({"vector": vector, "limit": limit, "with_payload": True}).encode("utf-8"),
        headers=headers,
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=30) as resp:
        result = json.loads(resp.read().decode("utf-8"))
    hits: list[SearchHit] = []
    for item in result.get("result", []):
        payload = item.get("payload") or {}
        row = payload.get("row") if isinstance(payload.get("row"), dict) else payload
        if isinstance(row, dict) and row:
            hits.append(SearchHit(
                score=round(float(item.get("score", 0)), 4),
                row=row,
                rule=detect_rule(row),
                retrieval="qdrant",
            ))
    return hits


def qdrant_status() -> dict[str, Any]:
    if not QDRANT_URL:
        return {"enabled": False}
    try:
        headers = {"api-key": QDRANT_API_KEY} if QDRANT_API_KEY else {}
        req = urllib.request.Request(f"{QDRANT_URL}/collections/{QDRANT_COLLECTION}", headers=headers)
        with urllib.request.urlopen(req, timeout=5) as resp:
            data = json.loads(resp.read().decode("utf-8")).get("result", {})
        return {"enabled": True, "url": QDRANT_URL, "collection": QDRANT_COLLECTION,
                "points": data.get("points_count", 0)}
    except Exception as exc:
        return {"enabled": True, "url": QDRANT_URL, "error": str(exc)}


class LogRagIndex:
    def __init__(self, data_file: Path):
        started = time.time()
        self.data_file = data_file
        self.rows = self._load_rows(data_file)
        self.doc_tokens: list[Counter[str]] = []
        self.doc_freq: Counter[str] = Counter()
        self.postings: defaultdict[str, list[tuple[int, int]]] = defaultdict(list)
        self.category_index: defaultdict[str, list[int]] = defaultdict(list)
        self.idf: dict[str, float] = {}
        self.stats = self._build_stats()
        self._build_index()
        self.vector_index = VectorIndex(VECTOR_FILE, self.rows)
        self.loaded_at = datetime.now().isoformat(timespec="seconds")
        self.load_seconds = round(time.time() - started, 3)

    def _load_rows(self, path: Path) -> list[dict[str, Any]]:
        if not path.exists():
            raise FileNotFoundError(f"Dataset not found: {path}")
        rows = []
        with path.open(encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    rows.append(json.loads(line))
        return rows

    def _build_stats(self) -> dict[str, Any]:
        by_category = Counter(row.get("category", "unknown") for row in self.rows)
        by_severity = Counter(row.get("severity", "unknown") for row in self.rows)
        by_source = Counter(row.get("source", "unknown") for row in self.rows)
        risky = sum(1 for row in self.rows if row.get("category") != "benign")
        top_ips = Counter(row.get("src_ip") for row in self.rows if row.get("src_ip")).most_common(10)
        top_targets = Counter(row.get("dst_ip") for row in self.rows if row.get("dst_ip")).most_common(10)
        times = [row.get("timestamp") for row in self.rows if row.get("timestamp")]
        timeline = Counter((ts or "")[:13] + ":00" for ts in times)
        timeline_by_severity: dict[str, dict[str, int]] = defaultdict(lambda: {"critical": 0, "high": 0, "medium": 0, "info": 0})
        for row in self.rows:
            bucket = (row.get("timestamp") or "")[:13] + ":00"
            if not bucket.strip(":0"):
                continue
            severity = row.get("severity", "info")
            timeline_by_severity[bucket][severity] = timeline_by_severity[bucket].get(severity, 0) + 1

        return {
            "total": len(self.rows),
            "risky": risky,
            "risk_rate": round((risky / max(len(self.rows), 1)) * 100, 2),
            "by_category": dict(by_category),
            "by_severity": dict(by_severity),
            "by_source": dict(by_source),
            "top_ips": top_ips,
            "top_targets": top_targets,
            "timeline": dict(sorted(timeline.items())),
            "timeline_by_severity": dict(sorted(timeline_by_severity.items())),
            "categories": CATEGORY_META,
        }

    def _row_text(self, row: dict[str, Any]) -> str:
        parts = [
            row.get("message", ""),
            row.get("category", ""),
            row.get("severity", ""),
            row.get("source", ""),
            row.get("src_ip", ""),
            row.get("dst_ip", ""),
            row.get("path", ""),
            row.get("user", ""),
        ]
        return " ".join(str(p) for p in parts if p is not None)

    def _build_index(self) -> None:
        for idx, row in enumerate(self.rows):
            counts = Counter(tokenize(self._row_text(row)))
            self.doc_tokens.append(counts)
            self.doc_freq.update(counts.keys())
            self.category_index[row.get("category", "unknown")].append(idx)
            for term, tf in counts.items():
                self.postings[term].append((idx, tf))
        total = max(len(self.rows), 1)
        self.idf = {term: math.log((1 + total) / (1 + df)) + 1 for term, df in self.doc_freq.items()}

    def add_rows(self, raw_rows: list[dict[str, Any]]) -> dict[str, Any]:
        """Classify + append rows into the live in-memory index, refresh stats.

        Returns an ingest summary (counts + detection accuracy vs any label)."""
        by_cat: Counter[str] = Counter()
        by_sev: Counter[str] = Counter()
        correct = labelled = 0
        start = len(self.rows)
        events: list[dict[str, Any]] = []

        for raw in raw_rows:
            rule = detect_rule(raw)
            category = rule["category"]
            severity = raw.get("severity") or CATEGORY_DEFAULT_SEVERITY.get(category, "info")
            expected = raw.get("category")
            row = dict(raw)
            row["category"] = category
            row["severity"] = severity
            row.setdefault("timestamp", datetime.now().isoformat(timespec="seconds"))

            idx = len(self.rows)
            self.rows.append(row)
            counts = Counter(tokenize(self._row_text(row)))
            self.doc_tokens.append(counts)
            self.doc_freq.update(counts.keys())
            self.category_index[category].append(idx)
            for term, tf in counts.items():
                self.postings[term].append((idx, tf))

            by_cat[category] += 1
            by_sev[severity] += 1
            if expected:
                labelled += 1
                if expected == category:
                    correct += 1
            if len(events) < 400:
                events.append({
                    "timestamp": row.get("timestamp"),
                    "source": row.get("source"),
                    "category": category,
                    "label": rule["label"],
                    "severity": severity,
                    "mitre": rule["mitre"],
                    "risk_score": rule["risk_score"],
                    "src_ip": row.get("src_ip"),
                    "dst_ip": row.get("dst_ip"),
                    "message": row.get("message"),
                })

        total = max(len(self.rows), 1)
        self.idf = {term: math.log((1 + total) / (1 + df)) + 1 for term, df in self.doc_freq.items()}
        self.stats = self._build_stats()

        ingested = len(self.rows) - start
        return {
            "ingested": ingested,
            "total": len(self.rows),
            "by_category": dict(by_cat),
            "by_severity": dict(by_sev),
            "risky": sum(v for k, v in by_cat.items() if k != "benign"),
            "critical": by_sev.get("critical", 0),
            "detection_accuracy": round(correct / labelled * 100, 1) if labelled else None,
            "events": events,
        }

    def filter_rows(self, params: dict[str, list[str]]) -> dict[str, Any]:
        limit = min(int((params.get("limit") or ["100"])[0]), 500)
        offset = max(int((params.get("offset") or ["0"])[0]), 0)
        category = (params.get("category") or [""])[0]
        severity = (params.get("severity") or [""])[0]
        source = (params.get("source") or [""])[0]
        query = (params.get("q") or [""])[0].strip().lower()

        filtered = []
        for row in self.rows:
            if category and row.get("category") != category:
                continue
            if severity and row.get("severity") != severity:
                continue
            if source and row.get("source") != source:
                continue
            if query and query not in self._row_text(row).lower():
                continue
            filtered.append(row)

        page = filtered[offset : offset + limit]
        return {
            "total": len(filtered),
            "offset": offset,
            "limit": limit,
            "rows": [{**row, "rule": detect_rule(row)} for row in page],
        }

    def _alert_key(self, row: dict[str, Any]) -> str:
        if row.get("category") == "priv_esc":
            return "|".join([row.get("category", ""), row.get("user", ""), row.get("host", "")])
        return "|".join(
            [
                row.get("category", ""),
                row.get("src_ip", ""),
                row.get("dst_ip", ""),
                row.get("user", ""),
                str(row.get("beacon_interval", "")),
            ]
        )

    def critical_alerts(self, params: dict[str, list[str]]) -> dict[str, Any]:
        limit = min(int((params.get("limit") or ["20"])[0]), 100)
        since = (params.get("since") or [""])[0]
        groups: dict[str, list[dict[str, Any]]] = defaultdict(list)

        for row in self.rows:
            if row.get("severity") != "critical":
                continue
            if since and str(row.get("timestamp", "")) <= since:
                continue
            groups[self._alert_key(row)].append(row)

        alerts = []
        for key, items in groups.items():
            items.sort(key=lambda item: str(item.get("timestamp", "")))
            first = items[0]
            last = items[-1]
            meta = CATEGORY_META.get(first.get("category"), CATEGORY_META["benign"])
            alert_id = hashlib.sha1(f"{key}|{first.get('timestamp')}|{last.get('timestamp')}".encode("utf-8")).hexdigest()[:12]
            rule = detect_rule(last)
            alerts.append(
                {
                    "id": f"ALT-{alert_id}",
                    "title": meta["label"],
                    "category": first.get("category"),
                    "severity": "critical",
                    "mitre": meta["mitre"],
                    "status": "new",
                    "count": len(items),
                    "first_seen": first.get("timestamp"),
                    "last_seen": last.get("timestamp"),
                    "src_ip": first.get("src_ip", "-"),
                    "dst_ip": first.get("dst_ip", "-"),
                    "user": first.get("user", "-"),
                    "message": last.get("message"),
                    "rule": rule["reasons"],
                    "recommended_action": meta["runbook"],
                }
            )

        alerts.sort(key=lambda item: str(item["last_seen"]), reverse=True)
        return {"total": len(alerts), "limit": limit, "alerts": alerts[:limit]}

    def _search_keyword(self, query: str, limit: int = 8) -> list[SearchHit]:
        terms, category_terms = expand_query(query)
        if not terms:
            return []

        query_counts = Counter(terms)
        scores: defaultdict[int, float] = defaultdict(float)

        for term, qtf in query_counts.items():
            for idx, tf in self.postings.get(term, []):
                scores[idx] += (1 + math.log(tf)) * self.idf.get(term, 1.0) * qtf

        for category in category_terms:
            for idx in self.category_index.get(category, []):
                scores[idx] += 12

        lowered = query.lower()
        for idx in list(scores.keys()):
            if lowered and lowered in self._row_text(self.rows[idx]).lower():
                scores[idx] += 8

        ranked = sorted(scores.items(), key=lambda item: item[1], reverse=True)[:limit]
        return [SearchHit(score=round(score, 3), row=self.rows[idx], rule=detect_rule(self.rows[idx])) for idx, score in ranked]

    def hits_from_rows(self, rows: list[dict[str, Any]], limit: int = 8) -> list[SearchHit]:
        return [
            SearchHit(score=1.0, row=row, rule=detect_rule(row), retrieval="selected_incident")
            for row in rows[:limit]
        ]

    def search(self, query: str, limit: int = 8, retrieval: str = "auto") -> list[SearchHit]:
        # Production vector RAG: real Qdrant DB first.
        if retrieval in {"auto", "vector", "qdrant"} and QDRANT_URL and os.getenv("OPENROUTER_API_KEY"):
            try:
                hits = qdrant_search(query, limit)
                if hits or retrieval in {"vector", "qdrant"}:
                    return hits
            except Exception:
                if retrieval == "qdrant":
                    raise
        # Legacy local SQLite vector index (kept as fallback).
        if retrieval in {"auto", "vector"} and self.vector_index.ready and os.getenv("OPENROUTER_API_KEY"):
            try:
                hits = self.vector_index.search(query, limit)
                if hits or retrieval == "vector":
                    return hits
            except Exception:
                if retrieval == "vector":
                    raise
        return self._search_keyword(query, limit)

    def answer_local(self, question: str, hits: list[SearchHit]) -> dict[str, Any]:
        rows = [hit.row for hit in hits]
        by_cat = Counter(row.get("category") for row in rows)
        by_sev = Counter(row.get("severity") for row in rows)
        categories = [cat for cat, _ in by_cat.most_common()]

        if not rows:
            answer = (
                "## ยังไม่พบหลักฐาน\n\n"
                "- ลองถามด้วย IP, category เช่น `brute_force`, `c2_beacon`\n"
                "- หรือเลือก incident ก่อนแล้วกดถามจากหลักฐาน"
            )
        else:
            top = rows[0]
            cat = top.get("category", "benign")
            meta = CATEGORY_META.get(cat, CATEGORY_META["benign"])
            answer = (
                f"## สรุปความเสี่ยง\n\n"
                f"- พบหลักฐานที่เกี่ยวข้อง **{len(rows)} รายการ**\n"
                f"- กลุ่มเด่น: **{', '.join(categories[:3])}**\n"
                f"- รายการสำคัญสุด: **{meta['label']}** severity=`{top.get('severity')}` source=`{top.get('source')}`\n"
                f"- เวลา: `{top.get('timestamp')}`\n\n"
                f"## เหตุผล\n\n"
                f"- {', '.join(detect_rule(top)['reasons'])}\n\n"
                f"## Next action\n\n"
                f"- {meta['runbook']}"
            )

        return {
            "answer": answer,
            "mode": "local",
            "summary": {
                "categories": dict(by_cat),
                "severities": dict(by_sev),
            },
        }

    def build_context(self, hits: list[SearchHit]) -> str:
        lines = []
        for i, hit in enumerate(hits, 1):
            row = hit.row
            meta = CATEGORY_META.get(row.get("category"), CATEGORY_META["benign"])
            lines.append(
                "\n".join(
                    [
                        f"[{i}] score={hit.score}",
                        f"time={row.get('timestamp')} source={row.get('source')} severity={row.get('severity')}",
                        f"category={row.get('category')} label={meta['label']} mitre={meta['mitre']}",
                        f"src_ip={row.get('src_ip', '-')} dst_ip={row.get('dst_ip', '-')}",
                        f"message={row.get('message')}",
                        f"rule_reason={'; '.join(hit.rule['reasons'])}",
                        f"runbook={meta['runbook']}",
                    ]
                )
            )
        return "\n\n".join(lines)


def call_openrouter(question: str, context: str, model: str) -> str:
    api_key = os.getenv("OPENROUTER_API_KEY")
    if not api_key:
        raise RuntimeError("OPENROUTER_API_KEY is not set")

    payload = {
        "model": model or DEFAULT_MODEL,
        "messages": [
            {
                "role": "system",
                "content": (
                    "You are ROGLAG, a defensive security log analyst. "
                    "Answer in Thai. Use only the provided retrieved logs and runbook context. "
                    "ALWAYS output ALL four Markdown sections in this order and never stop early: "
                    "## สรุป, ## หลักฐาน, ## ความเสี่ยง, ## Next action. "
                    "Keep each section short; finish every section. "
                    "Use short bullet points and inline code for IPs, users, MITRE IDs, and log fields. "
                    "Do not include safety labels such as 'User Safety'. "
                    "Do not provide offensive instructions."
                ),
            },
            {
                "role": "user",
                "content": f"คำถาม: {question}\n\nRetrieved context:\n{context}",
            },
        ],
        "temperature": 0.2,
        "max_tokens": 900,
    }
    data = json.dumps(payload).encode("utf-8")
    req = urllib.request.Request(
        "https://openrouter.ai/api/v1/chat/completions",
        data=data,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "http://localhost:8765",
            "X-Title": "ROGLAG Demo",
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=45) as resp:
            result = json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"OpenRouter error {exc.code}: {detail}") from exc

    return result["choices"][0]["message"]["content"]


def ensure_answer_sections(answer: str, hits: list[SearchHit]) -> str:
    """Guarantee the answer always ends with a concrete Next action + risk line,
    even when a flaky free LLM returns a truncated/incomplete response."""
    text = (answer or "").strip()
    if not hits:
        return text
    meta = CATEGORY_META.get(hits[0].row.get("category"), CATEGORY_META["benign"])
    lowered = text.lower()
    if "ความเสี่ยง" not in text:
        text += f"\n\n## ความเสี่ยง\n\n- ระดับเด่น **{meta['label']}** · MITRE `{meta['mitre']}` · risk `{meta['risk']}`"
    if "next action" not in lowered:
        text += f"\n\n## Next action\n\n- {meta['runbook']}"
    return text


INDEX = LogRagIndex(DATA_FILE)


class Handler(BaseHTTPRequestHandler):
    server_version = "ROGLAG/0.1"

    def do_HEAD(self) -> None:
        parsed = urlparse(self.path)
        if parsed.path in {"/api/health", "/api/stats", "/api/logs", "/api/alerts"}:
            self.send_response(200)
            self.send_header("Content-Type", "application/json; charset=utf-8")
            self.end_headers()
            return
        if parsed.path in {"", "/"}:
            file_path = WEB_DIR / "index.html"
        else:
            file_path = (WEB_DIR / parsed.path.lstrip("/")).resolve()
            if not str(file_path).startswith(str(WEB_DIR.resolve())):
                self.send_error(403)
                return
        if not file_path.exists() or file_path.is_dir():
            self.send_error(404)
            return
        self.send_response(200)
        self.send_header("Content-Length", str(file_path.stat().st_size))
        self.end_headers()

    def do_GET(self) -> None:
        parsed = urlparse(self.path)
        if parsed.path == "/api/health":
            json_response(
                self,
                {
                    "ok": True,
                    "dataset": str(DATA_FILE.relative_to(ROOT)),
                    "records": len(INDEX.rows),
                    "loaded_at": INDEX.loaded_at,
                    "load_seconds": INDEX.load_seconds,
                    "openrouter_ready": bool(os.getenv("OPENROUTER_API_KEY")),
                    "default_model": DEFAULT_MODEL,
                    "embed_model": DEFAULT_EMBED_MODEL,
                    "vector_ready": INDEX.vector_index.ready,
                    "vector_records": INDEX.vector_index.count,
                    "vector_dimensions": INDEX.vector_index.dimensions,
                    "vector_error": INDEX.vector_index.error,
                    "qdrant": qdrant_status(),
                },
            )
            return
        if parsed.path == "/api/stats":
            json_response(self, INDEX.stats)
            return
        if parsed.path == "/api/logs":
            json_response(self, INDEX.filter_rows(parse_qs(parsed.query)))
            return
        if parsed.path == "/api/alerts":
            json_response(self, INDEX.critical_alerts(parse_qs(parsed.query)))
            return
        if parsed.path == "/api/generate/csv":
            self.serve_generated_csv(parse_qs(parsed.query))
            return
        self.serve_static(parsed.path)

    def do_POST(self) -> None:
        parsed = urlparse(self.path)
        try:
            body = read_json_body(self)
            if parsed.path == "/api/mock/emit":
                if not GENERATOR_OK:
                    json_response(self, {"error": "generator not available"}, 503)
                    return
                params = parse_qs(parsed.query)
                count = min(int((params.get("count") or ["20"])[0]), 500)
                profile = (params.get("profile") or ["mixed"])[0]
                summary = INDEX.add_rows(generate_live(count, profile=profile))
                json_response(self, summary)
                return
            if parsed.path == "/api/ingest/csv":
                rows = parse_csv_text(str(body.get("content") or ""))
                if not rows:
                    json_response(self, {"error": "no parseable rows (need a message or source column)"}, 400)
                    return
                if len(rows) > 50000:
                    rows = rows[:50000]
                summary = INDEX.add_rows(rows)
                summary["filename"] = body.get("filename")
                json_response(self, summary)
                return
            if parsed.path == "/api/search":
                query = str(body.get("query", ""))
                limit = min(int(body.get("limit", 8)), 30)
                retrieval = str(body.get("retrieval") or "auto")
                hits = INDEX.search(query, limit, retrieval=retrieval)
                json_response(self, {"hits": [hit.__dict__ for hit in hits]})
                return
            if parsed.path == "/api/ask":
                question = str(body.get("question", "")).strip()
                limit = min(int(body.get("limit", 8)), 20)
                model = str(body.get("model") or DEFAULT_MODEL)
                retrieval = str(body.get("retrieval") or "auto")
                if not question:
                    json_response(self, {"error": "question is required"}, 400)
                    return
                context_rows = body.get("context_rows")
                if isinstance(context_rows, list) and context_rows:
                    hits = INDEX.hits_from_rows([row for row in context_rows if isinstance(row, dict)], limit)
                else:
                    hits = INDEX.search(question, limit, retrieval=retrieval)
                context = INDEX.build_context(hits)
                local = INDEX.answer_local(question, hits)
                mode = "local"
                answer = local["answer"]
                llm_error = None
                if body.get("use_llm") and os.getenv("OPENROUTER_API_KEY"):
                    try:
                        answer = call_openrouter(question, context, model)
                        mode = "openrouter"
                    except Exception as exc:  # keep demo alive when free providers rate-limit.
                        llm_error = str(exc)
                answer = ensure_answer_sections(answer, hits)
                json_response(
                    self,
                    {
                        "answer": answer,
                        "mode": mode,
                        "model": model,
                        "retrieval": hits[0].retrieval if hits else "none",
                        "llm_error": llm_error,
                        "local_summary": local["summary"],
                        "hits": [hit.__dict__ for hit in hits],
                    },
                )
                return
            json_response(self, {"error": "not found"}, 404)
        except Exception as exc:
            json_response(self, {"error": str(exc)}, 500)

    def serve_generated_csv(self, params: dict[str, list[str]]) -> None:
        if not GENERATOR_OK:
            json_response(self, {"error": "generator not available"}, 503)
            return
        count = min(int((params.get("count") or ["200"])[0]), 20000)
        profile = (params.get("profile") or ["mixed"])[0]
        records = gen_records(count, profile=profile, seed=None)  # unseeded -> unique each time
        body = to_csv(records).encode("utf-8")
        fname = f"roglag_logs_{profile}_{len(records)}.csv"
        self.send_response(200)
        self.send_header("Content-Type", "text/csv; charset=utf-8")
        self.send_header("Content-Disposition", f'attachment; filename="{fname}"')
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def serve_static(self, path: str) -> None:
        if path in {"", "/"}:
            file_path = WEB_DIR / "index.html"
        else:
            file_path = (WEB_DIR / path.lstrip("/")).resolve()
            if not str(file_path).startswith(str(WEB_DIR.resolve())):
                self.send_error(403)
                return
        if not file_path.exists() or file_path.is_dir():
            self.send_error(404)
            return
        content_type = "text/plain; charset=utf-8"
        if file_path.suffix == ".html":
            content_type = "text/html; charset=utf-8"
        elif file_path.suffix == ".css":
            content_type = "text/css; charset=utf-8"
        elif file_path.suffix == ".js":
            content_type = "application/javascript; charset=utf-8"
        elif file_path.suffix == ".svg":
            content_type = "image/svg+xml; charset=utf-8"
        elif file_path.suffix == ".png":
            content_type = "image/png"
        data = file_path.read_bytes()
        self.send_response(200)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def log_message(self, fmt: str, *args: Any) -> None:
        sys.stderr.write("[%s] %s\n" % (self.log_date_time_string(), fmt % args))


def main() -> None:
    port = int(os.getenv("PORT", "8765"))
    host = os.getenv("HOST", "127.0.0.1")
    httpd = ThreadingHTTPServer((host, port), Handler)
    print(f"ROGLAG demo running at http://{host}:{port}")
    print(f"Loaded {len(INDEX.rows)} records from {DATA_FILE}")
    print(f"OpenRouter: {'ready' if os.getenv('OPENROUTER_API_KEY') else 'not configured'}")
    print(
        "Vector index: "
        f"{'ready' if INDEX.vector_index.ready else 'not built'} "
        f"({INDEX.vector_index.count} records, {INDEX.vector_index.dimensions} dims)"
    )
    httpd.serve_forever()


if __name__ == "__main__":
    main()