"""
cortex/batch_worker.py — Background worker for batch ingestion jobs (#1326).

Runs as a daemon thread started at app startup. Polls for pending jobs in the
batch_jobs table and processes them using LocalBatchIngestor.

Architecture:
    - Worker thread polls every 5 seconds for pending jobs
    - Creates its OWN CortexPlugin instance in the worker thread's event loop
      to avoid cross-event-loop aiosqlite issues (#1334)
    - Uses LocalBatchIngestor (in-process, no HTTP self-calls)
    - Updates progress in batch_jobs row after each chunk
    - On startup, resets any 'running' jobs to 'pending' for resume
"""
from __future__ import annotations

import asyncio
import json
import logging
import os
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional

import aiosqlite

logger = logging.getLogger(__name__)

POLL_INTERVAL_S = 5


class BatchWorker:
    """Background worker that processes batch ingest jobs.

    The worker runs in a daemon thread with its own asyncio event loop and
    its own CortexPlugin instance.  This avoids sharing aiosqlite connections
    across event loops, which causes silent failures (#1334).
    """

    def __init__(
        self,
        db_path: str,
        plugin=None,
        *,
        config_path: Optional[str] = None,
    ) -> None:
        self._db_path = db_path
        # Store config for creating a fresh plugin in the worker thread.
        # The caller's plugin is only used as a config/path reference.
        self._config_path = config_path
        if self._config_path is None and plugin is not None:
            # Try to extract config_path from the caller's plugin
            cfg = getattr(plugin, "config", None)
            self._config_path = getattr(cfg, "_config_path", None) if cfg else None
        self._thread: Optional[threading.Thread] = None
        self._stop_event = threading.Event()
        self._started_at: Optional[str] = None
        self._last_error: Optional[str] = None
        self._jobs_processed: int = 0
        self._state: str = "idle"  # idle | starting | running | crashed | stopped

    @property
    def is_alive(self) -> bool:
        """Return True if the worker thread is currently alive."""
        return self._thread is not None and self._thread.is_alive()

    @property
    def status(self) -> dict:
        """Return worker status dict for health/status endpoints."""
        return {
            "alive": self.is_alive,
            "state": self._state,
            "started_at": self._started_at,
            "jobs_processed": self._jobs_processed,
            "last_error": self._last_error,
            "db_path": self._db_path,
        }

    def start(self) -> None:
        """Start the background worker thread."""
        self._stop_event.clear()
        self._state = "starting"
        self._started_at = datetime.now(timezone.utc).isoformat()
        self._thread = threading.Thread(
            target=self._run_loop,
            name="batch-ingest-worker",
            daemon=True,
        )
        self._thread.start()
        logger.info("Batch ingest worker thread started (db=%s)", self._db_path)

    def stop(self) -> None:
        """Signal the worker to stop."""
        self._stop_event.set()
        self._state = "stopped"
        if self._thread and self._thread.is_alive():
            self._thread.join(timeout=10)
        logger.info("Batch ingest worker stopped")

    def _run_loop(self) -> None:
        """Main worker loop — runs in a background thread."""
        # Create a fresh event loop for this thread
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            loop.run_until_complete(self._async_loop())
        except Exception as exc:
            self._state = "crashed"
            self._last_error = str(exc)
            logger.exception("Batch worker loop crashed: %s", exc)
        finally:
            loop.close()

    async def _async_loop(self) -> None:
        """Async loop: init plugin, reset stale jobs, then poll for pending work."""
        # Create a fresh CortexPlugin in the worker's own event loop (#1334).
        # This ensures all aiosqlite connections are bound to THIS loop.
        worker_plugin = None
        try:
            from cortex.api.plugin import CortexPlugin
            worker_plugin = CortexPlugin(
                config_path=self._config_path,
                db_path=self._db_path,
            )
            if worker_plugin.storage is not None:
                await worker_plugin.storage.initialize()
            logger.info(
                "Batch worker: CortexPlugin initialized in worker thread "
                "(storage=%s, init_errors=%s)",
                "ok" if worker_plugin.storage else "NONE",
                worker_plugin._init_errors or "none",
            )
        except Exception as exc:
            self._state = "crashed"
            self._last_error = f"Plugin init failed: {exc}"
            logger.error(
                "Batch worker: failed to create CortexPlugin — "
                "worker cannot process jobs: %s",
                exc,
            )
            return

        if worker_plugin.storage is None:
            self._state = "crashed"
            self._last_error = "Plugin storage is None — cannot process jobs"
            logger.error(
                "Batch worker: plugin.storage is None after init — "
                "worker cannot process jobs. Init errors: %s",
                worker_plugin._init_errors,
            )
            return

        self._state = "running"
        await self._reset_running_jobs()

        while not self._stop_event.is_set():
            try:
                job = await self._pick_next_job()
                if job:
                    await self._process_job(job, worker_plugin)
                    self._jobs_processed += 1
                else:
                    # No work — wait before polling again
                    self._stop_event.wait(POLL_INTERVAL_S)
            except Exception as exc:
                self._last_error = str(exc)
                logger.exception("Error in batch worker poll cycle: %s", exc)
                self._stop_event.wait(POLL_INTERVAL_S)

        # Clean shutdown: close the worker's plugin storage
        try:
            if worker_plugin.storage is not None:
                await worker_plugin.storage.close()
        except Exception:
            pass

    async def _reset_running_jobs(self) -> None:
        """Reset any 'running' jobs to 'pending' on startup (crash recovery)."""
        async with aiosqlite.connect(self._db_path) as db:
            now = datetime.now(timezone.utc).isoformat()
            cursor = await db.execute(
                "UPDATE batch_jobs SET status = 'pending', updated_at = ? "
                "WHERE status = 'running'",
                (now,),
            )
            count = cursor.rowcount
            await db.commit()
            if count:
                logger.info("Batch worker: reset %d stale running job(s) to pending", count)
            else:
                logger.debug("Batch worker: no stale running jobs to reset")

    async def _pick_next_job(self) -> Optional[dict]:
        """Pick the oldest pending job and mark it running."""
        async with aiosqlite.connect(self._db_path) as db:
            db.row_factory = aiosqlite.Row
            async with db.execute(
                "SELECT * FROM batch_jobs WHERE status = 'pending' "
                "ORDER BY created_at ASC LIMIT 1"
            ) as cur:
                row = await cur.fetchone()
                if not row:
                    return None

                job = dict(row)

            now = datetime.now(timezone.utc).isoformat()
            await db.execute(
                "UPDATE batch_jobs SET status = 'running', updated_at = ? WHERE id = ?",
                (now, job["id"]),
            )
            await db.commit()
            return job

    async def _process_job(self, job: dict, worker_plugin) -> None:
        """Process a single batch job using the worker's own plugin instance."""
        job_id = job["id"]
        logger.info("Processing batch job %s (source_type=%s)", job_id, job["source_type"])

        try:
            config = json.loads(job["config_json"])
            source_type = job["source_type"]
            owner_id = config.get("owner_id")
            if not owner_id or owner_id == "default":
                await self._fail_job(job_id, "owner_id is required and cannot be 'default'")
                return

            # Build chunks from the source
            chunks = self._build_chunks(job_id, source_type, config)
            if chunks is None:
                await self._fail_job(job_id, "Failed to build chunks from source")
                return

            total_chunks = len(chunks)
            await self._update_job(job_id, total_chunks=total_chunks)

            if total_chunks == 0:
                await self._complete_job(job_id, 0, 0, [])
                return

            # Process chunks using LocalBatchIngestor with the WORKER's plugin
            from cortex.batch_ingest import IdempotencyStore, LocalBatchIngestor

            data_dir = os.environ.get("CORTEX_DATA_DIR", "/tmp/cortex-data")
            idempotency_path = os.path.join(data_dir, "batch_uploads", job_id, "idempotency.json")
            idempotency = IdempotencyStore(state_path=idempotency_path)

            # Category exclusion from job config (e.g. exclude_categories=["decisions","engineering"])
            exclude_cats = config.get("exclude_categories")
            exclude_set = set(exclude_cats) if exclude_cats else None

            ingestor = LocalBatchIngestor(
                plugin=worker_plugin,
                owner_id=owner_id,
                dry_run=False,
                idempotency=idempotency,
                source_channel="batch_ingest",
                exclude_categories=exclude_set,
            )

            processed = 0
            claims_stored = 0
            errors = []

            for i, chunk in enumerate(chunks):
                if self._stop_event.is_set():
                    await self._update_job_status(job_id, "cancelled")
                    return

                try:
                    if idempotency.is_processed(chunk.chunk_hash):
                        processed += 1
                        await self._update_progress(job_id, processed, claims_stored, errors)
                        continue

                    result = await ingestor._remember_chunk(chunk)

                    if result.ok:
                        claims_stored += result.claims_stored
                        idempotency.mark_processed(chunk.chunk_hash)
                    else:
                        errors.append({
                            "chunk_index": i,
                            "error": result.error or "Unknown error",
                        })
                except Exception as exc:
                    errors.append({
                        "chunk_index": i,
                        "error": str(exc),
                    })

                processed += 1
                await self._update_progress(job_id, processed, claims_stored, errors)

            # Save idempotency state
            idempotency.save()

            await self._complete_job(job_id, processed, claims_stored, errors)
            logger.info(
                "Batch job %s completed: %d/%d chunks, %d claims, %d errors",
                job_id, processed, total_chunks, claims_stored, len(errors),
            )

        except Exception as exc:
            logger.exception("Batch job %s failed", job_id)
            await self._fail_job(job_id, str(exc))

    def _build_chunks(self, job_id: str, source_type: str, config: dict):
        """Build chunk list from the job's source file."""
        from cortex.batch_ingest import FileSource, Mem0JsonlSource, TextSource

        data_dir = os.environ.get("CORTEX_DATA_DIR", "/tmp/cortex-data")
        upload_dir = os.path.join(data_dir, "batch_uploads", job_id)

        try:
            if source_type == "file":
                filename = config.get("filename", "")
                file_path = os.path.join(upload_dir, filename)
                source = FileSource(file_path)
                return list(source.iter_chunks())
            elif source_type == "mem0_jsonl":
                filename = config.get("filename", "")
                file_path = os.path.join(upload_dir, filename)
                source = Mem0JsonlSource(file_path)
                return list(source.iter_chunks())
            elif source_type == "text":
                text = config.get("text", "")
                source = TextSource(text, label=f"batch-{job_id}")
                return list(source.iter_chunks())
            else:
                logger.error("Unknown source_type: %s", source_type)
                return None
        except Exception as exc:
            logger.exception("Failed to build chunks for job %s", job_id)
            return None

    async def _update_job(self, job_id: str, **kwargs) -> None:
        """Update arbitrary fields on the batch_jobs row."""
        async with aiosqlite.connect(self._db_path) as db:
            now = datetime.now(timezone.utc).isoformat()
            sets = ["updated_at = ?"]
            vals = [now]
            for k, v in kwargs.items():
                sets.append(f"{k} = ?")
                vals.append(v)
            vals.append(job_id)
            await db.execute(
                f"UPDATE batch_jobs SET {', '.join(sets)} WHERE id = ?",
                vals,
            )
            await db.commit()

    async def _update_progress(self, job_id: str, processed: int, claims_stored: int, errors: list) -> None:
        """Update job progress after each chunk."""
        async with aiosqlite.connect(self._db_path) as db:
            now = datetime.now(timezone.utc).isoformat()
            await db.execute(
                "UPDATE batch_jobs SET processed_chunks = ?, claims_stored = ?, "
                "errors_json = ?, updated_at = ? WHERE id = ?",
                (processed, claims_stored, json.dumps(errors), now, job_id),
            )
            await db.commit()

    async def _update_job_status(self, job_id: str, status: str) -> None:
        """Update job status."""
        async with aiosqlite.connect(self._db_path) as db:
            now = datetime.now(timezone.utc).isoformat()
            await db.execute(
                "UPDATE batch_jobs SET status = ?, updated_at = ? WHERE id = ?",
                (status, now, job_id),
            )
            await db.commit()

    async def _complete_job(self, job_id: str, processed: int, claims_stored: int, errors: list) -> None:
        """Mark job as completed."""
        async with aiosqlite.connect(self._db_path) as db:
            now = datetime.now(timezone.utc).isoformat()
            await db.execute(
                "UPDATE batch_jobs SET status = 'completed', processed_chunks = ?, "
                "claims_stored = ?, errors_json = ?, updated_at = ?, completed_at = ? "
                "WHERE id = ?",
                (processed, claims_stored, json.dumps(errors), now, now, job_id),
            )
            await db.commit()

    async def _fail_job(self, job_id: str, error: str) -> None:
        """Mark job as failed."""
        async with aiosqlite.connect(self._db_path) as db:
            now = datetime.now(timezone.utc).isoformat()
            errors = [{"error": error}]
            await db.execute(
                "UPDATE batch_jobs SET status = 'failed', errors_json = ?, "
                "updated_at = ? WHERE id = ?",
                (json.dumps(errors), now, job_id),
            )
            await db.commit()
