#!/usr/bin/env python3
"""ePulz.io LAN Agent v2 - bulletproof edition.

Configuration via environment variables:
  EPULZIO_API           https://epulz.io (default)
  EPULZIO_AGENT_TOKEN   epulzio_xxxxxxxxxxxx (required)
  EPULZIO_POLL_SECONDS  jobs poll interval (default 30)
  EPULZIO_HEALTH_PORT   local health endpoint port (default 9090, 0 to disable)
  EPULZIO_BUFFER_DIR    failed-results buffer dir (default /var/lib/epulzio-agent)
  EPULZIO_MAX_PARALLEL  max parallel checks (default 10)
  EPULZIO_HTTP_PROXY    optional outbound HTTP proxy
  EPULZIO_LOG_LEVEL     debug, info, warning, error (default info)

Install (systemd):
  sudo bash <(curl -fsSL https://epulz.io/install-agent.sh) <token>

Stdlib only - no pip dependencies. Python 3.8+.
"""
from __future__ import annotations

import json
import logging
import os
import platform
import queue
import signal
import socket
import subprocess
import sys
import threading
import time
import urllib.error
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from typing import Optional

VERSION = "2.0.0"

# ─────────── Configuration ───────────

API = os.environ.get("EPULZIO_API", "https://epulz.io").rstrip("/")
TOKEN = os.environ.get("EPULZIO_AGENT_TOKEN", "").strip()
POLL_SECONDS = max(5, int(os.environ.get("EPULZIO_POLL_SECONDS", "30")))
HEALTH_PORT = int(os.environ.get("EPULZIO_HEALTH_PORT", "9090"))
BUFFER_DIR = Path(os.environ.get("EPULZIO_BUFFER_DIR", "/var/lib/epulzio-agent"))
MAX_PARALLEL = max(1, int(os.environ.get("EPULZIO_MAX_PARALLEL", "10")))
HTTP_PROXY = os.environ.get("EPULZIO_HTTP_PROXY", "").strip()
LOG_LEVEL = os.environ.get("EPULZIO_LOG_LEVEL", "info").upper()

# Watchdog: hard timeout for any single poll/process iteration
WATCHDOG_TIMEOUT_S = max(120, POLL_SECONDS * 4)

# ─────────── Logging ───────────

logging.basicConfig(
    level=getattr(logging, LOG_LEVEL, logging.INFO),
    format="%(asctime)s %(levelname)s %(message)s",
    datefmt="%Y-%m-%dT%H:%M:%S%z",
    stream=sys.stdout,
)
log = logging.getLogger("epulzio-agent")

# ─────────── Token check ───────────

if not TOKEN:
    log.error("Missing EPULZIO_AGENT_TOKEN env variable. Get a token at %s/dashboard/agents", API)
    sys.exit(2)

if not TOKEN.startswith("epulzio_"):
    log.error("Invalid token format: must start with 'epulzio_'. Generate a new token at %s/dashboard/agents", API)
    sys.exit(1)

# ─────────── Internal state (thread-safe) ───────────

class AgentState:
    def __init__(self):
        self.lock = threading.Lock()
        self.started_at = time.time()
        self.last_poll_ok_at: Optional[float] = None
        self.last_poll_error: Optional[str] = None
        self.last_error_at: Optional[float] = None
        self.jobs_completed = 0
        self.jobs_24h = []  # list of timestamps
        self.results_buffered = 0
        self.heartbeat_ok = False
        self.last_latest_version: Optional[str] = None

    def record_poll_ok(self):
        with self.lock:
            self.last_poll_ok_at = time.time()
            self.last_poll_error = None

    def record_poll_error(self, err: str):
        with self.lock:
            self.last_poll_error = err
            self.last_error_at = time.time()

    def record_jobs(self, n: int):
        with self.lock:
            self.jobs_completed += n
            now = time.time()
            self.jobs_24h.append((now, n))
            self.jobs_24h = [(t, c) for t, c in self.jobs_24h if now - t < 86400]

    def jobs_last_24h(self) -> int:
        with self.lock:
            now = time.time()
            return sum(c for t, c in self.jobs_24h if now - t < 86400)

    def snapshot(self) -> dict:
        with self.lock:
            now = time.time()
            return {
                "version": VERSION,
                "uptime_s": int(now - self.started_at),
                "last_poll_ok_at": self.last_poll_ok_at,
                "last_poll_ago_s": int(now - self.last_poll_ok_at) if self.last_poll_ok_at else None,
                "last_poll_error": self.last_poll_error,
                "last_error_at": self.last_error_at,
                "jobs_completed_total": self.jobs_completed,
                "jobs_last_24h": sum(c for t, c in self.jobs_24h if now - t < 86400),
                "results_buffered": self.results_buffered,
                "heartbeat_ok": self.heartbeat_ok,
                "latest_version_available": self.last_latest_version,
            }


STATE = AgentState()
STOP = threading.Event()

# ─────────── Result buffering (resilient to API outages) ───────────

class ResultBuffer:
    """Persists failed result uploads to disk so they can be retried later.
    Each result is one JSON line in a .jsonl file. Max 10 MB total."""
    MAX_BYTES = 10 * 1024 * 1024

    def __init__(self, path: Path):
        self.path = path
        self.lock = threading.Lock()
        self.path.parent.mkdir(parents=True, exist_ok=True)
        if not self.path.exists():
            self.path.touch()

    def append(self, results: list[dict]) -> None:
        with self.lock:
            try:
                # Truncate if oversized (drop oldest entries)
                if self.path.stat().st_size > self.MAX_BYTES:
                    lines = self.path.read_text(encoding="utf-8").splitlines()
                    kept = lines[-1000:]  # keep last 1000 entries only
                    self.path.write_text("\n".join(kept) + "\n", encoding="utf-8")
                with self.path.open("a", encoding="utf-8") as f:
                    for r in results:
                        r["buffered_at"] = int(time.time())
                        f.write(json.dumps(r, ensure_ascii=False) + "\n")
                STATE.results_buffered = sum(1 for _ in self.path.read_text(encoding="utf-8").splitlines())
            except Exception as e:
                log.warning("Failed to buffer results: %s", e)

    def drain(self) -> list[dict]:
        """Return all buffered results and clear the buffer."""
        with self.lock:
            try:
                if self.path.stat().st_size == 0:
                    return []
                lines = [l for l in self.path.read_text(encoding="utf-8").splitlines() if l.strip()]
                self.path.write_text("", encoding="utf-8")
                STATE.results_buffered = 0
                return [json.loads(l) for l in lines]
            except Exception as e:
                log.warning("Failed to drain buffer: %s", e)
                return []


BUFFER = ResultBuffer(BUFFER_DIR / "results.jsonl")

# ─────────── HTTP client (urllib with retry + proxy) ───────────

def _setup_proxy():
    if HTTP_PROXY:
        proxy = {"http": HTTP_PROXY, "https": HTTP_PROXY}
        handler = urllib.request.ProxyHandler(proxy)
        opener = urllib.request.build_opener(handler)
        urllib.request.install_opener(opener)
        log.info("HTTP proxy configured: %s", HTTP_PROXY)


_setup_proxy()


def api_call(method: str, path: str, body: Optional[dict] = None,
             timeout: int = 20, retries: int = 3) -> dict:
    """API call with exponential backoff retry.
    Retries on network errors and 5xx. Does NOT retry on 4xx (likely permanent error).
    """
    url = f"{API}{path}"
    data = json.dumps(body).encode("utf-8") if body is not None else None
    headers = {"X-Agent-Token": TOKEN, "User-Agent": f"ePulzio-Agent/{VERSION}"}
    if data is not None:
        headers["Content-Type"] = "application/json"
    last_err: Optional[Exception] = None
    for attempt in range(retries):
        try:
            req = urllib.request.Request(url, data=data, headers=headers, method=method)
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                raw = resp.read().decode("utf-8")
                return json.loads(raw) if raw else {}
        except urllib.error.HTTPError as e:
            if 400 <= e.code < 500:
                # Client error - don't retry (probably bad token or bad request)
                body_text = ""
                try:
                    body_text = e.read().decode("utf-8", errors="replace")[:200]
                except Exception:
                    pass
                log.error("API %s %s: HTTP %s %s", method, path, e.code, body_text)
                raise
            last_err = e
            log.warning("API %s %s: HTTP %s (attempt %d/%d)", method, path, e.code, attempt + 1, retries)
        except (urllib.error.URLError, socket.timeout, OSError) as e:
            last_err = e
            log.warning("API %s %s: %s (attempt %d/%d)", method, path, e, attempt + 1, retries)
        # Exponential backoff: 1s, 2s, 4s
        if attempt < retries - 1:
            time.sleep(2 ** attempt)
    # All retries exhausted
    raise last_err if last_err else RuntimeError("all retries failed")


# ─────────── Check implementations ───────────

def _ping_args(target: str, count: int, timeout_s: int) -> list[str]:
    is_win = platform.system().lower().startswith("win")
    count_flag = "-n" if is_win else "-c"
    timeout_flag = "-w" if is_win else "-W"
    timeout_val = str(timeout_s * 1000) if is_win else str(timeout_s)
    return ["ping", count_flag, str(count), timeout_flag, timeout_val, target]


def do_ping(target: str, timeout_s: int = 5) -> tuple[str, Optional[float], Optional[str]]:
    """ICMP ping. Returns (status, response_ms, error_msg)."""
    args = _ping_args(target, 3, timeout_s)
    start = time.time()
    try:
        proc = subprocess.run(args, capture_output=True, text=True,
                              timeout=timeout_s * 4 + 2)
        elapsed_ms = (time.time() - start) * 1000
        if proc.returncode == 0:
            for line in proc.stdout.splitlines():
                if "avg" in line.lower() or "average" in line.lower():
                    for tok in line.replace("=", " ").split():
                        if "/" in tok and tok.count("/") >= 2:
                            parts = tok.split("/")
                            try:
                                return "up", round(float(parts[1]), 3), None
                            except (ValueError, IndexError):
                                pass
            return "up", round(elapsed_ms / 3, 3), None
        err = (proc.stderr or proc.stdout).strip().splitlines()
        msg = err[-1] if err else f"ping exit {proc.returncode}"
        return "down", None, msg[:200]
    except subprocess.TimeoutExpired:
        return "down", None, f"ping timeout (>{timeout_s*4}s)"
    except FileNotFoundError:
        return "down", None, "'ping' binary not found (install iputils-ping)"
    except Exception as e:
        return "down", None, f"{type(e).__name__}: {e}"


def do_tcp(target: str, timeout_s: int = 5) -> tuple[str, Optional[float], Optional[str]]:
    """TCP connect to host:port."""
    if ":" not in target:
        return "down", None, "TCP target must be in host:port format"
    host, _, port_str = target.rpartition(":")
    try:
        port = int(port_str)
    except ValueError:
        return "down", None, f"invalid port: {port_str}"
    start = time.time()
    try:
        with socket.create_connection((host, port), timeout=timeout_s):
            ms = round((time.time() - start) * 1000, 3)
            return "up", ms, None
    except socket.timeout:
        return "down", None, f"TCP timeout after {timeout_s}s"
    except (ConnectionRefusedError, OSError) as e:
        return "down", None, f"TCP {host}:{port}: {e}"
    except Exception as e:
        return "down", None, f"{type(e).__name__}: {e}"


def do_http(target: str, timeout_s: int = 10) -> tuple[str, Optional[float], Optional[str]]:
    """HTTP GET. Accepts URL or host."""
    url = target if "://" in target else f"http://{target}"
    start = time.time()
    try:
        req = urllib.request.Request(url, headers={"User-Agent": f"ePulzio-Agent/{VERSION}"})
        with urllib.request.urlopen(req, timeout=timeout_s) as resp:
            ms = round((time.time() - start) * 1000, 3)
            code = resp.status
            if 200 <= code < 400:
                return "up", ms, None
            return "down", ms, f"HTTP {code}"
    except urllib.error.HTTPError as e:
        ms = round((time.time() - start) * 1000, 3)
        return "down", ms, f"HTTP {e.code}"
    except urllib.error.URLError as e:
        return "down", None, f"URL error: {e.reason}"
    except socket.timeout:
        return "down", None, f"HTTP timeout after {timeout_s}s"
    except Exception as e:
        return "down", None, f"{type(e).__name__}: {e}"


def run_check(job: dict) -> dict:
    mid = job["monitor_id"]
    typ = job.get("type", "")
    target = job.get("target", "")
    timeout = max(2, int(job.get("timeout_seconds", 5)))
    try:
        if typ == "lan_ping":
            status, ms, err = do_ping(target, timeout_s=min(timeout, 10))
        elif typ == "lan_tcp":
            status, ms, err = do_tcp(target, timeout_s=timeout)
        elif typ == "lan_http":
            status, ms, err = do_http(target, timeout_s=max(timeout, 10))
        else:
            status, ms, err = "down", None, f"agent does not support type '{typ}'"
    except Exception as e:
        status, ms, err = "down", None, f"check crashed: {type(e).__name__}: {e}"
    return {"monitor_id": mid, "status": status, "response_ms": ms, "error": err}


# ─────────── Health endpoint ───────────

class HealthHandler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path in ("/health", "/"):
            snap = STATE.snapshot()
            body = json.dumps(snap, indent=2).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
        else:
            self.send_response(404)
            self.end_headers()
            self.wfile.write(b"not found")

    def log_message(self, format, *args):
        # Suppress request logs to stdout (we have our own logger)
        pass


def start_health_server():
    if HEALTH_PORT <= 0:
        return
    try:
        server = HTTPServer(("127.0.0.1", HEALTH_PORT), HealthHandler)
        t = threading.Thread(target=server.serve_forever, name="health-srv", daemon=True)
        t.start()
        log.info("Health endpoint on http://127.0.0.1:%d/health", HEALTH_PORT)
    except OSError as e:
        log.warning("Cannot bind health endpoint on :%d (%s) - skipping", HEALTH_PORT, e)


# ─────────── Watchdog ───────────

def watchdog():
    """Kill process if main loop hasn't progressed for WATCHDOG_TIMEOUT_S.
    systemd will restart us (Restart=always). Prevents infinite hangs."""
    while not STOP.is_set():
        time.sleep(30)
        last = STATE.last_poll_ok_at
        if last is None:
            continue
        if time.time() - last > WATCHDOG_TIMEOUT_S:
            log.error("Watchdog: no successful poll for %ds - killing process", int(time.time() - last))
            os._exit(3)  # noqa: SystemExit equivalent, bypass cleanup


# ─────────── Main poll loop ───────────

def _send_heartbeat():
    try:
        r = api_call("POST", "/api/agent/heartbeat", {
            "agent_version": VERSION,
            "hostname": socket.gethostname(),
            "platform": f"{platform.system()} {platform.release()} ({platform.machine()})",
        })
        STATE.heartbeat_ok = True
        log.info("Heartbeat ok (agent_id=%s)", r.get("agent_id"))
    except Exception as e:
        STATE.heartbeat_ok = False
        log.warning("Heartbeat failed: %s", e)


def _poll_iteration(executor: ThreadPoolExecutor) -> None:
    """Single iteration: get jobs, run them in parallel, send results."""
    data = api_call("GET", "/api/agent/jobs")
    jobs = data.get("jobs", []) or []
    latest = data.get("latest_agent_version")
    if latest and latest != VERSION and latest != STATE.last_latest_version:
        STATE.last_latest_version = latest
        log.warning("Newer agent version available: %s (current %s)", latest, VERSION)
        log.warning("  Update: sudo curl -fsSL %s/static/agent/epulzio-agent.py -o /opt/epulzio-agent/agent.py && sudo systemctl restart epulzio-agent", API)

    # Drain buffered results from previous failures, prepend them to new results
    buffered = BUFFER.drain()
    if buffered:
        log.info("Replaying %d buffered results from previous outage", len(buffered))

    results: list[dict] = list(buffered)
    if jobs:
        log.info("Got %d jobs to check", len(jobs))
        futures = [executor.submit(run_check, j) for j in jobs]
        for fut in as_completed(futures, timeout=WATCHDOG_TIMEOUT_S - 30):
            try:
                results.append(fut.result())
            except Exception as e:
                log.warning("Check task crashed: %s", e)

    if results:
        try:
            api_call("POST", "/api/agent/results",
                     {"results": [r for r in results if "monitor_id" in r],
                      "agent_version": VERSION})
            up = sum(1 for r in results if r.get("status") == "up")
            STATE.record_jobs(len(results))
            log.info("Sent %d results (up=%d, down=%d, replayed=%d)",
                     len(results), up, len(results) - up, len(buffered))
        except Exception as e:
            log.error("Upload failed - buffering %d results for retry: %s", len(results), e)
            BUFFER.append(results)
            return  # Don't mark poll as ok if upload failed

    STATE.record_poll_ok()


def main():
    log.info("ePulz.io Agent v%s starting → %s", VERSION, API)
    log.info("Config: poll=%ds, parallel=%d, watchdog=%ds, buffer=%s",
             POLL_SECONDS, MAX_PARALLEL, WATCHDOG_TIMEOUT_S, BUFFER_DIR)

    # Signal handlers - clean shutdown on SIGTERM (systemd stop) and SIGINT
    def _handle_signal(signum, frame):
        log.info("Received signal %d - shutting down", signum)
        STOP.set()
    for sig in (signal.SIGTERM, signal.SIGINT):
        try:
            signal.signal(sig, _handle_signal)
        except Exception:
            pass

    start_health_server()

    # Watchdog thread
    threading.Thread(target=watchdog, name="watchdog", daemon=True).start()

    # Initial heartbeat
    _send_heartbeat()

    # Main loop with parallel check execution
    with ThreadPoolExecutor(max_workers=MAX_PARALLEL, thread_name_prefix="check") as executor:
        next_heartbeat_at = time.time() + 300  # heartbeat every 5 min
        while not STOP.is_set():
            iteration_start = time.time()
            try:
                _poll_iteration(executor)
            except urllib.error.HTTPError as e:
                STATE.record_poll_error(f"HTTP {e.code}")
                # 4xx already logged in api_call; just sleep and continue
            except Exception as e:
                STATE.record_poll_error(f"{type(e).__name__}: {e}")
                log.error("Poll iteration failed: %s", e)

            if time.time() >= next_heartbeat_at:
                _send_heartbeat()
                next_heartbeat_at = time.time() + 300

            # Sleep until next poll (interruptible by STOP signal)
            slept = time.time() - iteration_start
            remaining = max(0, POLL_SECONDS - slept)
            STOP.wait(timeout=remaining)

    log.info("Agent stopped cleanly")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        log.info("Interrupted")
    except Exception as e:
        log.exception("Fatal: %s", e)
        sys.exit(1)