from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Header, Depends, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from datetime import datetime, timedelta
from typing import Optional, List
import psycopg2
import psycopg2.extras
import os
import secrets
import ipaddress
import httpx
import subprocess
import asyncio
import re
import time
from datetime import datetime as dt
from psycopg2 import pool

app = FastAPI(title="Smart IDS API", version="1.0")

# ── CORS Middleware Configuration ──
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"]
)

# ── Retraining Job State ──
retrain_status = {
    "state": "idle",          # idle | running | success | rejected | error
    "message": "",
    "started_at": None,
    "finished_at": None,
    "new_accuracy": None,
    "old_accuracy": None,
}

SPARK_HOST = os.getenv("SPARK_HOST", "10.0.0.15")
SPARK_SSH_USER = os.getenv("SPARK_SSH_USER", "ubuntu")
DB_PASSWORD = os.getenv("DB_PASSWORD", "")


def run_retrain_job():
    global retrain_status
    retrain_status["state"] = "running"
    retrain_status["message"] = "Retraining job started..."
    retrain_status["started_at"] = dt.now().isoformat()
    retrain_status["finished_at"] = None

    remote_cmd = (
        f"sudo docker exec -u root -e DB_PASSWORD={DB_PASSWORD} spark-master "
        f"/opt/spark/bin/spark-submit "
        f"--master spark://{SPARK_HOST}:7077 "
        f"--total-executor-cores 2 --executor-memory 1G "
        f"--conf spark.driver.host={SPARK_HOST} "
        f"--conf spark.driver.bindAddress=0.0.0.0 "
        f"--conf spark.driver.port=37100 "
        f"--conf spark.blockManager.port=37101 "
        f"/opt/spark-apps/retrain_model.py"
    )

    ssh_cmd = [
        "ssh", "-o", "StrictHostKeyChecking=no",
        "-i", "/root/.ssh/id_rsa",
        f"{SPARK_SSH_USER}@{SPARK_HOST}",
        remote_cmd
    ]

    try:
        result = subprocess.run(
            ssh_cmd, capture_output=True, text=True, timeout=900
        )
        output = result.stdout + result.stderr

        match = re.search(r"RETRAIN_RESULT:\s*(\w+)(.*)", output)

        if not match:
            retrain_status["state"] = "error"
            retrain_status["message"] = "No result line found — check job logs."
        else:
            outcome = match.group(1)
            details = match.group(2)

            acc_matches = re.findall(r"(\w+)_?accuracy=([\d.]+)", details)
            for name, val in acc_matches:
                if "new" in name:
                    retrain_status["new_accuracy"] = float(val)
                elif "old" in name:
                    retrain_status["old_accuracy"] = float(val)
                elif name == "accuracy":
                    retrain_status["new_accuracy"] = float(val)

            if outcome in ("promoted", "would_promote"):
                retrain_status["state"] = "success"
                retrain_status["message"] = (
                    "New model trained and promoted successfully."
                    if outcome == "promoted"
                    else "New model would be promoted (dry run — no files changed)."
                )
            elif outcome == "rejected":
                retrain_status["state"] = "rejected"
                retrain_status["message"] = (
                    "Retraining finished, but the new model did not outperform "
                    "the current production model. The existing model was kept."
                )
            elif outcome == "insufficient_data":
                retrain_status["state"] = "rejected"
                retrain_status["message"] = (
                    "Not enough analyst-reviewed alerts yet to retrain safely."
                )
            else:
                retrain_status["state"] = "error"
                retrain_status["message"] = f"Unexpected result: {outcome}"

    except subprocess.TimeoutExpired:
        retrain_status["state"] = "error"
        retrain_status["message"] = "Retraining job timed out after 15 minutes."
    except Exception as e:
        retrain_status["state"] = "error"
        retrain_status["message"] = f"Failed to run retraining job: {e}"

    retrain_status["finished_at"] = dt.now().isoformat()


@app.post("/api/retrain")
async def trigger_retrain(background_tasks: BackgroundTasks):
    if retrain_status["state"] == "running":
        raise HTTPException(status_code=409, detail="A retraining job is already running.")

    background_tasks.add_task(run_retrain_job)
    return {"status": "started"}


@app.get("/api/retrain/status")
async def get_retrain_status():
    return retrain_status

ABUSEIPDB_API_KEY = os.getenv("ABUSEIPDB_API_KEY", "")
CACHE_TTL_HOURS = 24

def is_private_ip(ip: str) -> bool:
    try:
        return ipaddress.ip_address(ip).is_private
    except ValueError:
        return True

async def get_ip_reputation(ip: str) -> dict:
    if is_private_ip(ip):
        return {
            "status": "private_ip",
            "note": "Internal/simulated address — no public reputation data available",
            "country": None,
            "abuse_score": None,
            "is_tor": None,
            "total_reports": None
        }

    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute("SELECT data, checked_at FROM ip_reputation_cache WHERE ip = %s", (ip,))
        cached = cur.fetchone()
        cur.close()
        conn.close()

        if cached and (datetime.now() - cached[1]) < timedelta(hours=CACHE_TTL_HOURS):
            return cached[0]
    except Exception:
        pass

    try:
        async with httpx.AsyncClient() as client:
            resp = await client.get(
                "https://api.abuseipdb.com/api/v2/check",
                params={"ipAddress": ip, "maxAgeInDays": 90},
                headers={"Key": ABUSEIPDB_API_KEY, "Accept": "application/json"},
                timeout=5
            )
        data = resp.json().get("data", {})
        result = {
            "status": "ok",
            "abuse_score": data.get("abuseConfidenceScore"),
            "total_reports": data.get("totalReports"),
            "country": data.get("countryCode"),
            "isp": data.get("isp"),
            "is_tor": data.get("isTor", False),
            "usage_type": data.get("usageType")
        }
    except Exception as e:
        result = {"status": "error", "note": str(e)}

    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute("""
            INSERT INTO ip_reputation_cache (ip, data, checked_at)
            VALUES (%s, %s, NOW())
            ON CONFLICT (ip) DO UPDATE SET data = EXCLUDED.data, checked_at = NOW()
        """, (ip, psycopg2.extras.Json(result) if hasattr(psycopg2, 'extras') else str(result)))
        conn.commit()
        cur.close()
        conn.close()
    except Exception:
        pass

    return result


# ── Database Connection (original — used by all other endpoints) ──
def get_db():
    return psycopg2.connect(
        host=os.getenv("DB_HOST", "10.0.0.18"),
        port=5432,
        database="ids_db",
        user="ids_user",
        password="ids_pass123"
    )

# ── Database Connection Pool (used only by /api/alerts) ──
db_pool = pool.ThreadedConnectionPool(
    minconn=5,
    maxconn=50,
    host=os.getenv("DB_HOST", "10.0.0.18"),
    port=5432,
    database="ids_db",
    user="ids_user",
    password="ids_pass123"
)

def get_db_pooled():
    while True:
        try:
            return db_pool.getconn()
        except pool.PoolError:
            time.sleep(0.05)

def release_db_pooled(conn):
    db_pool.putconn(conn)

# ── Simple session-based auth (single admin account) ──
ADMIN_USERNAME = os.getenv("ADMIN_USERNAME", "admin")
ADMIN_PASSWORD = os.getenv("ADMIN_PASSWORD", "IDPS@2026")
active_tokens = set()

class LoginRequest(BaseModel):
    username: str
    password: str

@app.post("/api/login")
def login(data: LoginRequest):
    if data.username == ADMIN_USERNAME and data.password == ADMIN_PASSWORD:
        token = secrets.token_hex(24)
        active_tokens.add(token)
        return {"status": "ok", "token": token, "username": data.username}
    raise HTTPException(status_code=401, detail="Invalid username or password")

@app.post("/api/logout")
def logout(authorization: Optional[str] = Header(None)):
    if authorization and authorization.startswith("Bearer "):
        active_tokens.discard(authorization.replace("Bearer ", ""))
    return {"status": "logged out"}

def verify_token(authorization: Optional[str] = Header(None)):
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing authentication token")
    token = authorization.replace("Bearer ", "")
    if token not in active_tokens:
        raise HTTPException(status_code=401, detail="Invalid or expired session")
    return token

# ── WebSocket Connection Manager ──
class ConnectionManager:
    def __init__(self):
        self.active: list = []

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self.active.append(ws)

    def disconnect(self, ws: WebSocket):
        if ws in self.active:
            self.active.remove(ws)

    async def broadcast(self, message: dict):
        dead = []
        for ws in self.active:
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(ws)

manager = ConnectionManager()

@app.websocket("/ws/alerts")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        manager.disconnect(websocket)

# ── Create Tables on Startup ──
@app.on_event("startup")
def startup():
    try:
        conn = get_db()
        cur = conn.cursor()

        cur.execute("""
            CREATE TABLE IF NOT EXISTS alerts (
                id SERIAL PRIMARY KEY,
                attack_type VARCHAR(50),
                src_ip VARCHAR(50),
                severity VARCHAR(20),
                timestamp TIMESTAMP,
                created_at TIMESTAMP DEFAULT NOW(),
                verdict VARCHAR(20) DEFAULT 'pending',
                reviewed_by VARCHAR(50),
                reviewed_at TIMESTAMP,
                occurrence_count INTEGER DEFAULT 1,
                last_seen TIMESTAMP DEFAULT NOW()
            )
        """)

        cur.execute("""
            CREATE TABLE IF NOT EXISTS historical_analysis (
                id SERIAL PRIMARY KEY,
                src_ip VARCHAR(50),
                total_connections INTEGER,
                unique_ports INTEGER,
                avg_bytes FLOAT,
                anomaly_score FLOAT,
                is_suspicious BOOLEAN,
                analysis_date DATE DEFAULT CURRENT_DATE,
                created_at TIMESTAMP DEFAULT NOW()
            )
        """)

        cur.execute("""
            CREATE TABLE IF NOT EXISTS blocked_ips (
                id SERIAL PRIMARY KEY,
                ip_address VARCHAR(50) UNIQUE,
                reason VARCHAR(100),
                blocked_at TIMESTAMP DEFAULT NOW()
            )
        """)

        cur.execute("""
            CREATE TABLE IF NOT EXISTS ip_reputation_cache (
                ip VARCHAR(50) PRIMARY KEY,
                data JSONB,
                checked_at TIMESTAMP DEFAULT NOW()
            )
        """)

        conn.commit()
        cur.close()
        conn.close()
        print("Database tables ready")
    except Exception as e:
        print(f"Database error: {e}")

# ── Pydantic Models ──
class Alert(BaseModel):
    attack_type: str
    src_ip: str
    severity: str
    timestamp: Optional[str] = None
    occurrence_count: Optional[int] = 1
    sample_features: Optional[dict] = None

class VerdictUpdate(BaseModel):
    verdict: str
    reviewed_by: Optional[str] = "admin"

class HistoricalResult(BaseModel):
    src_ip: str
    total_connections: int
    unique_ports: int
    avg_bytes: float
    anomaly_score: float
    is_suspicious: bool

class BlockIP(BaseModel):
    ip_address: str
    reason: Optional[str] = "Manual block"

# ── Health Check ──
@app.get("/")
def root():
    return {"status": "Smart IDS API is running", "version": "1.0"}

@app.get("/api/health")
def health():
    try:
        conn = get_db()
        conn.close()
        db_status = "connected"
    except Exception:
        db_status = "disconnected"
    return {
        "status": "healthy",
        "database": db_status,
        "timestamp": datetime.now().isoformat()
    }

# ── Real-time Alerts ──
@app.post("/api/alerts")
async def create_alert(alert: Alert):
    try:
        conn = get_db_pooled()
        cur = conn.cursor()

        cur.execute("""
            SELECT id, occurrence_count FROM alerts
            WHERE src_ip = %s
              AND attack_type = %s
              AND (verdict = 'pending' OR verdict IS NULL)
              AND last_seen >= NOW() - INTERVAL '60 seconds'
            ORDER BY created_at DESC
            LIMIT 1
        """, (alert.src_ip, alert.attack_type))
        existing = cur.fetchone()

        if existing:
            existing_id, existing_count = existing
            new_count = existing_count + (alert.occurrence_count or 1)
            cur.execute("""
                UPDATE alerts
                SET occurrence_count = %s, last_seen = NOW(), sample_features = %s
                WHERE id = %s
                RETURNING id, created_at
            """, (
                new_count,
                psycopg2.extras.Json(alert.sample_features) if alert.sample_features else None,
                existing_id
            ))
            row = cur.fetchone()
            conn.commit()
            cur.close()
            release_db_pooled(conn)

            payload = {
                "id": row[0],
                "attack_type": alert.attack_type,
                "src_ip": alert.src_ip,
                "severity": alert.severity,
                "timestamp": alert.timestamp or datetime.now().isoformat(),
                "occurrence_count": new_count,
                "verdict": "pending"
            }
            await manager.broadcast({"type": "alert_updated", "data": payload})
            return {"status": "updated", "id": row[0], "occurrence_count": new_count}

        cur.execute("""
            INSERT INTO alerts (attack_type, src_ip, severity, timestamp, occurrence_count, last_seen, sample_features)
            VALUES (%s, %s, %s, %s, %s, NOW(), %s)
            RETURNING id, created_at
        """, (
            alert.attack_type,
            alert.src_ip,
            alert.severity,
            alert.timestamp or datetime.now().isoformat(),
            alert.occurrence_count or 1,
            psycopg2.extras.Json(alert.sample_features) if alert.sample_features else None
        ))
        row = cur.fetchone()
        conn.commit()
        cur.close()
        release_db_pooled(conn)

        payload = {
            "id": row[0],
            "attack_type": alert.attack_type,
            "src_ip": alert.src_ip,
            "severity": alert.severity,
            "timestamp": alert.timestamp or datetime.now().isoformat(),
            "occurrence_count": alert.occurrence_count or 1,
            "verdict": "pending"
        }
        await manager.broadcast({"type": "new_alert", "data": payload})
        return {"status": "saved", "data": alert.dict()}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/alerts")
def get_alerts(
    limit: int = 100,
    severity: Optional[str] = None,
    attack_type: Optional[str] = None,
    verdict: Optional[str] = None,
    src_ip: Optional[str] = None
):
    try:
        conn = get_db_pooled()
        cur = conn.cursor()

        conditions = []
        params = []

        if severity:
            conditions.append("severity = %s")
            params.append(severity)
        if attack_type:
            conditions.append("attack_type = %s")
            params.append(attack_type)
        if verdict:
            conditions.append("verdict = %s")
            params.append(verdict)
        if src_ip:
            conditions.append("src_ip ILIKE %s")
            params.append(f"%{src_ip}%")

        where_clause = f"WHERE {' AND '.join(conditions)}" if conditions else ""
        params.append(limit)

        cur.execute(f"""
            SELECT id, attack_type, src_ip, severity, timestamp, created_at, verdict, occurrence_count
            FROM alerts
            {where_clause}
            ORDER BY created_at DESC
            LIMIT %s
        """, params)
        rows = cur.fetchall()
        cur.close()
        release_db_pooled(conn)
        return {"count": len(rows), "alerts": [
            {
                "id": r[0], "attack_type": r[1], "src_ip": r[2],
                "severity": r[3], "timestamp": str(r[4]),
                "created_at": str(r[5]), "verdict": r[6] or "pending",
                "occurrence_count": r[7] or 1
            } for r in rows
        ]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.patch("/api/alerts/{alert_id}/verdict")
def update_verdict(alert_id: int, data: VerdictUpdate, _: str = Depends(verify_token)):
    if data.verdict not in ("true_positive", "false_positive", "pending"):
        raise HTTPException(status_code=400, detail="Invalid verdict value")
    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute("""
            UPDATE alerts
            SET verdict = %s, reviewed_by = %s, reviewed_at = NOW()
            WHERE id = %s
        """, (data.verdict, data.reviewed_by, alert_id))
        conn.commit()
        cur.close()
        conn.close()
        return {"status": "updated", "alert_id": alert_id, "verdict": data.verdict}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/alerts/{alert_id}/details")
async def get_alert_details(alert_id: int):
    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute("SELECT id, attack_type, src_ip, severity, timestamp, sample_features FROM alerts WHERE id = %s", (alert_id,))
        alert = cur.fetchone()
        cur.close()
        conn.close()

        if not alert:
            raise HTTPException(status_code=404, detail="Alert not found")

        reputation = await get_ip_reputation(alert[2])

        return {
            "id": alert[0],
            "attack_type": alert[1],
            "src_ip": alert[2],
            "severity": alert[3],
            "timestamp": str(alert[4]),
            "sample_features": alert[5],
            "ip_reputation": reputation
        }
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/alerts/reviewed")
def get_reviewed_alerts():
    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute("""
            SELECT id, attack_type, src_ip, severity, timestamp, verdict, reviewed_at
            FROM alerts
            WHERE verdict IN ('true_positive', 'false_positive')
            ORDER BY reviewed_at DESC
        """)
        rows = cur.fetchall()
        cur.close()
        conn.close()
        return {"count": len(rows), "alerts": [
            {"id": r[0], "attack_type": r[1], "src_ip": r[2],
             "severity": r[3], "timestamp": str(r[4]),
             "verdict": r[5], "reviewed_at": str(r[6])}
            for r in rows
        ]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/alerts/reviewed/csv")
def export_reviewed_csv():
    import csv
    import io
    from fastapi.responses import StreamingResponse
    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute("""
            SELECT id, attack_type, src_ip, severity, timestamp, verdict, reviewed_by, reviewed_at
            FROM alerts
            WHERE verdict IN ('true_positive', 'false_positive')
            ORDER BY reviewed_at DESC
        """)
        rows = cur.fetchall()
        cur.close()
        conn.close()

        buffer = io.StringIO()
        writer = csv.writer(buffer)
        writer.writerow(["id", "attack_type", "src_ip", "severity", "timestamp", "verdict", "reviewed_by", "reviewed_at"])
        writer.writerows(rows)
        buffer.seek(0)

        return StreamingResponse(
            iter([buffer.getvalue()]),
            media_type="text/csv",
            headers={"Content-Disposition": "attachment; filename=reviewed_alerts_for_retraining.csv"}
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ── Statistics ──
@app.get("/api/stats")
def get_stats():
    try:
        conn = get_db()
        cur = conn.cursor()

        cur.execute("SELECT COUNT(*) FROM alerts")
        total = cur.fetchone()[0]

        cur.execute("""
            SELECT attack_type, COUNT(*) as count
            FROM alerts
            GROUP BY attack_type
            ORDER BY count DESC
        """)
        by_type = cur.fetchall()

        cur.execute("""
            SELECT severity, COUNT(*) as count
            FROM alerts
            GROUP BY severity
        """)
        by_severity = cur.fetchall()

        cur.execute("""
            SELECT src_ip, COUNT(*) as count
            FROM alerts
            GROUP BY src_ip
            ORDER BY count DESC
            LIMIT 10
        """)
        top_attackers = cur.fetchall()

        cur.execute("SELECT COUNT(*) FROM blocked_ips")
        blocked = cur.fetchone()[0]

        cur.execute("SELECT COUNT(*) FROM historical_analysis WHERE is_suspicious = true")
        suspicious = cur.fetchone()[0]

        cur.close()
        conn.close()

        return {
            "total_alerts": total,
            "blocked_ips": blocked,
            "suspicious_ips_historical": suspicious,
            "by_type": [{"type": r[0], "count": r[1]} for r in by_type],
            "by_severity": [{"severity": r[0], "count": r[1]} for r in by_severity],
            "top_attackers": [{"ip": r[0], "count": r[1]} for r in top_attackers]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/stats/summary")
def get_summary():
    try:
        conn = get_db()
        cur = conn.cursor()

        cur.execute("""
            SELECT COUNT(*) FROM alerts
            WHERE created_at >= CURRENT_DATE
        """)
        today_count = cur.fetchone()[0]

        cur.execute("SELECT verdict, COUNT(*) FROM alerts GROUP BY verdict")
        verdict_counts = {row[0]: row[1] for row in cur.fetchall()}

        cur.execute("SELECT COUNT(*) FROM blocked_ips")
        blocked_count = cur.fetchone()[0]

        cur.execute("SELECT MAX(created_at) FROM alerts")
        last_alert = cur.fetchone()[0]

        cur.close()
        conn.close()

        seconds_since_last = None
        is_live = False
        if last_alert:
            seconds_since_last = int((datetime.now() - last_alert).total_seconds())
            is_live = seconds_since_last < 30

        return {
            "alerts_today": today_count,
            "true_positive": verdict_counts.get("true_positive", 0),
            "false_positive": verdict_counts.get("false_positive", 0),
            "pending": verdict_counts.get("pending", 0),
            "blocked_ips": blocked_count,
            "data_source": {
                "is_live": is_live,
                "seconds_since_last_alert": seconds_since_last,
                "last_alert_at": str(last_alert) if last_alert else None
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/stats/timeseries")
def get_timeseries(hours: int = 12):
    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute("""
            SELECT date_trunc('hour', created_at) AS bucket,
                   attack_type,
                   COUNT(*) AS count
            FROM alerts
            WHERE created_at >= NOW() - (%s * INTERVAL '1 hour')
            GROUP BY bucket, attack_type
            ORDER BY bucket ASC
        """, (hours,))
        rows = cur.fetchall()
        cur.close()
        conn.close()
        return {"data": [
            {"time": str(r[0]), "attack_type": r[1], "count": r[2]}
            for r in rows
        ]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/stats/top-ips")
def get_top_ips(limit: int = 10):
    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute("""
            SELECT src_ip,
                   COUNT(*) AS total,
                   COUNT(*) FILTER (WHERE verdict = 'true_positive') AS confirmed
            FROM alerts
            GROUP BY src_ip
            ORDER BY total DESC
            LIMIT %s
        """, (limit,))
        rows = cur.fetchall()
        cur.close()
        conn.close()
        return {"ips": [
            {"ip": r[0], "total": r[1], "confirmed": r[2]}
            for r in rows
        ]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ── Model Performance ──
MODEL_METADATA = {
    "name": "Smart-IDPS Two-Stage Classifier",
    "algorithm": "Stage 1: Logistic Regression (Benign vs Attack) -> Stage 2: Random Forest (Attack Type)",
    "trained_on": "CSE-CIC-IDS2018 (CIC, Canadian Institute for Cybersecurity)",
    "training_rows": 456051,
    "training_accuracy": 98.0,
    "version": "v1.0",
    "last_trained": "2026-06-05"
}

RETRAIN_PRECISION_THRESHOLD = 0.85
RETRAIN_SAMPLE_THRESHOLD = 200
RETRAIN_MIN_REVIEWED = 20

@app.get("/api/model/performance")
def get_model_performance():
    try:
        conn = get_db()
        cur = conn.cursor()

        cur.execute("""
            SELECT
                COUNT(*) FILTER (WHERE verdict = 'true_positive') AS tp,
                COUNT(*) FILTER (WHERE verdict = 'false_positive') AS fp,
                COUNT(*) FILTER (WHERE verdict = 'pending' OR verdict IS NULL) AS pending
            FROM alerts
        """)
        tp, fp, pending = cur.fetchone()
        reviewed = tp + fp

        overall_precision = round(tp / reviewed, 4) if reviewed > 0 else None

        cur.execute("""
            SELECT attack_type,
                   COUNT(*) FILTER (WHERE verdict = 'true_positive') AS tp,
                   COUNT(*) FILTER (WHERE verdict = 'false_positive') AS fp
            FROM alerts
            WHERE verdict IN ('true_positive', 'false_positive')
            GROUP BY attack_type
        """)
        by_type = []
        for row in cur.fetchall():
            t_tp, t_fp = row[1], row[2]
            t_reviewed = t_tp + t_fp
            by_type.append({
                "attack_type": row[0],
                "confirmed": t_tp,
                "false_positive": t_fp,
                "precision": round(t_tp / t_reviewed, 4) if t_reviewed > 0 else None
            })

        cur.close()
        conn.close()

        reasons = []
        should_retrain = False
        if reviewed >= RETRAIN_MIN_REVIEWED and overall_precision is not None and overall_precision < RETRAIN_PRECISION_THRESHOLD:
            should_retrain = True
            reasons.append(f"Analyst-confirmed precision ({overall_precision*100:.1f}%) is below threshold ({RETRAIN_PRECISION_THRESHOLD*100:.0f}%)")
        if pending >= RETRAIN_SAMPLE_THRESHOLD:
            should_retrain = True
            reasons.append(f"{pending} new alerts are pending review, exceeding the {RETRAIN_SAMPLE_THRESHOLD}-sample threshold")
        if not reasons:
            reasons.append("Model performance is within acceptable range; no retraining needed at this time")

        return {
            "metadata": MODEL_METADATA,
            "live_stats": {
                "reviewed_total": reviewed,
                "confirmed_attacks": tp,
                "false_positives": fp,
                "pending_review": pending,
                "overall_precision": overall_precision
            },
            "by_attack_type": by_type,
            "retraining": {
                "recommended": should_retrain,
                "reasons": reasons,
                "thresholds": {
                    "min_precision": RETRAIN_PRECISION_THRESHOLD,
                    "sample_size_trigger": RETRAIN_SAMPLE_THRESHOLD
                }
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ── IP Blocking ──
@app.post("/api/block-ip")
def block_ip(data: BlockIP, _: str = Depends(verify_token)):
    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute("""
            INSERT INTO blocked_ips (ip_address, reason)
            VALUES (%s, %s)
            ON CONFLICT (ip_address) DO NOTHING
        """, (data.ip_address, data.reason))
        conn.commit()
        cur.close()
        conn.close()
        return {"status": "blocked", "ip": data.ip_address}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/blocked-ips")
def get_blocked_ips():
    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute("SELECT ip_address, reason, blocked_at FROM blocked_ips ORDER BY blocked_at DESC")
        rows = cur.fetchall()
        cur.close()
        conn.close()
        return {"count": len(rows), "blocked_ips": [
            {"ip": r[0], "reason": r[1], "blocked_at": str(r[2])}
            for r in rows
        ]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.delete("/api/block-ip/{ip}")
def unblock_ip(ip: str, _: str = Depends(verify_token)):
    try:
        conn = get_db()
        cur = conn.cursor()
        cur.execute("DELETE FROM blocked_ips WHERE ip_address = %s", (ip,))
        conn.commit()
        cur.close()
        conn.close()
        return {"status": "unblocked", "ip": ip}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# ── Historical Analysis ──
@app.post("/api/historical")
def save_historical(results: List[HistoricalResult]):
    try:
        conn = get_db()
        cur = conn.cursor()
        for r in results:
            cur.execute("""
                INSERT INTO historical_analysis
                (src_ip, total_connections, unique_ports, avg_bytes, anomaly_score, is_suspicious)
                VALUES (%s, %s, %s, %s, %s, %s)
            """, (
                r.src_ip, r.total_connections, r.unique_ports,
                r.avg_bytes, r.anomaly_score, r.is_suspicious
            ))
        conn.commit()
        cur.close()
        conn.close()
        return {"status": "saved", "count": len(results)}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/historical")
def get_historical(suspicious_only: bool = False):
    try:
        conn = get_db()
        cur = conn.cursor()
        if suspicious_only:
            cur.execute("""
                SELECT src_ip, total_connections, unique_ports,
                       avg_bytes, anomaly_score, is_suspicious, created_at
                FROM historical_analysis
                WHERE is_suspicious = true
                ORDER BY anomaly_score DESC
            """)
        else:
            cur.execute("""
                SELECT src_ip, total_connections, unique_ports,
                       avg_bytes, anomaly_score, is_suspicious, created_at
                FROM historical_analysis
                ORDER BY created_at DESC
                LIMIT 100
            """)
        rows = cur.fetchall()
        cur.close()
        conn.close()
        return {"count": len(rows), "results": [
            {
                "src_ip": r[0], "total_connections": r[1],
                "unique_ports": r[2], "avg_bytes": r[3],
                "anomaly_score": r[4], "is_suspicious": r[5],
                "created_at": str(r[6])
            } for r in rows
        ]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
