import sqlite3
from datetime import datetime

DB_PATH = "/shared_data/detections.db"

def init_db():
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    #  WAL mode 
    cursor.execute("PRAGMA journal_mode=WAL;")

    cursor.execute("""
        CREATE TABLE IF NOT EXISTS detections (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            timestamp TEXT NOT NULL,
            src_ip TEXT NOT NULL,
            dst_ip TEXT NOT NULL,
            dst_port INTEGER,
            stage TEXT NOT NULL,
            decision TEXT NOT NULL,
            reconstruction_error REAL NOT NULL,
            threshold REAL NOT NULL,
            processing_latency_ms REAL,
            flow_duration_sec REAL
        )
    """)

    # Migration:
    try:
        cursor.execute("ALTER TABLE detections ADD COLUMN processing_latency_ms REAL")
    except sqlite3.OperationalError:
        pass
    try:
        cursor.execute("ALTER TABLE detections ADD COLUMN flow_duration_sec REAL")
    except sqlite3.OperationalError:
        pass


    cursor.execute("""
        CREATE TABLE IF NOT EXISTS system_status (
            id INTEGER PRIMARY KEY CHECK (id = 1),
            active_flows_count INTEGER NOT NULL,
            last_updated TEXT NOT NULL
        )
    """)

    
    cursor.execute("""
        INSERT OR IGNORE INTO system_status (id, active_flows_count, last_updated)
        VALUES (1, 0, ?)
    """, (datetime.now().isoformat(),))

    conn.commit()
    conn.close()
    print("[DB] Database initialized successfully.")


def log_detection(src_ip, dst_ip, dst_port, stage, decision, reconstruction_error, threshold, processing_latency_ms=None, flow_duration_sec=None):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    cursor.execute("""
        INSERT INTO detections
        (timestamp, src_ip, dst_ip, dst_port, stage, decision, reconstruction_error, threshold, processing_latency_ms, flow_duration_sec)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    """, (
        datetime.now().isoformat(),
        src_ip,
        dst_ip,
        dst_port,
        stage,
        decision,
        float(reconstruction_error),
        float(threshold),
        float(processing_latency_ms) if processing_latency_ms is not None else None,
        float(flow_duration_sec) if flow_duration_sec is not None else None
    ))

    conn.commit()
    conn.close()


def update_active_flow_count(count):
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    cursor.execute("""
        UPDATE system_status
        SET active_flows_count = ?, last_updated = ?
        WHERE id = 1
    """, (count, datetime.now().isoformat()))

    conn.commit()
    conn.close()