"""Persistent extraction progress and per-batch worker ownership."""
import json
import time
import uuid

from db.connection import execute, query

LEASE_SECONDS = 300


def claim(batch_id, total):
    """Claim a batch atomically, allowing retry after failure or expired lease."""
    owner = str(uuid.uuid4())
    now = time.time()
    changed = execute(
        """INSERT INTO extraction_jobs (batch_id, owner, state, total, updated_at)
           VALUES (%s, %s, 'running', %s, %s)
           ON CONFLICT(batch_id) DO UPDATE SET owner=excluded.owner, state='running',
             total=excluded.total, done=0, failed=0, classes_json='{}',
             result_json=NULL, error=NULL, updated_at=excluded.updated_at
           WHERE extraction_jobs.state != 'running' OR extraction_jobs.updated_at < %s""",
        (batch_id, owner, total, now, now - LEASE_SECONDS),
    )
    return owner if changed else None


def heartbeat(batch_id, owner):
    return execute(
        "UPDATE extraction_jobs SET updated_at=%s WHERE batch_id=%s AND owner=%s AND state='running'",
        (time.time(), batch_id, owner),
    )


def progress(batch_id, owner, completed, result):
    rows = query(
        "SELECT failed, classes_json FROM extraction_jobs WHERE batch_id=%s AND owner=%s AND state='running'",
        (batch_id, owner),
    )
    if not rows:
        return
    failed = rows[0]["failed"]
    classes = json.loads(rows[0]["classes_json"])
    if not result.get("ok"):
        failed += 1
    else:
        doc_class = (result.get("extraction") or {}).get("doc_class", "unknown")
        classes[doc_class] = classes.get(doc_class, 0) + 1
    execute(
        """UPDATE extraction_jobs SET done=%s, failed=%s, classes_json=%s, updated_at=%s
           WHERE batch_id=%s AND owner=%s AND state='running'""",
        (completed, failed, json.dumps(classes), time.time(), batch_id, owner),
    )


def finish(batch_id, owner, result=None, error=None):
    state = "done" if result and result.get("ok") and not error else "failed"
    return execute(
        """UPDATE extraction_jobs SET state=%s, result_json=%s, error=%s, updated_at=%s
           WHERE batch_id=%s AND owner=%s AND state='running'""",
        (state, json.dumps(result, ensure_ascii=False) if result is not None else None,
         error or ("extraction failed" if state == "failed" else None), time.time(), batch_id, owner),
    )


def get(batch_id):
    """Mark an abandoned lease interrupted; callers may explicitly retry it."""
    execute(
        """UPDATE extraction_jobs SET state='interrupted', error='worker stopped; retry the batch'
           WHERE batch_id=%s AND state='running' AND updated_at < %s""",
        (batch_id, time.time() - LEASE_SECONDS),
    )
    rows = query("SELECT * FROM extraction_jobs WHERE batch_id=%s", (batch_id,))
    if not rows:
        return None
    job = rows[0]
    job["classes"] = json.loads(job.pop("classes_json"))
    result_json = job.pop("result_json")
    job["result"] = json.loads(result_json) if result_json else None
    return job