"""SQLite job store. Short-lived connections so the API and worker processes can share it."""
from __future__ import annotations

import json
import sqlite3
import time
from pathlib import Path

SCHEMA = """
CREATE TABLE IF NOT EXISTS jobs (
  id TEXT PRIMARY KEY, dataset_id TEXT NOT NULL, status TEXT NOT NULL,
  stage TEXT, stage_label TEXT, progress REAL DEFAULT 0, message TEXT,
  config TEXT, error TEXT, created REAL, started REAL, finished REAL, cancel INTEGER DEFAULT 0
);
CREATE INDEX IF NOT EXISTS jobs_status ON jobs(status);
"""
TERMINAL = {"done", "failed", "cancelled"}


def connect(path: Path) -> sqlite3.Connection:
    con = sqlite3.connect(path, timeout=30, isolation_level=None)
    con.row_factory = sqlite3.Row
    con.execute("PRAGMA journal_mode=WAL")
    return con


def init(path: Path):
    with connect(path) as con:
        con.executescript(SCHEMA)
        # jobs interrupted by a restart can never finish
        con.execute("UPDATE jobs SET status='failed', error=?, finished=? WHERE status IN ('queued','running')",
                    (json.dumps({"message": "Interrupted by a server restart; please run again."}), time.time()))


def create_job(path: Path, job_id: str, dataset_id: str, config: dict):
    with connect(path) as con:
        con.execute("INSERT INTO jobs (id, dataset_id, status, progress, message, config, created) VALUES (?,?,?,?,?,?,?)",
                    (job_id, dataset_id, "queued", 0.0, "Waiting for a free worker", json.dumps(config), time.time()))


def update(path: Path, job_id: str, **fields):
    if not fields:
        return
    if "error" in fields and not isinstance(fields["error"], (str, type(None))):
        fields["error"] = json.dumps(fields["error"])
    cols = ", ".join(f"{k}=?" for k in fields)
    with connect(path) as con:
        con.execute(f"UPDATE jobs SET {cols} WHERE id=?", (*fields.values(), job_id))


def get(path: Path, job_id: str) -> dict | None:
    with connect(path) as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
    if row is None:
        return None
    d = dict(row)
    d["config"] = json.loads(d["config"] or "{}")
    d["error"] = json.loads(d["error"]) if d["error"] else None
    d["cancel"] = bool(d["cancel"])
    return d


def count(path: Path, statuses: tuple[str, ...]) -> int:
    with connect(path) as con:
        q = f"SELECT COUNT(*) FROM jobs WHERE status IN ({','.join('?' * len(statuses))})"
        return con.execute(q, statuses).fetchone()[0]


def queue_position(path: Path, job_id: str) -> int:
    with connect(path) as con:
        row = con.execute("SELECT created FROM jobs WHERE id=?", (job_id,)).fetchone()
        if row is None:
            return 0
        return con.execute("SELECT COUNT(*) FROM jobs WHERE status='queued' AND created < ?", (row[0],)).fetchone()[0]


def old_jobs(path: Path, before: float) -> list[str]:
    with connect(path) as con:
        return [r[0] for r in con.execute("SELECT id FROM jobs WHERE created < ? AND status IN ('done','failed','cancelled')", (before,))]


def delete(path: Path, job_id: str):
    with connect(path) as con:
        con.execute("DELETE FROM jobs WHERE id=?", (job_id,))