"""Main agent daemon — heartbeat, probes, job execution."""
from __future__ import annotations

import asyncio
import hashlib
import logging
import platform
import signal
import time
from datetime import datetime
from typing import Optional

from mesh.config import MeshConfig
from mesh.agent.hardware_scanner import scan_hardware
from mesh.agent.health_reporter import collect_health
from mesh.agent.probe_runner import run_probes, extract_atlas_data
from mesh.models import (
    HardwareInventory, HealthSnapshot, NodeHeartbeat,
    NodeRegistration, NodeStatus
)
from mesh.net.client import MeshClient

logger = logging.getLogger(__name__)


class NodeAgent:
    """Agent daemon running on each mesh node."""

    def __init__(self, config: MeshConfig):
        self.config = config
        self.hostname = config.get_node_name()
        self.node_id = self._generate_node_id()
        self.hardware: Optional[HardwareInventory] = None
        self.status = NodeStatus.ONLINE
        self.registered = False
        self._running = False
        self._last_probe_time = 0.0
        self._controller_alive = False
        self._client: Optional[MeshClient] = None

    def _generate_node_id(self) -> str:
        """Stable node ID from configured name + arch."""
        raw = f"{self.hostname}-{platform.machine()}"
        return hashlib.sha256(raw.encode()).hexdigest()[:12]

    async def start(self):
        """Start the agent: scan hardware, register, begin loops."""
        self._running = True
        self._client = MeshClient(timeout=10.0)
        await self._client.__aenter__()
        logger.info("Node agent starting: id=%s hostname=%s", self.node_id, self.hostname)

        # Scan hardware
        self.hardware = scan_hardware(hostname_override=self.hostname, config=self.config)
        logger.info("Hardware: %d CPUs, %d GPUs, %d FPGAs, %dMB RAM",
                     self.hardware.cpu.cores,
                     len(self.hardware.gpus),
                     len(self.hardware.fpgas),
                     self.hardware.memory_total_mb)

        # Start background tasks.
        # In NAT mode, also start job polling (agent pulls jobs from controller
        # since the controller can't push to us behind NAT).
        tasks = [
            self._registration_loop(),
            self._heartbeat_loop(),
            self._probe_loop(),
            self._optimization_loop(),
        ]
        if self.config.node.nat_mode:
            tasks.append(self._job_poll_loop())
        await asyncio.gather(*tasks)

    async def stop(self):
        """Graceful shutdown: deregister from controller, close client."""
        self._running = False
        logger.info("Node agent stopping")

        # Deregister from controller
        if self.registered and self._client:
            try:
                url = f"{self.config.controller_url}/nodes/{self.node_id}?archive=true"
                status, _ = await self._client.delete(url)
                if status == 200:
                    logger.info("Deregistered from controller")
                else:
                    logger.debug("Deregistration returned %d", status)
            except Exception as e:
                logger.debug("Deregistration failed: %s", e)

        # Close shared client
        if self._client:
            await self._client.close()
            self._client = None

    async def _registration_loop(self):
        """Register with controller, retry on failure."""
        while self._running:
            if not self.registered:
                try:
                    await self._register()
                except Exception as e:
                    logger.warning("Registration failed: %s", e)
            await asyncio.sleep(30)

    async def _register(self):
        """Register this node with the controller.

        WAN support: if public_url is configured, use it as the agent_url so the
        controller can reach this agent over the internet. If not set, fall back to
        the local hostname + port (LAN behavior).

        NAT mode: if nat_mode is set, we tell the controller this node is push-only.
        The agent_url is still sent (for identification) but the controller should not
        attempt to connect to it for job dispatch or probes.
        """
        # Determine the URL the controller should use to reach us
        if self.config.node.public_url:
            agent_url = self.config.node.public_url
        else:
            agent_port = self.config.node.agent_port
            # Use the actual LAN IP rather than the display name,
            # which may not be DNS-resolvable.
            import socket
            try:
                host = self.config.node.host or socket.gethostname()
                if host == "0.0.0.0":
                    # Get the actual LAN IP by connecting to the controller
                    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                    try:
                        s.connect(("172.16.25.1", 80))
                        host = s.getsockname()[0]
                    except Exception:
                        host = socket.getfqdn()
                    finally:
                        s.close()
            except Exception:
                host = socket.getfqdn()
            agent_url = f"http://{host}:{agent_port}"

        reg = NodeRegistration(
            node_id=self.node_id,
            hostname=self.hostname,
            agent_url=agent_url,
            hardware=self.hardware,
            tags=self.config.node.tags,
            nat_mode=self.config.node.nat_mode,
        )

        url = f"{self.config.controller_url}/nodes/register"
        status, data = await self._client.post(url, json=reg.model_dump(mode="json"))
        if status == 200:
            self.registered = True
            self._controller_alive = True
            logger.info("Registered with controller")
        else:
            logger.warning("Registration rejected (%d): %s", status, str(data)[:200])

    async def _heartbeat_loop(self):
        """Send health snapshots to controller."""
        interval = self.config.controller.heartbeat_interval_s
        while self._running:
            await asyncio.sleep(interval)
            if not self.registered:
                continue
            try:
                health = collect_health()
                hb = NodeHeartbeat(
                    node_id=self.node_id,
                    health=health,
                )
                url = f"{self.config.controller_url}/nodes/{self.node_id}/heartbeat"
                status, _ = await self._client.post(url, json=hb.model_dump(mode="json"))
                self._controller_alive = status == 200
            except Exception as e:
                self._controller_alive = False
                logger.debug("Heartbeat failed: %s", e)

    async def _probe_loop(self):
        """Run probes periodically and sync to controller."""
        while self._running and not self.registered:
            await asyncio.sleep(5)

        probe_interval = self.config.controller.probe_interval_s
        while self._running:
            now = time.time()
            if now - self._last_probe_time >= probe_interval:
                logger.info("Starting scheduled probe run")
                try:
                    result = await asyncio.get_event_loop().run_in_executor(
                        None, lambda: run_probes(iterations=20)
                    )
                    self._last_probe_time = time.time()

                    if "error" not in result and result.get("db_path"):
                        await self._sync_atlas(result["db_path"])
                except Exception as e:
                    logger.error("Probe run failed: %s", e)

            await asyncio.sleep(60)

    async def _sync_atlas(self, db_path: str):
        """Send probe results to controller."""
        data = extract_atlas_data(db_path)
        if not data["data_points"]:
            return

        payload = {
            "node_id": self.node_id,
            "atlas_db_name": f"{self.hostname}_atlas.db",
            "data_points": data["data_points"],
            "invariants": data["invariants"],
        }

        url = f"{self.config.controller_url}/atlas/sync"
        try:
            status, _ = await self._client.post(url, json=payload)
            if status == 200:
                logger.info("Atlas synced: %d points, %d invariants",
                             len(data["data_points"]), len(data["invariants"]))
            else:
                logger.warning("Atlas sync failed (%d)", status)
        except Exception as e:
            logger.warning("Atlas sync error: %s", e)

    async def _optimization_loop(self):
        """System-wide tuning + event-driven process optimization."""
        while self._running and not self.registered:
            await asyncio.sleep(5)

        from mesh.agent.system_optimizer import SystemOptimizer, ProcessMonitor

        optimizer = SystemOptimizer(
            node_name=self.hostname,
            controller_url=self.config.controller_url,
        )

        proc_monitor = ProcessMonitor(optimizer)
        proc_monitor.start()
        logger.info("Process monitor started — new processes will be optimized automatically")

        syswide_interval = self.config.controller.probe_interval_s

        try:
            while self._running:
                try:
                    results = await asyncio.get_event_loop().run_in_executor(
                        None, lambda: optimizer.apply_all(dry_run=False)
                    )
                    ok_count = sum(1 for v in results.values() if "OK" in str(v))
                    logger.info("Full optimization applied: %d/%d settings OK", ok_count, len(results))
                except Exception as e:
                    logger.warning("System optimization error: %s", e)

                await asyncio.sleep(syswide_interval)
        finally:
            proc_monitor.stop()

    async def _job_poll_loop(self):
        """NAT mode: poll the controller for pending job assignments.

        Since the controller cannot push jobs to us (we're behind NAT), we
        periodically ask the controller if there are any jobs assigned to us.
        The controller queues jobs for NAT-mode nodes instead of dispatching them.
        """
        poll_interval = self.config.node.nat_poll_interval_s
        while self._running and not self.registered:
            await asyncio.sleep(5)

        logger.info("NAT mode: job polling started (interval=%.0fs)", poll_interval)
        while self._running:
            await asyncio.sleep(poll_interval)
            if not self.registered or not self._controller_alive:
                continue
            try:
                url = f"{self.config.controller_url}/nodes/{self.node_id}/pending-jobs"
                status, data = await self._client.get(url)
                if status == 200 and data:
                    jobs = data if isinstance(data, list) else []
                    for job_data in jobs:
                        logger.info("NAT poll: received job %s", job_data.get("job_id", "?"))
                        # Execute via the same path as the agent API's submit endpoint
                        from mesh.agent.job_executor import execute_job
                        from mesh.models import JobInfo as _JobInfo, JobStatus as _JS
                        job = _JobInfo(**job_data)
                        job.status = _JS.RUNNING
                        result = await execute_job(job)
                        # Report result back to controller
                        result_url = f"{self.config.controller_url}/jobs/{job.job_id}/result"
                        await self._client.post(result_url, json={
                            "job_id": job.job_id,
                            "status": result.status.value,
                            "result": result.result,
                            "error": result.error,
                            "duration_s": result.duration_s,
                        })
            except Exception as e:
                logger.debug("NAT job poll failed: %s", e)

    @property
    def is_controller_alive(self) -> bool:
        return self._controller_alive
