# heartbeat.py
from __future__ import annotations

import asyncio
import socket
import logging
import time

import functools
from dataclasses import dataclass
from typing import Any, Dict, Optional, Literal

# Use the InfluxDBConnectionPool singleton
from ..influx.influx_client import pool as influx_pool

log = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


def _now_ts_s() -> int:
    return int(time.time())


def _default_host() -> str:
    # hostname for tags
    return socket.gethostname()





@dataclass(frozen=True, slots=True)
class HeartbeatConfig:
    """
    Configuration for heartbeat points.

    - service: required logical daemon/service name
    - version: required version string (e.g. git tag, semver, setuptools-scm)
    - host: default hostname if not set
    - measurement: defaults to "heartbeat"
    """
    service: str
    version: str
    host: str = ""
    measurement: str = "heartbeat"

    def normalized_host(self) -> str:
        return self.host or _default_host()



class Heartbeat:
    """
    Reusable heartbeat emitter using InfluxDBConnectionPool.
    Emits points with payload grammar:
    {
        "measurement": "heartbeat",
        "fields": {"alive": 1/0, "start_ts": <epoch seconds>},
        "tags": {"host": "...", "service": "...", "version": "..."}
    }
    """

    def __init__(
        self,
        cfg: HeartbeatConfig,
        *,
        start_ts: Optional[int] = None,
        target: Literal["ha", "mons"] = "ha",
        database: Optional[str] = None,
        secrets: Optional[dict] = None,
    ) -> None:
        self._cfg = cfg
        self._start_ts = _now_ts_s() if start_ts is None else int(start_ts)
        self._target = target
        self._database = database or ("building_power" if target == "ha" else "building")
        if secrets is not None and influx_pool._secrets is None:
            influx_pool.initialize(secrets)
        self._closed = False

    def close(self) -> None:
        self._closed = True
    
    def __enter__(self) -> "Heartbeat":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    async def __aenter__(self) -> "Heartbeat":
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        # executor shutdown is sync; do it in a thread so it doesn't block loop
        await asyncio.to_thread(self.close)

    @property
    def start_ts(self) -> int:
        return self._start_ts

    def payload(self, *, alive: int) -> Dict[str, Any]:
        """
        Build the payload dict EXACTLY matching your grammar.
        """
        alive_i = int(alive)
        if alive_i not in (0, 1):
            raise ValueError("alive must be 0 or 1")

        return {
            "measurement": self._cfg.measurement,
            "fields": {
                "alive": alive_i,
                "start_ts": int(self._start_ts),
            },
            "tags": {
                "host": self._cfg.normalized_host(),
                "service": self._cfg.service,
                "version": self._cfg.version,
            },
        }

    # -------- Sync API --------

    def write(self, *, alive: int = 1) -> bool:
        """
        Synchronously write a single heartbeat point using the pool.
        Returns True on success, False on failure.
        """
        point = self.payload(alive=alive)
        if self._target == "ha":
            return influx_pool.write_ha(self._database, point)
        else:
            return influx_pool.write_mons(self._database, point)

    # -------- Async API --------

    async def write_async(self, *, alive: int = 1) -> bool:
        """
        asyncio-friendly write using the pool in a thread.
        """
        point = self.payload(alive=alive)
        if self._target == "ha":
            return await asyncio.to_thread(influx_pool.write_ha, self._database, point)
        else:
            return await asyncio.to_thread(influx_pool.write_mons, self._database, point)


    @classmethod
    def create_with_pool(
        cls,
        *,
        service: str,
        version: str,
        target: Literal["ha", "mons"] = "mons",
        database: Optional[str] = None,
        secrets: Optional[dict] = None,
    ) -> "Heartbeat":
        """
        Factory method to create a Heartbeat instance using the InfluxDBConnectionPool.
        """
        cfg = HeartbeatConfig(service=service, version=version)
        return cls(cfg, target=target, database=database, secrets=secrets)


async def heartbeat_loop(
    service: str,
    version: str,
    secrets: dict,
    stop_event: asyncio.Event,
    interval_s: float = 30.0,
    target: str = "mons",
    database: str | None = None
) -> None:
    """
    Background coroutine that writes heartbeat to InfluxDB.
    
    Args:
        service: Service name for the heartbeat tag
        version: Service version
        secrets: Secrets dict containing InfluxDB credentials
        stop_event: Event to signal shutdown
        interval_s: Interval between heartbeats in seconds
        target: Target InfluxDB ("mons" or "ha")
        database: Database name (defaults to secrets['ops_db_name'])
    """
    import logging
    log = logging.getLogger(__name__)
    
    if database is None:
        database = secrets.get('ops_db_name', 'ops')
    
    hb = Heartbeat.create_with_pool(
        service=service,
        version=version,
        target=target,
        database=database,
        secrets=secrets
    )
    
    try:
        while not stop_event.is_set():
            await hb.write_async(alive=1)
            log.debug("Heartbeat written (alive=1) for service=%s", service)
            
            # Sleep with cancellation support
            try:
                await asyncio.wait_for(stop_event.wait(), timeout=interval_s)
                # If we get here, stop_event was set, break the loop
                break
            except asyncio.TimeoutError:
                # Normal - timeout means continue the loop
                pass
                
    except asyncio.CancelledError:
        log.info("Heartbeat loop cancelled, writing final heartbeat (alive=0) for service=%s", service)
        await hb.write_async(alive=0)
        raise



