"""Background telemetry collector — gathers metrics and sends heartbeats."""
from __future__ import annotations

import asyncio
import json
import logging
import time
from datetime import datetime, timedelta
from pathlib import Path
from typing import Any, Dict, List, Optional

from .client import LicenseClient

logger = logging.getLogger("mesh.licensing.telemetry")

_TELEMETRY_LOG_PATH = "data/telemetry_log.json"
_PRUNE_DAYS = 30
_DEFAULT_INTERVAL_S = 21600  # 6 hours


class TelemetryCollector:
    """Collects aggregate metrics and sends them to the license server.

    Runs as an asyncio background task.  No PII is collected — only
    aggregate counts and percentages.
    """

    def __init__(
        self,
        coordinator,
        license_client: LicenseClient,
        log_path: str = _TELEMETRY_LOG_PATH,
        interval_s: float = _DEFAULT_INTERVAL_S,
        mesh_version: str = "1.0.0",
    ):
        self.coordinator = coordinator
        self.license_client = license_client
        self.log_path = Path(log_path)
        self.interval_s = interval_s
        self.mesh_version = mesh_version
        self._start_time = time.monotonic()
        self._task: Optional[asyncio.Task] = None

    def start(self) -> asyncio.Task:
        """Start the background collection loop."""
        self._task = asyncio.create_task(self._loop())
        logger.info("Telemetry collector started (interval=%.0fs)", self.interval_s)
        return self._task

    async def stop(self):
        if self._task and not self._task.done():
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        logger.info("Telemetry collector stopped")

    async def _loop(self):
        while True:
            try:
                await asyncio.sleep(self.interval_s)
                await self.collect_and_send()
            except asyncio.CancelledError:
                break
            except Exception as exc:
                logger.error("Telemetry collection error: %s", exc)

    async def collect_and_send(self):
        """Gather metrics and send heartbeat."""
        snapshot = self._gather_metrics()

        # Send heartbeat
        ok = self.license_client.heartbeat(
            node_count=snapshot["node_count"],
            hardware_summary=snapshot["hardware_summary"],
            version=self.mesh_version,
            uptime_hours=snapshot["uptime_hours"],
            avg_cpu_pct=snapshot["avg_cpu_pct"],
            avg_memory_pct=snapshot["avg_memory_pct"],
            total_jobs_run=snapshot["total_jobs_run"],
            total_retrains=snapshot["total_retrains"],
        )

        if ok:
            logger.debug("Telemetry heartbeat sent successfully")
        else:
            logger.debug("Telemetry heartbeat failed (server unreachable)")

        # Log locally
        self._log_locally(snapshot)

    def _gather_metrics(self) -> Dict[str, Any]:
        """Gather aggregate metrics from the coordinator."""
        nodes = []
        try:
            nodes = self.coordinator.get_online_nodes()
        except Exception:
            pass

        hardware_summary: Dict[str, Any] = {}
        avg_cpu = 0.0
        avg_mem = 0.0

        for node in nodes:
            hw = getattr(node, "hardware", None)
            if hw:
                for gpu in getattr(hw, "gpus", []):
                    key = f"{gpu.vendor}_{gpu.name}".replace(" ", "_")
                    hardware_summary[key] = hardware_summary.get(key, 0) + 1

            health = getattr(node, "last_health", None)
            if health:
                avg_cpu += getattr(health, "cpu_pct", 0.0)
                avg_mem += getattr(health, "memory_used_pct", 0.0)

        n = max(len(nodes), 1)
        uptime_hours = (time.monotonic() - self._start_time) / 3600.0

        # Job/retrain counts from coordinator if available
        total_jobs = 0
        total_retrains = 0
        if hasattr(self.coordinator, "total_jobs_run"):
            total_jobs = self.coordinator.total_jobs_run
        if hasattr(self.coordinator, "total_retrains"):
            total_retrains = self.coordinator.total_retrains

        return {
            "timestamp": datetime.utcnow().isoformat(),
            "node_count": len(nodes),
            "hardware_summary": hardware_summary,
            "uptime_hours": round(uptime_hours, 2),
            "avg_cpu_pct": round(avg_cpu / n, 1),
            "avg_memory_pct": round(avg_mem / n, 1),
            "total_jobs_run": total_jobs,
            "total_retrains": total_retrains,
        }

    def _log_locally(self, snapshot: Dict[str, Any]):
        """Append snapshot to local telemetry log, pruning entries older than 30 days."""
        try:
            self.log_path.parent.mkdir(parents=True, exist_ok=True)
            entries: List[Dict[str, Any]] = []
            if self.log_path.exists():
                try:
                    entries = json.loads(self.log_path.read_text())
                except (json.JSONDecodeError, ValueError):
                    entries = []

            entries.append(snapshot)

            # Prune old entries
            cutoff = (datetime.utcnow() - timedelta(days=_PRUNE_DAYS)).isoformat()
            entries = [e for e in entries if e.get("timestamp", "") >= cutoff]

            self.log_path.write_text(json.dumps(entries, indent=2))
        except Exception as exc:
            logger.debug("Failed to write telemetry log: %s", exc)
