database.py

"""
database.py
Handles all SQLite persistence for Cyrex emergency requests.
"""

import sqlite3
from datetime import datetime
from pathlib import Path
from typing import Optional

DB_PATH = Path(__file__).parent / "cyrex.db"

SCHEMA = """
CREATE TABLE IF NOT EXISTS requests (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
    contact TEXT NOT NULL,
    location TEXT NOT NULL,
    description TEXT NOT NULL,
    num_people INTEGER DEFAULT 1,
    disaster_type TEXT,
    urgency_level TEXT,
    urgency_score INTEGER,
    assigned_team TEXT,
    status TEXT DEFAULT 'Pending',
    reasoning TEXT,
    created_at TEXT,
    updated_at TEXT
);
"""


def get_connection() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn


def init_db() -> None:
    conn = get_connection()
    conn.execute(SCHEMA)
    conn.commit()
    conn.close()


def insert_request(data: dict) -> int:
    """Insert a new emergency request and return its id."""
    conn = get_connection()
    now = datetime.utcnow().isoformat(timespec="seconds")
    cur = conn.execute(
        """
        INSERT INTO requests
            (name, contact, location, description, num_people,
             disaster_type, urgency_level, urgency_score,
             assigned_team, status, reasoning, created_at, updated_at)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            data.get("name"),
            data.get("contact"),
            data.get("location"),
            data.get("description"),
            data.get("num_people", 1),
            data.get("disaster_type"),
            data.get("urgency_level"),
            data.get("urgency_score"),
            data.get("assigned_team"),
            data.get("status", "Pending"),
            data.get("reasoning"),
            now,
            now,
        ),
    )
    conn.commit()
    request_id = cur.lastrowid
    conn.close()
    return request_id


def update_status(request_id: int, status: str) -> None:
    conn = get_connection()
    conn.execute(
        "UPDATE requests SET status = ?, updated_at = ? WHERE id = ?",
        (status, datetime.utcnow().isoformat(timespec="seconds"), request_id),
    )
    conn.commit()
    conn.close()


def get_all_requests(status_filter: Optional[str] = None):
    conn = get_connection()
    if status_filter and status_filter != "All":
        rows = conn.execute(
            "SELECT * FROM requests WHERE status = ? ORDER BY urgency_score DESC, created_at ASC",
            (status_filter,),
        ).fetchall()
    else:
        rows = conn.execute(
            "SELECT * FROM requests ORDER BY urgency_score DESC, created_at ASC"
        ).fetchall()
    conn.close()
    return [dict(r) for r in rows]


def get_request(request_id: int):
    conn = get_connection()
    row = conn.execute("SELECT * FROM requests WHERE id = ?", (request_id,)).fetchone()
    conn.close()
    return dict(row) if row else None


def get_stats():
    conn = get_connection()
    total = conn.execute("SELECT COUNT(*) c FROM requests").fetchone()["c"]
    pending = conn.execute(
        "SELECT COUNT(*) c FROM requests WHERE status = 'Pending'"
    ).fetchone()["c"]
    in_progress = conn.execute(
        "SELECT COUNT(*) c FROM requests WHERE status = 'Team Dispatched'"
    ).fetchone()["c"]
    resolved = conn.execute(
        "SELECT COUNT(*) c FROM requests WHERE status = 'Resolved'"
    ).fetchone()["c"]
    critical = conn.execute(
        "SELECT COUNT(*) c FROM requests WHERE urgency_level = 'Critical' AND status != 'Resolved'"
    ).fetchone()["c"]
    conn.close()
    return {
        "total": total,
        "pending": pending,
        "in_progress": in_progress,
        "resolved": resolved,
        "critical": critical,
    }