import asyncio
import json
import uuid
from datetime import UTC, datetime, timedelta

import aiosqlite

from haiku.rag.ingester.queue.models import (
    Job,
    JobOp,
    JobStatus,
    SyncRow,
    SyncStateRow,
)


def _utcnow_iso() -> str:
    return datetime.now(UTC).isoformat()


def _parse_dt(value: str | None) -> datetime | None:
    return datetime.fromisoformat(value) if value else None


def _row_to_job(row: aiosqlite.Row) -> Job:
    extra_text = row["extra"]
    return Job(
        id=row["id"],
        source_id=row["source_id"],
        uri=row["uri"],
        op=JobOp(row["op"]),
        content_hash=row["content_hash"],
        revision=row["revision"],
        status=JobStatus(row["status"]),
        attempts=row["attempts"],
        max_attempts=row["max_attempts"],
        last_error=row["last_error"],
        extra=json.loads(extra_text) if extra_text else None,
        enqueued_at=datetime.fromisoformat(row["enqueued_at"]),
        scheduled_at=datetime.fromisoformat(row["scheduled_at"]),
        claimed_at=_parse_dt(row["claimed_at"]),
        claimed_by=row["claimed_by"],
        completed_at=_parse_dt(row["completed_at"]),
    )


def _row_to_sync_state(row: aiosqlite.Row) -> SyncStateRow:
    return SyncStateRow(
        source_id=row["source_id"],
        uri=row["uri"],
        revision=row["revision"],
        content_hash=row["content_hash"],
        last_seen_at=datetime.fromisoformat(row["last_seen_at"]),
        last_ingested_at=_parse_dt(row["last_ingested_at"]),
    )


class JobRepo:
    def __init__(
        self,
        conn: aiosqlite.Connection,
        lock: asyncio.Lock | None = None,
    ):
        # Row access by name in helpers below.
        conn.row_factory = aiosqlite.Row
        self._conn = conn
        # Serialize repo calls on the shared connection so cursors from one
        # coroutine don't sit "in progress" when another tries to commit.
        # aiosqlite executes statements on a single worker thread, but
        # individual cursors don't finalize until closed or GC'd — SQLite
        # then refuses commit() with "SQL statements in progress". When
        # JobRepo and SyncStateRepo share the same connection, callers must
        # pass the same lock instance so cross-repo calls also serialize.
        self._lock = lock or asyncio.Lock()
        # Notified after a successful enqueue so workers can wake
        # immediately instead of polling on a fixed sleep interval.
        self.job_available = asyncio.Condition()

    async def enqueue(
        self,
        source_id: str,
        uri: str,
        op: JobOp = JobOp.UPSERT,
        *,
        revision: str | None = None,
        content_hash: str | None = None,
        max_attempts: int = 5,
        extra: dict | None = None,
    ) -> Job | None:
        """Enqueue an upsert/delete job. Returns the inserted Job, or None if a
        live (queued/claimed) job already exists for the same (source_id, uri,
        op). The partial unique index enforces atomicity."""
        job_id = str(uuid.uuid4())
        now = _utcnow_iso()
        extra_json = json.dumps(extra) if extra is not None else None
        async with self._lock:
            async with self._conn.execute(
                """
                INSERT INTO jobs (
                    id, source_id, uri, op, content_hash, revision, status,
                    attempts, max_attempts, last_error, extra,
                    enqueued_at, scheduled_at
                ) VALUES (?, ?, ?, ?, ?, ?, 'queued', 0, ?, NULL, ?, ?, ?)
                ON CONFLICT DO NOTHING
                RETURNING *
                """,
                (
                    job_id,
                    source_id,
                    uri,
                    op.value,
                    content_hash,
                    revision,
                    max_attempts,
                    extra_json,
                    now,
                    now,
                ),
            ) as cursor:
                row = await cursor.fetchone()
            await self._conn.commit()
        if row is not None:
            async with self.job_available:
                self.job_available.notify_all()
        return _row_to_job(row) if row else None

    async def claim_next(self, worker_id: str) -> Job | None:
        """Atomically claim the oldest queued job whose scheduled_at <= now.
        Implemented as a single UPDATE ... RETURNING — no SELECT/UPDATE race."""
        now = _utcnow_iso()
        async with self._lock:
            async with self._conn.execute(
                """
                UPDATE jobs
                SET status = 'claimed',
                    claimed_at = ?,
                    claimed_by = ?,
                    attempts = attempts + 1
                WHERE id = (
                    SELECT id FROM jobs
                    WHERE status = 'queued' AND scheduled_at <= ?
                    ORDER BY scheduled_at
                    LIMIT 1
                )
                RETURNING *
                """,
                (now, worker_id, now),
            ) as cursor:
                row = await cursor.fetchone()
            await self._conn.commit()
        return _row_to_job(row) if row else None

    async def get_job(self, job_id: str) -> Job | None:
        async with self._lock:
            async with self._conn.execute(
                "SELECT * FROM jobs WHERE id = ?", (job_id,)
            ) as cursor:
                row = await cursor.fetchone()
        return _row_to_job(row) if row else None

    async def mark_succeeded(self, job_id: str, claimed_by: str) -> bool:
        """Transition a still-claimed job to `succeeded`. Guarded on
        `status='claimed' AND claimed_by=?` so a reaper-resurrected job
        picked up by a different worker isn't clobbered by the original
        slow worker. Returns True when the row was updated."""
        async with self._lock:
            async with self._conn.execute(
                "UPDATE jobs SET status='succeeded', completed_at=? "
                "WHERE id=? AND status='claimed' AND claimed_by=? RETURNING id",
                (_utcnow_iso(), job_id, claimed_by),
            ) as cursor:
                row = await cursor.fetchone()
            await self._conn.commit()
        return row is not None

    async def mark_dead(self, job_id: str, error: str, claimed_by: str) -> bool:
        """Transition a still-claimed job to `dead`. See `mark_succeeded`
        for the guard semantics."""
        async with self._lock:
            async with self._conn.execute(
                "UPDATE jobs SET status='dead', completed_at=?, last_error=? "
                "WHERE id=? AND status='claimed' AND claimed_by=? RETURNING id",
                (_utcnow_iso(), error, job_id, claimed_by),
            ) as cursor:
                row = await cursor.fetchone()
            await self._conn.commit()
        return row is not None

    async def reschedule(
        self, job_id: str, delay_seconds: float, error: str, claimed_by: str
    ) -> bool:
        """Reset a still-claimed job back to `queued` with a future
        scheduled_at. Guarded on `status='claimed' AND claimed_by=?` so a
        slow worker can't clobber a re-claim that happened after the reaper
        reset its claim. Returns True when the row was updated."""
        scheduled = (datetime.now(UTC) + timedelta(seconds=delay_seconds)).isoformat()
        async with self._lock:
            async with self._conn.execute(
                """
                UPDATE jobs
                SET status='queued',
                    scheduled_at=?,
                    claimed_at=NULL,
                    claimed_by=NULL,
                    last_error=?
                WHERE id=? AND status='claimed' AND claimed_by=?
                RETURNING id
                """,
                (scheduled, error, job_id, claimed_by),
            ) as cursor:
                row = await cursor.fetchone()
            await self._conn.commit()
        return row is not None

    async def retry(self, job_id: str) -> Job:
        """Reset a `dead` or `queued` job: status='queued', attempts=0,
        error cleared, scheduled for immediate re-claim. Refuses `claimed`
        rows (would race with the worker still processing) and `succeeded`
        rows (re-ingest via UPSERT instead). Raises KeyError when the row
        is missing or in a non-retryable state."""
        now = _utcnow_iso()
        async with self._lock:
            async with self._conn.execute(
                """
                UPDATE jobs
                SET status='queued',
                    attempts=0,
                    last_error=NULL,
                    claimed_at=NULL,
                    claimed_by=NULL,
                    completed_at=NULL,
                    scheduled_at=?
                WHERE id=? AND status IN ('dead', 'queued')
                RETURNING *
                """,
                (now, job_id),
            ) as cursor:
                row = await cursor.fetchone()
            await self._conn.commit()
        if not row:
            raise KeyError(f"Job {job_id!r} not found or not retryable")
        return _row_to_job(row)

    async def cancel(self, job_id: str) -> bool:
        """True iff a queued/claimed row was removed; terminal jobs aren't
        cancellable (succeeded/dead rows are kept for history)."""
        async with self._lock:
            async with self._conn.execute(
                "DELETE FROM jobs WHERE id=? AND status IN ('queued', 'claimed') RETURNING id",
                (job_id,),
            ) as cursor:
                row = await cursor.fetchone()
            await self._conn.commit()
        return row is not None

    async def list_jobs(
        self,
        *,
        status: JobStatus | None = None,
        source_id: str | None = None,
        uri: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> list[Job]:
        clauses: list[str] = []
        params: list[object] = []
        if status is not None:
            clauses.append("status = ?")
            params.append(status.value)
        if source_id is not None:
            clauses.append("source_id = ?")
            params.append(source_id)
        if uri is not None:
            clauses.append("uri = ?")
            params.append(uri)
        where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
        params.extend([limit, offset])
        async with self._lock:
            async with self._conn.execute(
                f"SELECT * FROM jobs {where} ORDER BY enqueued_at DESC LIMIT ? OFFSET ?",
                params,
            ) as cursor:
                rows = await cursor.fetchall()
        return [_row_to_job(r) for r in rows]

    async def has_pending(self, source_id: str) -> bool:
        """True iff at least one queued/claimed job exists for the source.

        Cheap probe used by pollers to skip sweeps when there's already
        outstanding work — the queue's unique index would dedupe new enqueues
        anyway, so a sweep into a saturated queue is pure wasted listing work.
        """
        async with self._lock:
            async with self._conn.execute(
                "SELECT 1 FROM jobs WHERE source_id=? AND status IN ('queued','claimed') LIMIT 1",
                (source_id,),
            ) as cursor:
                row = await cursor.fetchone()
        return row is not None

    async def counts_by_status(self) -> dict[str, int]:
        async with self._lock:
            async with self._conn.execute(
                "SELECT status, COUNT(*) AS n FROM jobs GROUP BY status"
            ) as cursor:
                rows = await cursor.fetchall()
        return {row["status"]: row["n"] for row in rows}

    async def counts_by_status_since(self, since: datetime) -> dict[str, int]:
        """status -> count of jobs that reached a terminal state at or after
        `since` (by completed_at). Only succeeded/dead set completed_at, so
        those are the only keys returned. Lets a one-shot batch report the
        work it finished, independent of terminal rows from earlier runs."""
        async with self._lock:
            async with self._conn.execute(
                "SELECT status, COUNT(*) AS n FROM jobs "
                "WHERE completed_at >= ? GROUP BY status",
                (since.isoformat(),),
            ) as cursor:
                rows = await cursor.fetchall()
        return {row["status"]: row["n"] for row in rows}

    async def count_succeeded_since(self, seconds: int) -> int:
        """How many jobs reached `succeeded` in the last `seconds` seconds.
        Drives the dashboard's rolling-throughput chips."""
        threshold = (datetime.now(UTC) - timedelta(seconds=seconds)).isoformat()
        async with self._lock:
            async with self._conn.execute(
                "SELECT COUNT(*) AS n FROM jobs WHERE status='succeeded' AND completed_at >= ?",
                (threshold,),
            ) as cursor:
                row = await cursor.fetchone()
        return int(row["n"]) if row else 0

    async def oldest_queued_age_seconds(self) -> float | None:
        """Age (in seconds) of the oldest job sitting in `queued` whose
        scheduled_at is in the past. Returns None when nothing is waiting.
        Tells operators whether work is backing up."""
        now = datetime.now(UTC)
        async with self._lock:
            async with self._conn.execute(
                "SELECT MIN(scheduled_at) AS oldest FROM jobs "
                "WHERE status='queued' AND scheduled_at <= ?",
                (now.isoformat(),),
            ) as cursor:
                row = await cursor.fetchone()
        if not row or row["oldest"] is None:
            return None
        return (now - datetime.fromisoformat(row["oldest"])).total_seconds()

    async def counts_by_source(self, *statuses: str) -> dict[str, int]:
        """source_id → count of jobs in any of the given statuses. Drives the
        dashboard's per-source DLQ and backlog summaries."""
        if not statuses:
            return {}
        placeholders = ",".join("?" * len(statuses))
        async with self._lock:
            async with self._conn.execute(
                f"SELECT source_id, COUNT(*) AS n FROM jobs "
                f"WHERE status IN ({placeholders}) GROUP BY source_id",
                statuses,
            ) as cursor:
                rows = await cursor.fetchall()
        return {row["source_id"]: row["n"] for row in rows}

    async def release_if_claimed(self, job_id: str, claimed_by: str) -> bool:
        """Reset a still-claimed job back to queued, immediately reclaimable.
        Guarded on `status='claimed' AND claimed_by=?` so the cancel-cleanup
        of a slow worker doesn't strip the claim of a different worker that
        re-claimed after a reaper reset. Decrements attempts to undo the
        increment from `claim_next`, since a cancellation isn't a failed
        attempt. Returns True if the row was released."""
        now = _utcnow_iso()
        async with self._lock:
            async with self._conn.execute(
                """
                UPDATE jobs
                SET status='queued',
                    claimed_at=NULL,
                    claimed_by=NULL,
                    scheduled_at=?,
                    attempts=MAX(0, attempts - 1)
                WHERE id=? AND status='claimed' AND claimed_by=?
                RETURNING id
                """,
                (now, job_id, claimed_by),
            ) as cursor:
                row = await cursor.fetchone()
            await self._conn.commit()
        return row is not None

    async def prune_dead(self, source_id: str, uri: str) -> int:
        """Delete dead jobs for the given (source_id, uri). Called after a
        successful DELETE to clear stale UPSERT failures for the same URI —
        the document is gone, so a "couldn't ingest this" entry is no longer
        actionable. Returns the number of rows removed."""
        async with self._lock:
            cursor = await self._conn.execute(
                "DELETE FROM jobs WHERE source_id=? AND uri=? AND status='dead'",
                (source_id, uri),
            )
            rowcount = cursor.rowcount or 0
            await cursor.close()
            await self._conn.commit()
        return rowcount

    async def reap_stale(self, claim_timeout_seconds: int) -> int:
        """Reset claimed jobs whose claimed_at is older than the timeout
        back to `queued`. Decrements `attempts` to undo the increment from
        `claim_next` — a crashed worker isn't a consumed attempt."""
        threshold = (
            datetime.now(UTC) - timedelta(seconds=claim_timeout_seconds)
        ).isoformat()
        async with self._lock:
            cursor = await self._conn.execute(
                """
                UPDATE jobs
                SET status='queued',
                    claimed_at=NULL,
                    claimed_by=NULL,
                    attempts=MAX(0, attempts - 1)
                WHERE status='claimed' AND claimed_at < ?
                """,
                (threshold,),
            )
            rowcount = cursor.rowcount or 0
            await cursor.close()
            await self._conn.commit()
        return rowcount


class SyncStateRepo:
    def __init__(
        self,
        conn: aiosqlite.Connection,
        lock: asyncio.Lock | None = None,
    ):
        conn.row_factory = aiosqlite.Row
        self._conn = conn
        # Pass the same lock instance JobRepo uses when both wrap one
        # connection. See JobRepo for the SQLite cursor + commit constraint.
        self._lock = lock or asyncio.Lock()

    async def get_revision_snapshot(self, source_id: str) -> dict[str, str]:
        """uri -> revision map for URIs that have a stored revision. Sources
        compare current revision against this map to decide UPSERT vs
        UNCHANGED. Rows without a revision (HTTP without ETag, or a worker
        that didn't complete) are excluded — they have no revision to
        compare against; the closing-loop DELETE diff uses list_known_uris
        instead."""
        async with self._lock:
            async with self._conn.execute(
                "SELECT uri, revision FROM sync_state WHERE source_id=? AND revision IS NOT NULL",
                (source_id,),
            ) as cursor:
                rows = await cursor.fetchall()
        return {row["uri"]: row["revision"] for row in rows}

    async def list_known_uris(self, source_id: str) -> set[str]:
        """Every URI the source has ever produced. Used by the closing-loop
        diff in discover() so a URI previously seen but no longer visible
        (FS file deleted, HTTP URL removed from config) emits DELETE."""
        async with self._lock:
            async with self._conn.execute(
                "SELECT uri FROM sync_state WHERE source_id=?",
                (source_id,),
            ) as cursor:
                rows = await cursor.fetchall()
        return {row["uri"] for row in rows}

    async def get_row(self, source_id: str, uri: str) -> SyncStateRow | None:
        async with self._lock:
            async with self._conn.execute(
                "SELECT * FROM sync_state WHERE source_id=? AND uri=?",
                (source_id, uri),
            ) as cursor:
                row = await cursor.fetchone()
        return _row_to_sync_state(row) if row else None

    async def upsert(
        self,
        source_id: str,
        uri: str,
        *,
        revision: str | None = None,
        content_hash: str | None = None,
        ingested: bool = False,
    ) -> None:
        """Insert-or-update the sync_state row. `ingested=True` stamps
        last_ingested_at; otherwise only last_seen_at is bumped.
        `revision=None` and `content_hash=None` leave any existing values
        untouched."""
        now = _utcnow_iso()
        ingested_at = now if ingested else None
        async with self._lock:
            async with self._conn.execute(
                """
                INSERT INTO sync_state (
                    source_id, uri, revision, content_hash, last_seen_at, last_ingested_at
                )
                VALUES (?, ?, ?, ?, ?, ?)
                ON CONFLICT(source_id, uri) DO UPDATE SET
                    revision = COALESCE(excluded.revision, revision),
                    content_hash = COALESCE(excluded.content_hash, content_hash),
                    last_seen_at = excluded.last_seen_at,
                    last_ingested_at = COALESCE(excluded.last_ingested_at, last_ingested_at)
                """,
                (source_id, uri, revision, content_hash, now, ingested_at),
            ):
                pass
            await self._conn.commit()

    async def batch_upsert(self, rows: list[SyncRow]) -> None:
        """Batch insert-or-update sync_state rows in a single transaction."""
        if not rows:
            return
        now = _utcnow_iso()
        async with self._lock:
            for source_id, uri, revision, content_hash, ingested in rows:
                ingested_at = now if ingested else None
                await self._conn.execute(
                    """
                    INSERT INTO sync_state (
                        source_id, uri, revision, content_hash,
                        last_seen_at, last_ingested_at
                    )
                    VALUES (?, ?, ?, ?, ?, ?)
                    ON CONFLICT(source_id, uri) DO UPDATE SET
                        revision = COALESCE(excluded.revision, revision),
                        content_hash = COALESCE(excluded.content_hash, content_hash),
                        last_seen_at = excluded.last_seen_at,
                        last_ingested_at = COALESCE(
                            excluded.last_ingested_at, last_ingested_at
                        )
                    """,
                    (source_id, uri, revision, content_hash, now, ingested_at),
                )
            await self._conn.commit()

    async def delete(self, source_id: str, uri: str) -> None:
        async with self._lock:
            async with self._conn.execute(
                "DELETE FROM sync_state WHERE source_id=? AND uri=?",
                (source_id, uri),
            ):
                pass
            await self._conn.commit()