"""Job records in Postgres.

An upload job and a batch run each live in the memory of the task running
them, as they always have; this module is the durable copy. The running
task writes every stage transition, every finished batch row, and a
heartbeat every few seconds. Any task can then answer a status poll or
serve a run's rows from the tables, and a task that dies leaves a record
whose heartbeat stops: a job with no heartbeat for ``STALE_SECONDS`` is
marked failed with that reason the next time anything reads it, instead of
sitting at its last stage forever.

The tables are created by ``ensure_tables`` at startup from ``jobs.sql``,
which is idempotent, so a running database gains them without a migration
step. Nothing here is imported by the batch runner; the runner takes a store
object with the same method names, and tests run it without one.
"""

from __future__ import annotations

import json
import os
import socket
import threading
import time
from pathlib import Path
from typing import Any, Callable, Dict, Iterable, List, LiteralString, Optional, cast

from loguru import logger

from quber.playground import db

JOBS_SQL = Path(__file__).with_name("jobs.sql")

#: Who is running a job: this process, named so a record can say which task
#: it belonged to. On Fargate the hostname is the task's.
OWNER = f"{socket.gethostname()}:{os.getpid()}"

STALE_SECONDS = 90
HEARTBEAT_SECONDS = 10
LOG_TAIL = 200
TERMINAL_STAGES = ("done", "failed")
STOPPED_REASON = "the task running this job stopped"

_UPLOAD_COLUMNS: tuple[str, ...] = (
    "job_id",
    "filename",
    "title",
    "folder",
    "filing_type",
    "year",
    "period",
    "version",
    "track",
    "replace",
    "replace_id",
    "replace_key",
    "doc_key",
    "content_hash",
    "doc_id",
    "source_uri",
    "stage",
    "stages",
    "error",
    "error_info",
    "log",
    "scan_pages",
    "scan_pages_reason",
)
_JSON_COLUMNS = {"stages", "error_info", "log"}


#: One lock for the table creation. CREATE TABLE IF NOT EXISTS is not safe
#: when two processes run it at the same instant, which two tasks of the
#: hosted service starting together do.
_TABLES_LOCK = 7_290_329


def ensure_tables() -> None:
    sql = cast(LiteralString, JOBS_SQL.read_text())
    with db.connect() as conn, conn.transaction():
        conn.execute("SELECT pg_advisory_xact_lock(%s)", (_TABLES_LOCK,))
        conn.execute(sql)


# ---------------------------------------------------------------- uploads


def save_upload(job: Dict[str, Any]) -> None:
    """Write the job as it stands, log tail included, and stamp the heartbeat."""
    values: List[Any] = []
    for column in _UPLOAD_COLUMNS:
        value = job.get(column)
        if column == "log":
            value = list(value or [])[-LOG_TAIL:]
        if column in _JSON_COLUMNS:
            value = json.dumps(value) if value is not None else None
        values.append(value)
    columns: LiteralString = ", ".join(_UPLOAD_COLUMNS)  # type: ignore[assignment]
    placeholders: LiteralString = ", ".join(["%s"] * len(_UPLOAD_COLUMNS))  # type: ignore[assignment]
    updates: LiteralString = ", ".join(f"{c} = EXCLUDED.{c}" for c in _UPLOAD_COLUMNS[1:])  # type: ignore[assignment]
    query: LiteralString = (
        "INSERT INTO ade_playground.upload_jobs ("  # noqa: S608 - column names are literals
        + columns
        + ", owner, heartbeat, updated_at) VALUES ("
        + placeholders
        + ", %s, now(), now()) ON CONFLICT (job_id) DO UPDATE SET "
        + updates
        + ", owner = EXCLUDED.owner, heartbeat = now(), updated_at = now()"
    )
    with db.connect() as conn:
        conn.execute(query, (*values, OWNER))


def _upload_from_row(row: tuple[Any, ...]) -> Dict[str, Any]:
    job = dict(zip(_UPLOAD_COLUMNS, row, strict=False))
    for column in _JSON_COLUMNS:
        value = job.get(column)
        if isinstance(value, str):
            job[column] = json.loads(value)
    job["log"] = job.get("log") or []
    return job


def mark_stale_uploads() -> int:
    """Fail every unfinished job whose heartbeat has stopped. Returns how many."""
    with db.connect() as conn:
        rows = conn.execute(
            """UPDATE ade_playground.upload_jobs
               SET stage = 'failed', error = %s, updated_at = now()
               WHERE stage <> ALL(%s)
                 AND (heartbeat IS NULL OR heartbeat < now() - make_interval(secs => %s))
               RETURNING job_id""",
            (STOPPED_REASON, list(TERMINAL_STAGES), STALE_SECONDS),
        ).fetchall()
    for (job_id,) in rows:
        logger.warning("upload job {} failed: {}", job_id, STOPPED_REASON)
    return len(rows)


def load_upload(job_id: str) -> Optional[Dict[str, Any]]:
    mark_stale_uploads()
    columns: LiteralString = ", ".join(_UPLOAD_COLUMNS)  # type: ignore[assignment]
    with db.connect() as conn:
        row = conn.execute(
            "SELECT " + columns + " FROM ade_playground.upload_jobs WHERE job_id = %s",  # noqa: S608
            (job_id,),
        ).fetchone()
    return _upload_from_row(row) if row else None


def live_uploads() -> List[Dict[str, Any]]:
    """Every upload job still running on any task."""
    mark_stale_uploads()
    columns: LiteralString = ", ".join(_UPLOAD_COLUMNS)  # type: ignore[assignment]
    with db.connect() as conn:
        rows = conn.execute(
            "SELECT " + columns + " FROM ade_playground.upload_jobs WHERE stage <> ALL(%s)",  # noqa: S608
            (list(TERMINAL_STAGES),),
        ).fetchall()
    return [_upload_from_row(r) for r in rows]


def heartbeat_uploads(job_ids: Iterable[str]) -> None:
    ids = list(job_ids)
    if not ids:
        return
    with db.connect() as conn:
        conn.execute(
            "UPDATE ade_playground.upload_jobs SET heartbeat = now() WHERE job_id = ANY(%s)",
            (ids,),
        )


# ---------------------------------------------------------------- batch runs


class BatchStore:
    """The batch runner's durable side. Method names are the runner's contract."""

    def save_run(self, job_id: str, doc_id: int, want: str, questions: List[str]) -> None:
        with db.connect() as conn:
            conn.execute(
                """INSERT INTO ade_playground.batch_runs (job_id, doc_id, want, questions, owner, heartbeat)
                   VALUES (%s, %s, %s, %s, %s, now())
                   ON CONFLICT (job_id) DO NOTHING""",
                (job_id, doc_id, want, json.dumps(questions), OWNER),
            )

    def save_row(
        self, job_id: str, index: int, question: str, error: Optional[str], response: Optional[Dict[str, Any]]
    ) -> None:
        with db.connect() as conn:
            conn.execute(
                """INSERT INTO ade_playground.batch_rows (job_id, index, question, error, response)
                   VALUES (%s, %s, %s, %s, %s)
                   ON CONFLICT (job_id, index) DO UPDATE
                   SET error = EXCLUDED.error, response = EXCLUDED.response, finished_at = now()""",
                (job_id, index, question, error, json.dumps(response) if response is not None else None),
            )
            conn.execute(
                "UPDATE ade_playground.batch_runs SET failed = failed + %s, heartbeat = now(), updated_at = now() WHERE job_id = %s",
                (1 if error else 0, job_id),
            )

    def finish_run(self, job_id: str) -> None:
        with db.connect() as conn:
            conn.execute(
                "UPDATE ade_playground.batch_runs SET running = false, updated_at = now() WHERE job_id = %s",
                (job_id,),
            )

    def mark_stale_runs(self) -> None:
        """A run whose task stopped: every unanswered question becomes an error
        row saying so, and the run is closed, so its counts add up."""
        with db.connect() as conn:
            stale = conn.execute(
                """SELECT job_id, questions FROM ade_playground.batch_runs
                   WHERE running AND (heartbeat IS NULL OR heartbeat < now() - make_interval(secs => %s))""",
                (STALE_SECONDS,),
            ).fetchall()
            for job_id, questions in stale:
                questions = json.loads(questions) if isinstance(questions, str) else questions
                done = {
                    r[0]
                    for r in conn.execute(
                        "SELECT index FROM ade_playground.batch_rows WHERE job_id = %s", (job_id,)
                    ).fetchall()
                }
                missing = [i for i in range(len(questions)) if i not in done]
                for index in missing:
                    conn.execute(
                        """INSERT INTO ade_playground.batch_rows (job_id, index, question, error)
                           VALUES (%s, %s, %s, %s) ON CONFLICT DO NOTHING""",
                        (job_id, index, questions[index], STOPPED_REASON),
                    )
                conn.execute(
                    "UPDATE ade_playground.batch_runs SET running = false, failed = failed + %s, updated_at = now() WHERE job_id = %s",
                    (len(missing), job_id),
                )
                logger.warning(
                    "batch run {} closed: {} ({} unanswered)", job_id, STOPPED_REASON, len(missing)
                )

    def load_run(self, job_id: str) -> Optional[Dict[str, Any]]:
        """The run and its rows as plain data, or None."""
        self.mark_stale_runs()
        with db.connect() as conn:
            run = conn.execute(
                "SELECT doc_id, want, questions, failed, running FROM ade_playground.batch_runs WHERE job_id = %s",
                (job_id,),
            ).fetchone()
            if run is None:
                return None
            rows = conn.execute(
                "SELECT index, question, error, response FROM ade_playground.batch_rows WHERE job_id = %s ORDER BY index",
                (job_id,),
            ).fetchall()
        questions = json.loads(run[2]) if isinstance(run[2], str) else run[2]
        return {
            "job_id": job_id,
            "doc_id": run[0],
            "want": run[1],
            "questions": questions,
            "failed": run[3],
            "running": run[4],
            "rows": [
                {
                    "index": r[0],
                    "question": r[1],
                    "error": r[2],
                    "response": json.loads(r[3]) if isinstance(r[3], str) else r[3],
                }
                for r in rows
            ],
        }

    def running_for(self, doc_id: int) -> bool:
        self.mark_stale_runs()
        with db.connect() as conn:
            row = conn.execute(
                "SELECT 1 FROM ade_playground.batch_runs WHERE doc_id = %s AND running LIMIT 1", (doc_id,)
            ).fetchone()
        return row is not None

    def heartbeat(self, job_ids: Iterable[str]) -> None:
        ids = list(job_ids)
        if not ids:
            return
        with db.connect() as conn:
            conn.execute(
                "UPDATE ade_playground.batch_runs SET heartbeat = now() WHERE job_id = ANY(%s)", (ids,)
            )


# ---------------------------------------------------------------- heartbeat


def start_heartbeat(
    upload_ids: Callable[[], Iterable[str]],
    run_ids: Callable[[], Iterable[str]],
    store: BatchStore,
) -> threading.Thread:
    """Beat, every ``HEARTBEAT_SECONDS``, for every job this process is still
    running. The callables name them at each beat, so a job that finished
    between beats is simply not beaten again."""

    def beat() -> None:
        while True:
            time.sleep(HEARTBEAT_SECONDS)
            try:
                heartbeat_uploads(upload_ids())
                store.heartbeat(run_ids())
            except Exception as exc:
                logger.warning("heartbeat failed: {}", exc)

    thread = threading.Thread(target=beat, name="job-heartbeat", daemon=True)
    thread.start()
    return thread