"""In-process runner loop — heartbeat, poll, and execute jobs without subprocesses."""

import asyncio
import collections
import inspect
import logging
import random
import signal
import threading
import time
from typing import Callable, Optional

from ..api_objects.agent_config.context import agent_config_context
from .. import id_helpers
from ..rest_api.client import OpikApi
from ..rest_api.core.api_error import ApiError
from ..rest_api.types.local_runner_job import LocalRunnerJob
from . import registry

LOGGER = logging.getLogger(__name__)

POLL_IDLE_INTERVAL_SECONDS = 0.5
_CANCELLED_JOBS_TTL_SECONDS = 300
_CANCELLED_JOBS_MAX_SIZE = 10_000


def _inject_trace_id(inputs: dict, trace_id: str) -> None:
    """Merge trace_id into inputs["opik_args"]["trace"]["id"].

    Skips injection when opik_args is explicitly set to None — callers that
    set the key to None signal they own that slot and we must not clobber it.
    """
    existing = inputs.get("opik_args")
    if existing is None and "opik_args" in inputs:
        return
    opik_args = dict(existing) if isinstance(existing, dict) else {}
    trace_args = dict(opik_args.get("trace") or {})
    trace_args["id"] = trace_id
    opik_args["trace"] = trace_args
    inputs["opik_args"] = opik_args


class InProcessRunnerLoop:
    def __init__(
        self,
        api: OpikApi,
        runner_id: str,
        shutdown_event: threading.Event,
        heartbeat_interval_seconds: float = 5.0,
        backoff_cap_seconds: float = 30.0,
    ) -> None:
        self._api = api
        self._runner_id = runner_id
        self._shutdown_event = shutdown_event
        self._heartbeat_interval_seconds = heartbeat_interval_seconds
        self._backoff_cap_seconds = backoff_cap_seconds
        self._cancelled_jobs: collections.OrderedDict[str, float] = (
            collections.OrderedDict()
        )
        self._lock = threading.Lock()
        self._job_queue: asyncio.Queue[LocalRunnerJob] = asyncio.Queue()
        self._loop: Optional[asyncio.AbstractEventLoop] = None

    def run(self) -> None:
        self._install_signal_handlers()

        heartbeat_thread = threading.Thread(
            target=self._heartbeat_loop,
            daemon=True,
        )
        heartbeat_thread.start()

        poll_thread = threading.Thread(
            target=self._poll_loop,
            daemon=True,
        )
        poll_thread.start()

        self._run_job_loop()

    def _poll_loop(self) -> None:
        backoff = 1.0

        while not self._shutdown_event.is_set():
            try:
                job = self._api.runners.next_job(self._runner_id)
            except ApiError as e:
                if e.status_code == 204:
                    job = None
                else:
                    LOGGER.debug("Poll error (API %s)", e.status_code, exc_info=True)
                    self._backoff_wait(backoff)
                    backoff = min(backoff * 2, self._backoff_cap_seconds)
                    continue
            except Exception:
                LOGGER.debug("Error polling for jobs", exc_info=True)
                self._backoff_wait(backoff)
                backoff = min(backoff * 2, self._backoff_cap_seconds)
                continue

            if job is None:
                backoff = 1.0
                self._shutdown_event.wait(POLL_IDLE_INTERVAL_SECONDS)
                continue

            backoff = 1.0
            if self._loop is not None:
                self._loop.call_soon_threadsafe(self._job_queue.put_nowait, job)

    def _heartbeat_loop(self) -> None:
        while not self._shutdown_event.is_set():
            try:
                resp = self._api.runners.heartbeat(self._runner_id)

                cancelled_job_ids = resp.cancelled_job_ids or []
                now = time.monotonic()
                with self._lock:
                    for jid in cancelled_job_ids:
                        self._cancelled_jobs[jid] = now
                    self._prune_cancelled_jobs(now)

            except ApiError as e:
                if e.status_code == 410:
                    LOGGER.info("Runner deregistered (410), shutting down")
                    self._shutdown_event.set()
                    return
                LOGGER.debug("Heartbeat error", exc_info=True)
            except Exception:
                LOGGER.debug("Heartbeat error", exc_info=True)

            self._shutdown_event.wait(self._heartbeat_interval_seconds)

    def _run_job_loop(self) -> None:
        self._loop = asyncio.new_event_loop()
        try:
            self._loop.run_until_complete(self._job_consumer())
        finally:
            self._loop.close()

    async def _job_consumer(self) -> None:
        tasks: set[asyncio.Task] = set()
        while not self._shutdown_event.is_set():
            try:
                job = await asyncio.wait_for(
                    self._job_queue.get(),
                    timeout=0.5,
                )
            except asyncio.TimeoutError:
                continue

            task = asyncio.create_task(self._execute_job(job))
            tasks.add(task)
            task.add_done_callback(tasks.discard)

        if tasks:
            await asyncio.gather(*tasks, return_exceptions=True)

    async def _execute_job(self, job: LocalRunnerJob) -> None:
        job_id = job.id or ""
        agent_name = job.agent_name or ""
        inputs = dict(job.inputs) if job.inputs else {}

        with self._lock:
            if job_id in self._cancelled_jobs:
                del self._cancelled_jobs[job_id]
                return

        entrypoints = registry.get_all()
        entry = entrypoints.get(agent_name)
        if entry is None:
            LOGGER.error("Unknown agent '%s' for job %s", agent_name, job_id)
            try:
                self._api.runners.report_job_result(
                    job_id=job_id,
                    status="failed",
                    error=f"Unknown agent: {agent_name}",
                )
            except Exception:
                LOGGER.debug("Failed to report error for job %s", job_id, exc_info=True)
            return

        func: Callable = entry["func"]
        mask_id = job.mask_id

        trace_id = id_helpers.generate_id()
        _inject_trace_id(inputs, trace_id)

        try:
            timeout = job.timeout
            if inspect.iscoroutinefunction(func):
                with agent_config_context(mask_id):
                    coro = func(**inputs)
                    if timeout:
                        result = await asyncio.wait_for(coro, timeout=timeout)
                    else:
                        result = await coro
            else:

                def _run_with_mask() -> object:
                    with agent_config_context(mask_id):
                        return func(**inputs)

                if timeout:
                    result = await asyncio.wait_for(
                        asyncio.get_running_loop().run_in_executor(
                            None, _run_with_mask
                        ),
                        timeout=timeout,
                    )
                else:
                    result = await asyncio.get_running_loop().run_in_executor(
                        None, _run_with_mask
                    )

            if not isinstance(result, (dict, str, int, float, bool, list, type(None))):
                result = str(result)

            try:
                self._api.runners.report_job_result(
                    job_id=job_id,
                    status="completed",
                    result={"result": result}
                    if not isinstance(result, dict)
                    else result,
                    trace_id=trace_id,
                )
            except Exception:
                LOGGER.debug(
                    "Failed to report result for job %s", job_id, exc_info=True
                )
        except asyncio.TimeoutError:
            LOGGER.warning("Job %s timed out", job_id)
            self._api.runners.report_job_result(
                job_id=job_id,
                status="failed",
                error="Job timed out",
                trace_id=trace_id,
            )
        except Exception as e:
            LOGGER.error("Job %s failed: %s", job_id, e, exc_info=True)
            try:
                self._api.runners.report_job_result(
                    job_id=job_id,
                    status="failed",
                    error=f"{type(e).__name__}: {e}",
                    trace_id=trace_id,
                )
            except Exception:
                LOGGER.debug("Failed to report error for job %s", job_id, exc_info=True)

    def _prune_cancelled_jobs(self, now: float) -> None:
        """Remove expired entries and enforce max size. Caller must hold self._lock."""
        cutoff = now - _CANCELLED_JOBS_TTL_SECONDS
        while self._cancelled_jobs:
            oldest_key, oldest_time = next(iter(self._cancelled_jobs.items()))
            if oldest_time > cutoff:
                break
            del self._cancelled_jobs[oldest_key]
        while len(self._cancelled_jobs) > _CANCELLED_JOBS_MAX_SIZE:
            self._cancelled_jobs.popitem(last=False)

    def _backoff_wait(self, backoff: float) -> None:
        wait = min(backoff, self._backoff_cap_seconds) * (0.5 + random.random() * 0.5)
        self._shutdown_event.wait(wait)

    def _install_signal_handlers(self) -> None:
        shutdown = self._shutdown_event

        def handler(signum: int, frame: object) -> None:
            LOGGER.info("Received signal %s, shutting down", signum)
            shutdown.set()

        try:
            signal.signal(signal.SIGTERM, handler)
            signal.signal(signal.SIGINT, handler)
        except ValueError:
            LOGGER.warning("Cannot install signal handlers outside main thread")
