"""Durable job state and atomic worker claims."""

import sqlite3
from contextlib import contextmanager
from uuid import uuid4


class Conflict(Exception):
    """An idempotency key was reused for a different request."""


class QueueStore:
    def __init__(self, path):
        self.path = path
        with self.connection() as db:
            db.execute("PRAGMA journal_mode=WAL")
            db.executescript("""
                CREATE TABLE IF NOT EXISTS jobs (
                    id TEXT PRIMARY KEY,
                    request_key TEXT UNIQUE NOT NULL,
                    scenario TEXT NOT NULL,
                    state TEXT NOT NULL,
                    attempts INTEGER NOT NULL DEFAULT 0,
                    available_at REAL NOT NULL,
                    lease_token TEXT,
                    lease_until REAL
                );
                CREATE TABLE IF NOT EXISTS events (
                    id INTEGER PRIMARY KEY,
                    job_id TEXT NOT NULL REFERENCES jobs(id),
                    at REAL NOT NULL,
                    message TEXT NOT NULL
                );
                CREATE INDEX IF NOT EXISTS jobs_ready ON jobs(state, available_at);
                CREATE INDEX IF NOT EXISTS jobs_leased ON jobs(state, lease_until);
            """)

    @contextmanager
    def connection(self, write=False):
        db = sqlite3.connect(self.path, timeout=5)
        db.row_factory = sqlite3.Row
        try:
            if write:
                db.execute("BEGIN IMMEDIATE")
            with db:
                yield db
        finally:
            db.close()

    @staticmethod
    def record(db, job_id, now, message):
        db.execute(
            "INSERT INTO events(job_id, at, message) VALUES (?, ?, ?)",
            (job_id, now, message),
        )

    def enqueue(self, key, scenario, now):
        with self.connection(write=True) as db:
            prior = db.execute("SELECT * FROM jobs WHERE request_key = ?", (key,)).fetchone()
            if prior:
                if prior["scenario"] != scenario:
                    raise Conflict("This key already belongs to a different scenario.")
                self.record(db, prior["id"], now, "Duplicate request returned the existing job")
                return prior["id"], False
            job_id = uuid4().hex
            db.execute(
                "INSERT INTO jobs(id, request_key, scenario, state, available_at) VALUES (?, ?, ?, 'queued', ?)",
                (job_id, key, scenario, now),
            )
            self.record(db, job_id, now, "Job queued")
            return job_id, True

    def claim(self, now, lease_seconds):
        with self.connection(write=True) as db:
            expired = db.execute(
                "SELECT id, attempts FROM jobs WHERE state = 'running' AND lease_until <= ?",
                (now,),
            ).fetchall()
            for row in expired:
                state = "dead" if row["attempts"] >= 3 else "queued"
                db.execute("""
                    UPDATE jobs SET state = ?, available_at = ?, lease_token = NULL,
                        lease_until = NULL WHERE id = ?
                """, (state, now, row["id"]))
                message = "Lease expired · attempts exhausted" if state == "dead" else "Lease expired · job recovered"
                self.record(db, row["id"], now, message)
            row = db.execute("""
                SELECT * FROM jobs WHERE state = 'queued' AND available_at <= ?
                ORDER BY available_at, rowid LIMIT 1
            """, (now,)).fetchone()
            if row is None:
                return None
            token = uuid4().hex
            attempt = row["attempts"] + 1
            db.execute("""
                UPDATE jobs SET state = 'running', attempts = ?, lease_token = ?, lease_until = ?
                WHERE id = ?
            """, (attempt, token, now + lease_seconds, row["id"]))
            self.record(db, row["id"], now, f"Worker claimed attempt {attempt}")
            return {**dict(row), "attempts": attempt, "lease_token": token}

    def finish(self, job, success, now):
        attempt = job["attempts"]
        delay = 2 ** attempt
        if success:
            state, message = "delivered", "Consumer acknowledged the event"
        elif attempt >= 3:
            state, message = "dead", "Consumer failed · attempts exhausted"
        else:
            state, message = "queued", f"Consumer failed · retry in {delay}s"
        with self.connection(write=True) as db:
            updated = db.execute("""
                UPDATE jobs SET state = ?, available_at = ?, lease_token = NULL, lease_until = NULL
                WHERE id = ? AND state = 'running' AND lease_token = ? AND lease_until > ?
            """, (state, now + (delay if state == "queued" else 0), job["id"], job["lease_token"], now))
            if updated.rowcount:
                self.record(db, job["id"], now, message)
            return updated.rowcount == 1

    def snapshot(self):
        with self.connection() as db:
            db.execute("BEGIN")
            counts = db.execute("SELECT state, COUNT(*) AS total FROM jobs GROUP BY state").fetchall()
            jobs = db.execute("""
                SELECT id, request_key, scenario, state, attempts, available_at
                FROM jobs ORDER BY rowid DESC LIMIT 24
            """).fetchall()
            events = db.execute("SELECT * FROM events ORDER BY id DESC LIMIT 32").fetchall()
        return {
            "counts": {row["state"]: row["total"] for row in counts},
            "jobs": [dict(row) for row in jobs],
            "events": [dict(row) for row in events],
        }