from __future__ import annotations

import asyncio
from contextlib import suppress
from uuid import uuid4

from ai_orchestrator.config.settings import Settings
from ai_orchestrator.orchestrators.base import OrchestratorEngine
from ai_orchestrator.schemas.task_request import TaskRequest
from ai_orchestrator.schemas.task_status import TaskLifecycleStatus


class TaskWorker:
    def __init__(
        self,
        *,
        settings: Settings,
        store,
        queue,
        orchestrators: dict[str, OrchestratorEngine],
        worker_id: str | None = None,
    ) -> None:
        self.settings = settings
        self.store = store
        self.queue = queue
        self.orchestrators = orchestrators
        self.worker_id = worker_id or f"worker_{uuid4().hex[:8]}"

    async def run_once(self) -> bool:
        job = self.queue.lease_next(
            worker_id=self.worker_id,
            lease_seconds=self.settings.job_lease_seconds,
        )
        if job is None:
            return False

        heartbeat = asyncio.create_task(self._heartbeat_loop(job.job_id))
        try:
            task = self.store.get(job.task_id)
            if task is None:
                self.queue.mark_failed(job.job_id, f"Task not found: {job.task_id}")
                return True

            if task.status == TaskLifecycleStatus.cancelled:
                self.queue.mark_cancelled(job.job_id)
                return True

            engine = self.orchestrators[job.engine_name]
            request = TaskRequest(goal=task.goal, repo_path=task.repo_path, constraints=task.constraints)
            await engine.submit_task(job.task_id, request)
            final_task = self.store.get(job.task_id)
            if final_task is None:
                self.queue.mark_failed(job.job_id, "Task disappeared during execution.")
            elif final_task.status == TaskLifecycleStatus.cancelled:
                self.queue.mark_cancelled(job.job_id)
            elif final_task.status == TaskLifecycleStatus.failed:
                self.queue.mark_failed(job.job_id, final_task.error_message or "Task failed.")
            else:
                self.queue.mark_completed(job.job_id)
        except Exception as exc:
            task = self.store.get(job.task_id)
            if task is not None and task.status not in TaskLifecycleStatus.terminal():
                task.mark_terminal(
                    TaskLifecycleStatus.failed,
                    summary="Task execution failed unexpectedly.",
                    error_message=str(exc),
                )
                self.store.save(task)
            self.queue.mark_failed(job.job_id, str(exc))
        finally:
            heartbeat.cancel()
            with suppress(asyncio.CancelledError):
                await heartbeat
        return True

    async def run_until_idle(self, idle_cycles: int = 3) -> None:
        seen_idle = 0
        while seen_idle < idle_cycles:
            did_work = await self.run_once()
            if did_work:
                seen_idle = 0
            else:
                seen_idle += 1
                await asyncio.sleep(self.settings.worker_poll_interval_seconds)

    async def run_forever(self) -> None:
        while True:
            did_work = await self.run_once()
            if not did_work:
                await asyncio.sleep(self.settings.worker_poll_interval_seconds)

    async def _heartbeat_loop(self, job_id: str) -> None:
        interval = max(1.0, self.settings.job_lease_seconds / 3)
        while True:
            await asyncio.sleep(interval)
            self.queue.heartbeat(
                job_id=job_id,
                worker_id=self.worker_id,
                lease_seconds=self.settings.job_lease_seconds,
            )
