from __future__ import annotations

import asyncio
import logging
from typing import Any

from govoplan_core.core.runtime_coordination import (
    RuntimeIdentity,
    heartbeat_runtime_node,
    register_runtime_node,
    runtime_identity,
    stop_runtime_node,
)
from govoplan_core.db.session import get_database


logger = logging.getLogger("govoplan.runtime")


class RuntimeNodeAgent:
    """Register one process in the shared runtime directory and heartbeat it."""

    def __init__(
        self,
        *,
        settings: object,
        software_version: str,
        module_ids: tuple[str, ...],
        role: str | None = None,
        node_id: str | None = None,
        queues: tuple[str, ...] | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        self.settings = settings
        self.identity: RuntimeIdentity = runtime_identity(
            settings,
            software_version=software_version,
            module_ids=module_ids,
            role=role,
            node_id=node_id,
            queues=queues,
        )
        self.metadata = dict(metadata or {})
        self.draining = False
        self.coordination_healthy = False
        self._task: asyncio.Task[None] | None = None
        self._stopping = False

    @property
    def heartbeat_seconds(self) -> int:
        return max(
            2,
            int(getattr(self.settings, "runtime_heartbeat_seconds", 15)),
        )

    async def start(self) -> None:
        await asyncio.to_thread(self._register)
        self._task = asyncio.create_task(
            self._heartbeat_loop(),
            name=f"govoplan-runtime-heartbeat:{self.identity.node_id}",
        )

    async def stop(self) -> None:
        self._stopping = True
        task = self._task
        self._task = None
        if task is not None:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass
        try:
            await asyncio.to_thread(self._mark_stopped)
        except Exception:  # noqa: BLE001 - shutdown must continue
            logger.exception(
                "runtime node stop marker failed node_id=%s",
                self.identity.node_id,
            )

    def _register(self) -> None:
        with get_database().SessionLocal() as session:
            node = register_runtime_node(
                session,
                self.identity,
                metadata=self.metadata,
            )
            session.commit()
        self.draining = node.state == "draining"
        self.coordination_healthy = True

    def _heartbeat(self) -> None:
        with get_database().SessionLocal() as session:
            node = heartbeat_runtime_node(
                session,
                self.identity,
                metadata=self.metadata,
            )
            session.commit()
        self.draining = node.state == "draining"
        self.coordination_healthy = True

    def _mark_stopped(self) -> None:
        with get_database().SessionLocal() as session:
            stop_runtime_node(session, self.identity)
            session.commit()

    async def _heartbeat_loop(self) -> None:
        while not self._stopping:
            await asyncio.sleep(self.heartbeat_seconds)
            await self._heartbeat_once()

    async def _heartbeat_once(self) -> bool:
        try:
            await asyncio.to_thread(self._heartbeat)
        except asyncio.CancelledError:
            raise
        except Exception:  # noqa: BLE001 - a later heartbeat can recover
            self.coordination_healthy = False
            logger.exception(
                "runtime heartbeat failed node_id=%s",
                self.identity.node_id,
            )
            return False
        self.coordination_healthy = True
        return True


__all__ = ["RuntimeNodeAgent"]
