"""Async recalc worker for formula fields.

Consumes recalc tasks from the queue, evaluates formulas, and writes results
back to records. Supports crash recovery (resets stale ``calculating`` tasks
on startup) and graceful shutdown.

Lifecycle (managed by app.py lifespan):
    worker = RecalcWorker(db, service)
    await worker.start()    # starts background task + crash recovery
    ...
    await worker.stop()     # waits for shutdown

The worker runs as an asyncio task, polling the queue every ``poll_interval``
seconds. Each task is processed in its own transaction.
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any

from agentkit.bitable.db import BitableDB
from agentkit.bitable.formula.engine import FormulaEngine
from agentkit.bitable.models import FieldType, RecalcStatus
from agentkit.bitable.repository import BitableRepository
from agentkit.bitable.service import BitableService

logger = logging.getLogger(__name__)

_DEFAULT_POLL_INTERVAL = 0.5  # seconds between queue polls
_DEFAULT_REAPER_INTERVAL = 300  # 5 minutes
_DEFAULT_STALE_THRESHOLD = 600  # 10 minutes


class RecalcWorker:
    """Background worker that processes formula recalc tasks.

    Usage::

        worker = RecalcWorker(db, service)
        await worker.start()
        # ... worker runs in background ...
        await worker.stop()
    """

    def __init__(
        self,
        db: BitableDB,
        service: BitableService,
        poll_interval: float = _DEFAULT_POLL_INTERVAL,
        reaper_interval: float = _DEFAULT_REAPER_INTERVAL,
        stale_threshold: float = _DEFAULT_STALE_THRESHOLD,
    ) -> None:
        self._db = db
        self._service = service
        self._repo = BitableRepository(db)
        self._poll_interval = poll_interval
        self._reaper_interval = reaper_interval
        self._stale_threshold = stale_threshold
        self._task: asyncio.Task[None] | None = None
        self._reaper_task: asyncio.Task[None] | None = None
        self._stop_event = asyncio.Event()
        # Per-table formula engines (cached, rebuilt when fields change)
        self._engines: dict[str, FormulaEngine] = {}

    async def start(self) -> None:
        """Start the worker. Performs crash recovery first."""
        # Crash recovery: reset stale 'calculating' tasks to 'pending'.
        # On startup, all calculating tasks are stale (worker was down).
        # Use threshold=0 to reset all calculating tasks immediately.
        reset_count = await self._repo.reset_stale_recalc_tasks(stale_threshold=0.0)
        if reset_count > 0:
            logger.info("RecalcWorker: reset %d stale tasks to pending", reset_count)

        self._stop_event.clear()
        self._task = asyncio.create_task(self._run(), name="recalc-worker")
        self._reaper_task = asyncio.create_task(self._run_reaper(), name="recalc-reaper")
        logger.info("RecalcWorker started")

    async def stop(self) -> None:
        """Gracefully stop the worker."""
        self._stop_event.set()
        if self._task is not None:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
        if self._reaper_task is not None:
            self._reaper_task.cancel()
            try:
                await self._reaper_task
            except asyncio.CancelledError:
                pass
            self._reaper_task = None
        logger.info("RecalcWorker stopped")

    async def _run(self) -> None:
        """Main worker loop — poll queue, process tasks."""
        while not self._stop_event.is_set():
            try:
                # Atomic claim (P1 #6): tasks are marked 'calculating' in the
                # same transaction, so concurrent workers never grab the same task.
                tasks = await self._repo.claim_recalc_tasks(limit=10)
                if not tasks:
                    await asyncio.sleep(self._poll_interval)
                    continue

                # P1 #7: sort tasks by topological order so formula-to-formula
                # dependencies resolve correctly within a batch. Tasks for the
                # same record are ordered so that if B depends on A, A is
                # processed first (A's result is written back before B reads it).
                tasks = await self._sort_by_topological_order(tasks)

                for task in tasks:
                    if self._stop_event.is_set():
                        break
                    await self.process_task(task)
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("RecalcWorker error in main loop")
                await asyncio.sleep(self._poll_interval)

    async def _sort_by_topological_order(self, tasks: list[Any]) -> list[Any]:
        """Sort claimed tasks so dependencies are processed first (P1 #7).

        Groups tasks by table_id, builds (or reuses) the engine to get the
        topological order, and assigns each task a sort key of
        ``(table_id, record_id, topo_index)``. Tasks for fields not in the
        DAG get topo_index = infinity (processed last).
        """
        if len(tasks) <= 1:
            return tasks

        # Build topo index per table: {table_id: {field_id: position}}
        topo_index: dict[str, dict[str, int]] = {}
        table_ids = {t.table_id for t in tasks}
        for tid in table_ids:
            engine = await self._get_or_build_engine(tid)
            if engine is None:
                topo_index[tid] = {}
                continue
            order = engine.topological_order()
            topo_index[tid] = {fid: i for i, fid in enumerate(order)}

        def _key(t: Any) -> tuple[str, str, int]:
            idx = topo_index.get(t.table_id, {}).get(t.field_id, 1 << 30)
            return (t.table_id, t.record_id, idx)

        return sorted(tasks, key=_key)

    async def _run_reaper(self) -> None:
        """Reaper loop — reset stale calculating tasks periodically.

        Only resets tasks older than ``stale_threshold`` (P1 #10), so active
        tasks being processed by a live worker are not interrupted.
        """
        while not self._stop_event.is_set():
            try:
                await asyncio.sleep(self._reaper_interval)
                count = await self._repo.reset_stale_recalc_tasks(
                    stale_threshold=self._stale_threshold
                )
                if count > 0:
                    logger.info(
                        "RecalcWorker reaper: reset %d stale tasks (threshold=%ss)",
                        count,
                        self._stale_threshold,
                    )
            except asyncio.CancelledError:
                break
            except Exception:
                logger.exception("RecalcWorker reaper error")

    async def process_task(self, task: Any) -> None:
        """Process a single recalc task: evaluate formula → write result.

        The task is expected to already be in ``calculating`` status when
        called from the worker loop (atomic claim sets it). When called
        synchronously via ``service.process_recalc_task``, this method
        marks it calculating first (idempotent — re-marking is harmless).
        """
        # Idempotent: mark calculating (no-op if already calculating via claim).
        await self._repo.update_recalc_status(task.id, RecalcStatus.calculating)

        try:
            field = await self._repo.get_field(task.field_id)
            if field is None or field.field_type != FieldType.formula:
                await self._repo.update_recalc_status(
                    task.id, RecalcStatus.error, "Field not found or not a formula"
                )
                return

            formula_expr = field.config.get("formula_expr", "")
            if not formula_expr:
                await self._repo.update_recalc_status(
                    task.id, RecalcStatus.error, "No formula_expr in field config"
                )
                return

            engine = await self._get_or_build_engine(task.table_id)
            if engine is None:
                await self._repo.update_recalc_status(
                    task.id, RecalcStatus.error, "No formula fields in table"
                )
                return

            record = await self._repo.get_record(task.record_id)
            if record is None:
                await self._repo.update_recalc_status(
                    task.id, RecalcStatus.error, "Record not found"
                )
                return

            deps = engine.get_dependencies(task.field_id)
            column_values: dict[str, list[Any]] = {}
            for dep_field_id in deps:
                column_values[dep_field_id] = await self._repo.get_column_values(
                    task.table_id, dep_field_id
                )

            result = engine.evaluate(
                task.field_id,
                row_values=record.values,
                column_values=column_values,
            )

            await self._repo.set_formula_value(task.record_id, task.field_id, result)
            await self._repo.update_recalc_status(task.id, RecalcStatus.done)

        except Exception as e:
            logger.exception("RecalcWorker: error processing task %s", task.id)
            await self._repo.update_recalc_status(task.id, RecalcStatus.error, str(e)[:500])

    async def _get_or_build_engine(self, table_id: str) -> FormulaEngine | None:
        """Get or build a FormulaEngine for a table.

        Returns None if the table has no formula fields.
        """
        if table_id in self._engines:
            return self._engines[table_id]

        fields = await self._repo.list_fields(table_id)
        formula_fields = [f for f in fields if f.field_type == FieldType.formula]

        if not formula_fields:
            return None

        engine = FormulaEngine()
        for f in formula_fields:
            formula_expr = f.config.get("formula_expr", "")
            if formula_expr:
                try:
                    engine.add_formula(f.id, formula_expr)
                except Exception:
                    logger.exception("RecalcWorker: failed to register formula for field %s", f.id)

        self._engines[table_id] = engine
        return engine

    def invalidate_engine(self, table_id: str) -> None:
        """Invalidate the cached formula engine for a table (call when fields change)."""
        self._engines.pop(table_id, None)
