"""SQLite database management for the Radix Edge Agent."""

from __future__ import annotations

import sqlite3
import threading
from contextlib import contextmanager
from pathlib import Path
from typing import Optional

from radix_edge.config import get_config

_local = threading.local()

SCHEMA = """
-- Jobs
CREATE TABLE IF NOT EXISTS jobs (
    job_id              TEXT PRIMARY KEY,
    user_id             TEXT NOT NULL,
    name                TEXT,
    image               TEXT NOT NULL,
    command             TEXT,
    args                TEXT,
    env                 TEXT,
    job_type            TEXT NOT NULL DEFAULT 'generic',
    priority            INTEGER NOT NULL DEFAULT 5,
    num_gpus            INTEGER NOT NULL DEFAULT 1,
    gpu_mem_required_mb INTEGER DEFAULT 0,
    timeout_seconds     INTEGER DEFAULT 3600,
    status              TEXT NOT NULL DEFAULT 'queued',
    scheduled_gpus      TEXT,
    container_id        TEXT,
    exit_code           INTEGER,
    error_message       TEXT,
    stdout_tail         TEXT,
    stderr_tail         TEXT,
    runtime_seconds     REAL,
    created_at          TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
    scheduled_at        TEXT,
    started_at          TEXT,
    completed_at        TEXT,
    metadata            TEXT,
    origin              TEXT NOT NULL DEFAULT 'local',
    upstream_job_id     TEXT,
    target_node         TEXT
);

CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status);
CREATE INDEX IF NOT EXISTS idx_jobs_user_status ON jobs(user_id, status);
CREATE INDEX IF NOT EXISTS idx_jobs_created ON jobs(created_at);
CREATE INDEX IF NOT EXISTS idx_jobs_origin ON jobs(origin, status);

-- GPU State (live)
CREATE TABLE IF NOT EXISTS gpus (
    gpu_index       INTEGER PRIMARY KEY,
    uuid            TEXT UNIQUE,
    name            TEXT NOT NULL,
    memory_total_mb INTEGER NOT NULL,
    memory_used_mb  INTEGER NOT NULL DEFAULT 0,
    memory_free_mb  INTEGER NOT NULL DEFAULT 0,
    utilization_pct REAL NOT NULL DEFAULT 0.0,
    temperature_c   REAL DEFAULT 0.0,
    power_draw_w    REAL DEFAULT 0.0,
    power_limit_w   REAL DEFAULT 0.0,
    pcie_gen        INTEGER DEFAULT 0,
    nvlink_peers    TEXT,
    assigned_job_id TEXT,
    status          TEXT NOT NULL DEFAULT 'idle',
    last_seen_at    TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
);

-- GPU Observation History
CREATE TABLE IF NOT EXISTS gpu_observations (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    gpu_index       INTEGER NOT NULL,
    utilization_pct REAL,
    memory_used_mb  INTEGER,
    temperature_c   REAL,
    power_draw_w    REAL,
    observed_at     TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
);

CREATE INDEX IF NOT EXISTS idx_gpu_obs_time ON gpu_observations(observed_at);

-- Scheduler Stats
CREATE TABLE IF NOT EXISTS scheduler_stats (
    job_type        TEXT NOT NULL,
    gpu_type        TEXT NOT NULL,
    mean_runtime    REAL NOT NULL,
    variance        REAL NOT NULL,
    count           INTEGER NOT NULL DEFAULT 0,
    last_updated    REAL NOT NULL,
    PRIMARY KEY (job_type, gpu_type)
);

CREATE TABLE IF NOT EXISTS scheduler_interference (
    job_type        TEXT NOT NULL,
    co_job_type     TEXT NOT NULL,
    slowdown_factor REAL NOT NULL DEFAULT 0.0,
    PRIMARY KEY (job_type, co_job_type)
);

-- User Quotas & Fair Share
CREATE TABLE IF NOT EXISTS user_quotas (
    user_id                 TEXT PRIMARY KEY,
    max_concurrent_gpu_jobs INTEGER NOT NULL DEFAULT 2,
    max_gpus_total          INTEGER NOT NULL DEFAULT 4,
    fair_share_weight       REAL NOT NULL DEFAULT 1.0,
    created_at              TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
    updated_at              TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
);

CREATE TABLE IF NOT EXISTS usage_accounting (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    user_id         TEXT NOT NULL,
    gpu_seconds     REAL NOT NULL DEFAULT 0.0,
    job_count       INTEGER NOT NULL DEFAULT 0,
    window_start    TEXT NOT NULL,
    window_end      TEXT NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_usage_user_window ON usage_accounting(user_id, window_start);

-- FEP State
CREATE TABLE IF NOT EXISTS fep_state (
    state_name        TEXT PRIMARY KEY,
    mean_vector       TEXT NOT NULL,
    std_vector        TEXT NOT NULL,
    observation_count REAL NOT NULL DEFAULT 0.0,
    updated_at        TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
);

CREATE TABLE IF NOT EXISTS fep_decisions (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    latent_state    TEXT NOT NULL,
    state_probs     TEXT NOT NULL,
    observations    TEXT NOT NULL,
    chosen_action   TEXT NOT NULL,
    free_energy     REAL NOT NULL,
    decided_at      TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
);

CREATE INDEX IF NOT EXISTS idx_fep_decisions_time ON fep_decisions(decided_at);

-- Brain Parameters
CREATE TABLE IF NOT EXISTS brain_params (
    key             TEXT PRIMARY KEY DEFAULT 'current',
    tau             REAL NOT NULL DEFAULT 0.8,
    alpha           REAL NOT NULL DEFAULT 0.1,
    ema_beta        REAL NOT NULL DEFAULT 0.15,
    w_queue_depth   REAL NOT NULL DEFAULT 0.5,
    w_gpu_saturation REAL NOT NULL DEFAULT 0.3,
    w_node_pressure REAL NOT NULL DEFAULT 0.2,
    updated_at      TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
);

CREATE TABLE IF NOT EXISTS brain_audit (
    id              INTEGER PRIMARY KEY AUTOINCREMENT,
    bottleneck      TEXT NOT NULL,
    old_params      TEXT NOT NULL,
    new_params      TEXT NOT NULL,
    metrics         TEXT NOT NULL,
    created_at      TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
);

-- General Config
CREATE TABLE IF NOT EXISTS config (
    key             TEXT PRIMARY KEY,
    value           TEXT NOT NULL,
    updated_at      TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
);
"""

SEED_DATA = """
INSERT OR IGNORE INTO brain_params (key) VALUES ('current');
"""


def _get_connection(db_path: Optional[str] = None) -> sqlite3.Connection:
    """Get a thread-local SQLite connection."""
    path = db_path or get_config().db_path
    if not hasattr(_local, "connections"):
        _local.connections = {}
    if path not in _local.connections:
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        conn = sqlite3.connect(path, check_same_thread=False)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA journal_mode=WAL")
        conn.execute("PRAGMA foreign_keys=ON")
        conn.execute("PRAGMA busy_timeout=5000")
        _local.connections[path] = conn
    return _local.connections[path]


@contextmanager
def get_db(db_path: Optional[str] = None):
    """Context manager for database access."""
    conn = _get_connection(db_path)
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise


def init_db(db_path: Optional[str] = None):
    """Initialize the database schema and seed data."""
    conn = _get_connection(db_path)
    conn.executescript(SCHEMA)
    conn.executescript(SEED_DATA)

    # Migration: add upstream and mesh columns for existing databases
    for col_sql in (
        "ALTER TABLE jobs ADD COLUMN origin TEXT NOT NULL DEFAULT 'local'",
        "ALTER TABLE jobs ADD COLUMN upstream_job_id TEXT",
        "ALTER TABLE jobs ADD COLUMN target_node TEXT",
    ):
        try:
            conn.execute(col_sql)
        except sqlite3.OperationalError:
            pass  # Column already exists

    conn.commit()


def close_db():
    """Close all thread-local connections."""
    if hasattr(_local, "connections"):
        for conn in _local.connections.values():
            conn.close()
        _local.connections.clear()
