# retoor <retoor@molodetz.nl>

from __future__ import annotations

import asyncio
import logging
from typing import Any, Awaitable, Callable, Optional

from .context import task_run_scope
from .controller import compute_followup
from .guards import BudgetProbe, Deferral, budget_deferral, retire_reason, run_deferral
from .limits import reserve_run
from .schedule import next_run, now_utc, to_iso
from .store import TaskStore, claim, due_rows

logger = logging.getLogger("devii.tasks.scheduler")


def _audit_task_execute(row: dict[str, Any], result: str, error: str | None) -> None:
    from devplacepy.services.audit import record as audit

    owner_kind = row.get("owner_kind") or "user"
    owner_id = row.get("owner_id") or ""
    audit.record_system(
        "devii.task.execute",
        actor_kind="user" if owner_kind == "user" else owner_kind,
        actor_uid=owner_id if owner_kind == "user" else None,
        actor_role="user" if owner_kind == "user" else owner_kind,
        origin="scheduler",
        via_agent=1,
        result="failure" if error else "success",
        target_type="task",
        target_uid=row.get("uid"),
        summary=f"Devii task {row.get('uid')} for {owner_id or owner_kind} executed",
        metadata={"run_count": row.get("run_count"), "error": error, "result": (result or "")[:140]},
        links=[audit.task(row.get("uid"))],
    )


def _audit_task_blocked(row: dict[str, Any], reason: str) -> None:
    from devplacepy.services.audit import record as audit

    owner_kind = row.get("owner_kind") or "user"
    owner_id = row.get("owner_id") or ""
    audit.record_system(
        "devii.task.blocked",
        actor_kind="user" if owner_kind == "user" else owner_kind,
        actor_uid=owner_id if owner_kind == "user" else None,
        actor_role="user" if owner_kind == "user" else owner_kind,
        origin="scheduler",
        via_agent=1,
        result="denied",
        target_type="task",
        target_uid=row.get("uid"),
        summary=f"Devii task {row.get('uid')} disabled - {reason}",
        metadata={"reason": reason, "run_count": row.get("run_count")},
        links=[audit.task(row.get("uid"))],
    )


PromptExecutor = Callable[[str], Awaitable[str]]
EventCallback = Callable[[str, dict[str, Any], str], None]
OwnerResolver = Callable[
    [dict[str, Any]], Optional[tuple[TaskStore, PromptExecutor, EventCallback]]
]

DEFAULT_TICK_SECONDS = 1.0
DEFAULT_MAX_CONCURRENT = 4
DEFAULT_MAX_FAILURES = 3
DUE_BATCH_LIMIT = 200


def _audit_task_deferred(row: dict[str, Any], reason: str, retry_at: str) -> None:
    from devplacepy.services.audit import record as audit

    owner_kind = row.get("owner_kind") or "user"
    owner_id = row.get("owner_id") or ""
    audit.record_system(
        "devii.task.deferred",
        actor_kind="user" if owner_kind == "user" else owner_kind,
        actor_uid=owner_id if owner_kind == "user" else None,
        actor_role="user" if owner_kind == "user" else owner_kind,
        origin="scheduler",
        via_agent=1,
        result="denied",
        target_type="task",
        target_uid=row.get("uid"),
        summary=f"Devii task {row.get('uid')} postponed - {reason}",
        metadata={"reason": reason, "retry_at": retry_at},
        links=[audit.task(row.get("uid"))],
    )


def defer(store: TaskStore, row: dict[str, Any], postponement: Deferral) -> None:
    retry_at = to_iso(postponement.retry_at)
    store.update(
        row["uid"],
        {
            "status": "pending",
            "next_run_at": retry_at,
            "last_error": postponement.reason,
        },
    )
    logger.info(
        "Task uid=%s postponed to %s: %s", row.get("uid"), retry_at, postponement.reason
    )
    _audit_task_deferred(row, postponement.reason, retry_at)


def retire(store: TaskStore, row: dict[str, Any], reason: str) -> None:
    store.update(
        row["uid"],
        {
            "enabled": False,
            "status": "disabled",
            "next_run_at": None,
            "last_error": reason,
        },
    )
    logger.info("Task uid=%s disabled: %s", row.get("uid"), reason)
    _audit_task_blocked(row, reason)


async def run_row(
    row: dict[str, Any],
    store: TaskStore,
    executor: PromptExecutor,
    on_event: EventCallback,
    max_failures: int = DEFAULT_MAX_FAILURES,
) -> None:
    uid = row["uid"]
    on_event("start", row, "")
    logger.info("Executing task uid=%s", uid)

    try:
        with task_run_scope():
            result = await executor(row["prompt"])
    except Exception as exc:  # noqa: BLE001 - surfaced into the task record
        logger.exception("Task uid=%s crashed", uid)
        store.update(uid, _failure_changes(row, str(exc), max_failures))
        on_event("error", row, str(exc))
        _audit_task_execute(row, "", str(exc))
        return

    changes, finished = compute_followup(row, now_utc())
    changes["last_result"] = result
    changes["last_error"] = None
    changes["failure_count"] = 0
    store.update(uid, changes)
    on_event("done" if not finished else "finished", row, result)
    _audit_task_execute(row, result, None)


def _failure_changes(
    row: dict[str, Any], error: str, max_failures: int
) -> dict[str, Any]:
    failures = int(row.get("failure_count") or 0) + 1
    changes: dict[str, Any] = {
        "failure_count": failures,
        "last_error": error,
        "last_run_at": to_iso(now_utc()),
    }
    upcoming = next_run(
        row.get("kind"), row.get("every_seconds"), row.get("cron"), now_utc()
    )
    if upcoming is None or (max_failures > 0 and failures >= max_failures):
        changes["status"] = "error"
        changes["enabled"] = False
        changes["next_run_at"] = None
    else:
        changes["status"] = "pending"
        changes["next_run_at"] = to_iso(upcoming)
    return changes


class Scheduler:
    def __init__(
        self,
        store: TaskStore,
        executor: PromptExecutor,
        on_event: EventCallback,
        tick_seconds: float = DEFAULT_TICK_SECONDS,
    ) -> None:
        self._store = store
        self._executor = executor
        self._on_event = on_event
        self._tick_seconds = tick_seconds
        self._loop_task: asyncio.Task[None] | None = None

    def start(self) -> None:
        if self._loop_task is None:
            self._store.recover_running()
            self._loop_task = asyncio.create_task(self._loop())
            logger.info("Scheduler started")

    async def stop(self) -> None:
        if self._loop_task is None:
            return
        self._loop_task.cancel()
        try:
            await self._loop_task
        except asyncio.CancelledError:
            pass
        self._loop_task = None
        logger.info("Scheduler stopped")

    async def _loop(self) -> None:
        while True:
            try:
                await self._tick()
            except Exception:  # noqa: BLE001 - the scheduler must never die
                logger.exception("Scheduler tick failed")
            await asyncio.sleep(self._tick_seconds)

    async def _tick(self) -> None:
        now = now_utc()
        for row in self._store.due(to_iso(now)):
            reason = retire_reason(row, now)
            if reason is not None:
                retire(self._store, row, reason)
                continue
            if not claim(self._store.db, row["uid"]):
                continue
            if not self._reserve(row, now):
                defer(self._store, row, run_deferral(self._store.db, row, now))
                continue
            await run_row(row, self._store, self._executor, self._on_event)

    def _reserve(self, row: dict[str, Any], reference: Any) -> bool:
        return reserve_run(
            self._store.db,
            str(row.get("owner_kind") or "user"),
            str(row.get("owner_id") or ""),
            row["uid"],
            reference,
        )


class GlobalScheduler:
    def __init__(
        self,
        db: Any,
        resolve: OwnerResolver,
        tick_seconds: float = DEFAULT_TICK_SECONDS,
        max_concurrent: int = DEFAULT_MAX_CONCURRENT,
        max_failures: int = DEFAULT_MAX_FAILURES,
        budget_exceeded: Optional[BudgetProbe] = None,
    ) -> None:
        self._db = db
        self._resolve = resolve
        self._tick_seconds = tick_seconds
        self._max_concurrent = max_concurrent
        self._max_failures = max_failures
        self._budget_exceeded = budget_exceeded
        self._loop_task: asyncio.Task[None] | None = None
        self._running: dict[str, asyncio.Task[None]] = {}

    @property
    def running(self) -> int:
        return len(self._running)

    def configure(self, max_concurrent: int, max_failures: int) -> None:
        self._max_concurrent = max_concurrent
        self._max_failures = max_failures

    def start(self) -> None:
        if self._loop_task is None:
            self._loop_task = asyncio.create_task(self._loop())
            logger.info("Global task scheduler started")

    async def stop(self) -> None:
        if self._loop_task is None:
            return
        self._loop_task.cancel()
        try:
            await self._loop_task
        except asyncio.CancelledError:
            pass
        self._loop_task = None
        for task in list(self._running.values()):
            task.cancel()
        self._running.clear()
        logger.info("Global task scheduler stopped")

    async def _loop(self) -> None:
        while True:
            try:
                await self._tick()
            except Exception:  # noqa: BLE001 - the scheduler must never die
                logger.exception("Global scheduler tick failed")
            await asyncio.sleep(self._tick_seconds)

    async def _tick(self) -> None:
        now = now_utc()
        free = self._max_concurrent - len(self._running)
        for row in due_rows(self._db, to_iso(now), DUE_BATCH_LIMIT):
            reason = retire_reason(row, now, self._max_failures)
            if reason is not None:
                retire(self._store_for(row), row, reason)
                continue
            if free <= 0:
                continue
            if str(row.get("owner_id") or "") in self._running:
                continue
            postponement = budget_deferral(row, now, self._budget_exceeded)
            if postponement is not None:
                defer(self._store_for(row), row, postponement)
                continue
            if self._dispatch(row, now):
                free -= 1

    def _dispatch(self, row: dict[str, Any], reference: Any) -> bool:
        resolved = self._resolve(row)
        if resolved is None:
            retire(self._store_for(row), row, "owner could not be resolved")
            return False
        store, executor, on_event = resolved

        if not claim(self._db, row["uid"]):
            return False
        if not reserve_run(
            self._db,
            str(row.get("owner_kind") or "user"),
            str(row.get("owner_id") or ""),
            row["uid"],
            reference,
        ):
            defer(store, row, run_deferral(self._db, row, reference))
            return False

        owner_id = str(row.get("owner_id") or "")
        task = asyncio.create_task(
            run_row(row, store, executor, on_event, self._max_failures)
        )
        self._running[owner_id] = task
        task.add_done_callback(lambda _: self._running.pop(owner_id, None))
        return True

    def _store_for(self, row: dict[str, Any]) -> TaskStore:
        return TaskStore(
            self._db,
            str(row.get("owner_kind") or "user"),
            str(row.get("owner_id") or ""),
        )
