"""Persistent Work Queue — guarantees every LLM-dependent job completes.

SQLite-backed queue that survives server restarts. When an LLM call fails
(model offline, API down, rate limit), the work is queued and retried with
exponential backoff until it succeeds. The system ALWAYS accomplishes the job.

Architecture principle: **No silent failures, no lost work.**
- Every LLM-dependent operation that fails is queued for retry.
- The queue processor checks LLM health before attempting work.
- When the LLM becomes available, all pending work is drained.
- Users are informed of queued work via heartbeat + API.

Job types map to service methods:
  ingest         → kb_service.ingest(slug)
  compile        → kb_service.compile_wiki(slug)
  reflect        → kb_service.reflect(slug)
  research       → kb_service.research(topic, query)
  lint           → kb_service.lint(slug)
  synthesize     → kb_service.synthesize(slug)
  consolidate    → kb_service.consolidate(slug)
  memory_daily   → memory_engine.generate_daily_memory()
  memory_compact → memory_engine.compact(layer)
  wiki_job       → kb_service full pipeline (ingest+compile+reflect+lint)
"""

from __future__ import annotations

import asyncio
import json
import logging
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

import aiosqlite

if TYPE_CHECKING:
    from agntspace.automation.heartbeat import Heartbeat

log = logging.getLogger(__name__)

# Retry delays: exponential backoff capped at 5 minutes, then repeat forever
_RETRY_DELAYS = [10, 30, 60, 120, 300]  # seconds

# Max attempts before marking as dead_letter (0 = infinite)
_MAX_ATTEMPTS = 0  # Never give up by default

# How often the queue processor checks for work (seconds)
_POLL_INTERVAL = 30

# When LLM is unhealthy and there are pending jobs, re-probe every N seconds
# instead of waiting for the 15-minute heartbeat cycle.
_REPROBE_INTERVAL = 60


class WorkQueue:
    """Persistent SQLite-backed work queue with guaranteed completion."""

    def __init__(
        self,
        db: aiosqlite.Connection,
        heartbeat: Heartbeat | None = None,
    ) -> None:
        self._db = db
        self._heartbeat = heartbeat
        self._handlers: dict[str, Any] = {}  # job_type → async callable
        self._task: asyncio.Task | None = None  # type: ignore[type-arg]
        self._running = False
        self._llm_healthy = True
        self._active_model: str = ""
        self._drain_event = asyncio.Event()  # Signal to wake processor immediately
        self._clock: callable | None = None  # Override for testing: returns datetime
        # Optional callback fired when LLM transitions from unhealthy → healthy.
        # Used to auto-restore quarantined watcher files.
        self._restore_quarantine_cb: callable | None = None
        # Optional LLM provider for self-healing health probes.
        # When set, the work queue re-probes the LLM every _REPROBE_INTERVAL
        # seconds instead of waiting up to 15 min for the heartbeat.
        self._llm: object | None = None
        self._last_reprobe: float = 0.0  # monotonic timestamp of last probe

    def _now(self) -> datetime:
        """Current time — overridable for testing."""
        if self._clock:
            return self._clock()
        return datetime.now(UTC)

    # ── Schema ──

    async def init_schema(self) -> None:
        """Create work queue tables if they don't exist."""
        await self._db.executescript("""
            CREATE TABLE IF NOT EXISTS work_queue (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                job_type TEXT NOT NULL,
                payload TEXT NOT NULL DEFAULT '{}',
                status TEXT NOT NULL DEFAULT 'pending',
                attempts INTEGER NOT NULL DEFAULT 0,
                max_attempts INTEGER NOT NULL DEFAULT 0,
                next_retry TEXT,
                created_at TEXT NOT NULL,
                updated_at TEXT NOT NULL,
                started_at TEXT,
                completed_at TEXT,
                error TEXT,
                result TEXT
            );

            CREATE INDEX IF NOT EXISTS idx_wq_status ON work_queue(status);
            CREATE INDEX IF NOT EXISTS idx_wq_next_retry ON work_queue(next_retry);
            CREATE INDEX IF NOT EXISTS idx_wq_job_type ON work_queue(job_type);
        """)
        await self._db.commit()

    # ── Handler Registration ──

    def register_handler(self, job_type: str, handler: Any) -> None:
        """Register an async handler for a job type.

        Handler signature: async def handler(payload: dict) -> dict | None
        """
        self._handlers[job_type] = handler
        log.debug("Work queue: registered handler for '%s'", job_type)

    # ── LLM Health ──

    def set_llm_health(self, healthy: bool, model: str = "") -> None:
        """Update LLM availability status. Called by health monitor."""
        was_unhealthy = not self._llm_healthy
        self._llm_healthy = healthy
        if model:
            self._active_model = model
        # Wake the processor if LLM just came back online
        if healthy and was_unhealthy:
            log.info("Work queue: LLM available (%s) — draining pending jobs", model or self._active_model)
            self._drain_event.set()
            # Also auto-restore quarantined files so the watcher retries them.
            if self._restore_quarantine_cb:
                try:
                    self._restore_quarantine_cb()
                except Exception:
                    log.exception("Work queue: auto-restore-quarantine callback failed")

    @property
    def llm_healthy(self) -> bool:
        return self._llm_healthy

    @property
    def active_model(self) -> str:
        return self._active_model

    async def _reprobe_llm_health(self) -> bool:
        """Lightweight LLM probe when unhealthy and pending jobs exist.

        Bridges the gap between the 30s poll interval and the 15-min
        heartbeat cycle so Ollama coming online is detected within ~60s
        instead of up to 15 minutes.
        """
        import time

        now = time.monotonic()
        if now - self._last_reprobe < _REPROBE_INTERVAL:
            return False  # too soon since last probe

        self._last_reprobe = now

        if not self._llm:
            return False  # no provider to probe

        try:
            from agntspace.llm.base import Message as LLMMessage

            await asyncio.wait_for(
                self._llm.complete(  # type: ignore[union-attr]
                    messages=[LLMMessage(role="user", content="ping")],
                    system="Reply with exactly: pong",
                    max_tokens=10,
                    temperature=0,
                ),
                timeout=10,
            )
            # LLM is back — update health
            model_name = getattr(self._llm, "model_name", getattr(self._llm, "_model", ""))
            self.set_llm_health(True, model=model_name)
            log.info("Work queue: LLM reprobe succeeded — draining pending jobs")
            return True
        except Exception:
            # If Ollama mode, try auto-starting it
            _model = getattr(self._llm, "_model", "") or ""
            if _model.startswith("ollama/") or _model.startswith("ollama_chat/"):
                from agntspace.llm.ollama_discovery import ensure_ollama_running
                if ensure_ollama_running():
                    log.info("Work queue: Ollama auto-started, marking healthy")
                    model_name = getattr(self._llm, "model_name", _model)
                    self.set_llm_health(True, model=model_name)
                    return True
            log.debug("Work queue: LLM reprobe still failing")
            return False

    # ── Enqueue ──

    async def enqueue(
        self,
        job_type: str,
        payload: dict | None = None,
        max_attempts: int = _MAX_ATTEMPTS,
        deduplicate: bool = True,
    ) -> int | None:
        """Add a job to the queue. Returns job ID, or None if deduplicated.

        Args:
            job_type: Handler key (e.g., 'ingest', 'compile')
            payload: JSON-serializable dict with job parameters
            max_attempts: 0 = retry forever (default)
            deduplicate: If True, skip if identical pending job exists
        """
        payload = payload or {}
        payload_json = json.dumps(payload, sort_keys=True)
        now = self._now().isoformat()

        if deduplicate:
            # Check for existing pending/failed job with same type and payload
            async with self._db.execute(
                "SELECT id FROM work_queue WHERE job_type = ? AND payload = ? AND status IN ('pending', 'failed')",
                (job_type, payload_json),
            ) as cursor:
                existing = await cursor.fetchone()
                if existing:
                    log.debug("Work queue: dedup — %s already queued (id=%d)", job_type, existing[0])
                    return None

        async with self._db.execute(
            """INSERT INTO work_queue (job_type, payload, status, max_attempts, next_retry, created_at, updated_at)
               VALUES (?, ?, 'pending', ?, ?, ?, ?)""",
            (job_type, payload_json, max_attempts, now, now, now),
        ) as cursor:
            job_id = cursor.lastrowid
        await self._db.commit()

        log.info("Work queue: enqueued %s (id=%d, payload=%s)", job_type, job_id, payload_json[:200])

        if self._heartbeat:
            await self._heartbeat.record_pulse(
                "automation",
                f"Queued {job_type}: {self._describe_job(job_type, payload)}",
            )

        # Wake the processor
        self._drain_event.set()
        return job_id

    # ── Queue Status ──

    async def get_status(self) -> dict:
        """Get queue overview for API/UI."""
        counts: dict[str, int] = {}
        async with self._db.execute(
            "SELECT status, COUNT(*) FROM work_queue GROUP BY status",
        ) as cursor:
            async for row in cursor:
                counts[row[0]] = row[1]

        pending_jobs: list[dict] = []
        async with self._db.execute(
            """SELECT id, job_type, payload, attempts, next_retry, error, created_at
               FROM work_queue WHERE status IN ('pending', 'failed')
               ORDER BY created_at ASC LIMIT 50""",
        ) as cursor:
            async for row in cursor:
                pending_jobs.append({
                    "id": row[0], "job_type": row[1],
                    "payload": json.loads(row[2]),
                    "attempts": row[3], "next_retry": row[4],
                    "last_error": row[5], "created_at": row[6],
                })

        return {
            "llm_healthy": self._llm_healthy,
            "active_model": self._active_model,
            "counts": counts,
            "pending_jobs": pending_jobs,
            "total_pending": counts.get("pending", 0) + counts.get("failed", 0),
        }

    async def get_history(self, limit: int = 50) -> list[dict]:
        """Get recently completed/dead_letter jobs."""
        jobs: list[dict] = []
        async with self._db.execute(
            """SELECT id, job_type, payload, status, attempts, completed_at, error, result
               FROM work_queue WHERE status IN ('completed', 'dead_letter')
               ORDER BY completed_at DESC LIMIT ?""",
            (limit,),
        ) as cursor:
            async for row in cursor:
                jobs.append({
                    "id": row[0], "job_type": row[1],
                    "payload": json.loads(row[2]),
                    "status": row[3], "attempts": row[4],
                    "completed_at": row[5], "error": row[6],
                    "result": row[7],
                })
        return jobs

    # ── Manual Controls ──

    async def retry_job(self, job_id: int) -> bool:
        """Reset a failed/dead_letter job to pending for immediate retry."""
        now = self._now().isoformat()
        async with self._db.execute(
            """UPDATE work_queue SET status = 'pending', next_retry = ?, updated_at = ?, error = NULL
               WHERE id = ? AND status IN ('failed', 'dead_letter')""",
            (now, now, job_id),
        ) as cursor:
            if cursor.rowcount == 0:
                return False
        await self._db.commit()
        self._drain_event.set()
        return True

    async def retry_all(self) -> int:
        """Reset all failed/dead_letter jobs to pending."""
        now = self._now().isoformat()
        async with self._db.execute(
            """UPDATE work_queue SET status = 'pending', next_retry = ?, updated_at = ?, error = NULL
               WHERE status IN ('failed', 'dead_letter')""",
            (now, now),
        ) as cursor:
            count = cursor.rowcount
        await self._db.commit()
        if count:
            self._drain_event.set()
        return count

    async def cancel_job(self, job_id: int) -> bool:
        """Cancel a pending/failed job."""
        async with self._db.execute(
            """DELETE FROM work_queue WHERE id = ? AND status IN ('pending', 'failed', 'dead_letter')""",
            (job_id,),
        ) as cursor:
            if cursor.rowcount == 0:
                return False
        await self._db.commit()
        return True

    async def clear_completed(self, older_than_days: int = 7) -> int:
        """Remove completed jobs older than N days."""
        from datetime import timedelta
        cutoff = (self._now() - timedelta(days=older_than_days)).isoformat()
        async with self._db.execute(
            "DELETE FROM work_queue WHERE status = 'completed' AND completed_at < ?",
            (cutoff,),
        ) as cursor:
            count = cursor.rowcount
        await self._db.commit()
        return count

    # ── Processing Loop ──

    def start(self) -> None:
        """Start the background queue processor."""
        if self._task and not self._task.done():
            return
        self._running = True
        self._task = asyncio.create_task(self._run())
        log.info("Work queue: processor started")

    def stop(self) -> None:
        """Stop the queue processor."""
        self._running = False
        self._drain_event.set()  # Wake to exit
        if self._task:
            self._task.cancel()
            self._task = None
        log.info("Work queue: processor stopped")

    async def _run(self) -> None:
        """Main processing loop — runs forever, drains queue when LLM is available."""
        await asyncio.sleep(5)  # Let other services initialize

        # Immediately drain any pending jobs from before shutdown
        try:
            processed = await self._process_batch()
            if processed > 0:
                log.info("Work queue: drained %d pending jobs on startup", processed)
        except Exception:
            log.debug("Work queue: startup drain failed, will retry in loop", exc_info=True)

        while self._running:
            try:
                # Wait for signal or poll interval
                try:
                    await asyncio.wait_for(self._drain_event.wait(), timeout=_POLL_INTERVAL)
                except TimeoutError:
                    pass
                self._drain_event.clear()

                if not self._running:
                    break

                # Process all ready jobs
                processed = await self._process_batch()
                if processed > 0:
                    log.info("Work queue: processed %d jobs", processed)

            except asyncio.CancelledError:
                break
            except Exception:
                log.exception("Work queue: processor error")
                await asyncio.sleep(10)

    async def _process_batch(self) -> int:
        """Process all jobs that are ready. Returns count of jobs attempted."""
        now = self._now().isoformat()
        processed = 0

        # Fetch ready jobs: pending with next_retry <= now, or NULL next_retry
        async with self._db.execute(
            """SELECT id, job_type, payload, attempts, max_attempts
               FROM work_queue
               WHERE status IN ('pending', 'failed')
                 AND (next_retry IS NULL OR next_retry <= ?)
               ORDER BY created_at ASC""",
            (now,),
        ) as cursor:
            jobs = await cursor.fetchall()

        # If LLM is unhealthy and there are pending jobs, try a quick reprobe
        # before giving up for another 30s cycle.  This closes the 15-minute
        # blind spot between heartbeat ticks.
        if jobs and not self._llm_healthy:
            await self._reprobe_llm_health()

        for row in jobs:
            if not self._running:
                break

            job_id, job_type, payload_json, attempts, max_attempts = row
            payload = json.loads(payload_json)

            # Check if handler exists
            handler = self._handlers.get(job_type)
            if not handler:
                log.warning("Work queue: no handler for job type '%s' (id=%d) — skipping", job_type, job_id)
                continue

            # Check LLM health before attempting LLM-dependent work
            if not self._llm_healthy:
                log.debug("Work queue: LLM unavailable — deferring %s (id=%d)", job_type, job_id)
                continue

            processed += 1
            await self._execute_job(job_id, job_type, payload, handler, attempts, max_attempts)

        return processed

    async def _execute_job(
        self,
        job_id: int,
        job_type: str,
        payload: dict,
        handler: Any,
        attempts: int,
        max_attempts: int,
    ) -> None:
        """Execute a single job with error handling and retry scheduling."""
        now = self._now().isoformat()
        attempts += 1

        # Mark as running
        await self._db.execute(
            "UPDATE work_queue SET status = 'running', started_at = ?, attempts = ?, updated_at = ? WHERE id = ?",
            (now, attempts, now, job_id),
        )
        await self._db.commit()

        try:
            result = await handler(payload)
            result_json = json.dumps(result) if result else None

            # Success
            completed_at = self._now().isoformat()
            await self._db.execute(
                """UPDATE work_queue
                   SET status = 'completed', completed_at = ?, result = ?, error = NULL, updated_at = ?
                   WHERE id = ?""",
                (completed_at, result_json, completed_at, job_id),
            )
            await self._db.commit()
            log.info("Work queue: completed %s (id=%d, attempts=%d)", job_type, job_id, attempts)

        except Exception as exc:
            error_msg = f"{type(exc).__name__}: {exc}"
            log.warning(
                "Work queue: %s failed (id=%d, attempt=%d): %s",
                job_type, job_id, attempts, error_msg,
            )

            # Check if we should give up (only if max_attempts > 0)
            if max_attempts > 0 and attempts >= max_attempts:
                await self._db.execute(
                    """UPDATE work_queue
                       SET status = 'dead_letter', error = ?, updated_at = ?
                       WHERE id = ?""",
                    (error_msg, self._now().isoformat(), job_id),
                )
                await self._db.commit()
                log.error("Work queue: %s exhausted max attempts (id=%d) — moved to dead letter", job_type, job_id)
                if self._heartbeat:
                    await self._heartbeat.record_pulse(
                        "error",
                        f"Work queue: {job_type} failed after {attempts} attempts: {error_msg[:200]}",
                    )
                return

            # Schedule retry with exponential backoff
            delay_idx = min(attempts - 1, len(_RETRY_DELAYS) - 1)
            delay = _RETRY_DELAYS[delay_idx]
            from datetime import timedelta
            next_retry = (self._now() + timedelta(seconds=delay)).isoformat()

            await self._db.execute(
                """UPDATE work_queue
                   SET status = 'failed', error = ?, next_retry = ?, updated_at = ?
                   WHERE id = ?""",
                (error_msg, next_retry, self._now().isoformat(), job_id),
            )
            await self._db.commit()
            log.info(
                "Work queue: %s scheduled for retry in %ds (id=%d, attempt=%d)",
                job_type, delay, job_id, attempts,
            )

            # Detect LLM-specific failures and update health
            from agntspace.llm.base import LLMError
            if isinstance(exc, LLMError) and not exc.retryable:
                self._llm_healthy = False
                log.warning("Work queue: LLM marked unhealthy — %s", error_msg[:100])
                if self._heartbeat:
                    await self._heartbeat.record_pulse(
                        "error",
                        f"LLM unavailable ({self._active_model}): {error_msg[:200]}. "
                        f"Queued work will retry when service is restored.",
                    )

    # ── Helpers ──

    @staticmethod
    def _describe_job(job_type: str, payload: dict) -> str:
        """Human-readable job description for notifications."""
        slug = payload.get("slug", "")
        if job_type == "ingest":
            return f"ingest files in KB '{slug}'"
        if job_type == "compile":
            return f"compile wiki for KB '{slug}'"
        if job_type == "reflect":
            return f"reflect on KB '{slug}'"
        if job_type == "research":
            return f"research '{payload.get('query', '')}' in '{payload.get('topic', '')}'"
        if job_type == "lint":
            return f"lint KB '{slug}'"
        if job_type == "wiki_job":
            return f"full wiki pipeline for KB '{slug}'"
        if job_type == "memory_daily":
            return "daily memory consolidation"
        if job_type == "memory_compact":
            return f"compact memory layer '{payload.get('layer', '')}'"
        return f"{job_type}: {json.dumps(payload)[:100]}"
