import http.server
import os
import re
import socketserver
import subprocess
import threading
from time import sleep, time

PORT = int(os.environ.get("JETSON_EXPORTER_PORT", "9100"))
NODE_NAME = os.environ.get("NODE_NAME") or os.uname().nodename
BASE_METRICS = {
    "gr3d_freq_percent": 0.0,
    "gpu_temp_c": 0.0,
    "cpu_temp_c": 0.0,
    "ram_used_mb": 0.0,
    "ram_total_mb": 0.0,
    "power_5v_in_mw": 0.0,
    "log_line_len": 0.0,
    "last_scrape_ts": 0.0,
    "last_sample_timestamp_seconds": 0.0,
    "gr3d_active_seconds_total": 0.0,
    "tegrastats_samples_total": 0.0,
}
COUNTER_METRICS = {"gr3d_active_seconds_total", "tegrastats_samples_total"}


def parse_line(line: str) -> dict:
    line = line.strip()
    updates = {}
    m = re.search(r"GR3D_FREQ\s+(\d+)%", line)
    if m:
        updates["gr3d_freq_percent"] = float(m.group(1))
    m = re.search(r"GPU@(\d+(?:\.\d+)?)C", line)
    if m:
        updates["gpu_temp_c"] = float(m.group(1))
    m = re.search(r"CPU@(\d+(?:\.\d+)?)C", line)
    if m:
        updates["cpu_temp_c"] = float(m.group(1))
    m = re.search(r"RAM\s+(\d+)/(\d+)MB", line)
    if m:
        updates["ram_used_mb"] = float(m.group(1))
        updates["ram_total_mb"] = float(m.group(2))
    m = re.search(r"(?:POM_5V_IN|VDD_IN)\s+(\d+)(?:mW)?/(\d+)(?:mW)?", line)
    if m:
        updates["power_5v_in_mw"] = float(m.group(1))
    return updates


class MetricStore:
    """Retain samples and integrate short GPU bursts between Prometheus scrapes."""

    def __init__(self):
        self.metrics = BASE_METRICS.copy()
        self.last_sample_at = None
        self.lock = threading.Lock()

    def record(self, line: str, now: float | None = None) -> None:
        """Record one tegrastats line and add utilization-weighted GPU seconds."""
        sampled_at = time() if now is None else now
        updates = parse_line(line)
        if not updates:
            return
        with self.lock:
            if self.last_sample_at is not None:
                elapsed = min(max(sampled_at - self.last_sample_at, 0.0), 2.0)
                previous = self.metrics["gr3d_freq_percent"]
                self.metrics["gr3d_active_seconds_total"] += elapsed * previous / 100.0
            self.metrics.update(updates)
            self.metrics["log_line_len"] = float(len(line))
            self.metrics["last_sample_timestamp_seconds"] = sampled_at
            self.metrics["tegrastats_samples_total"] += 1
            self.last_sample_at = sampled_at

    def snapshot(self, now: float | None = None) -> dict:
        """Return an atomic copy suitable for Prometheus exposition."""
        scraped_at = time() if now is None else now
        with self.lock:
            metrics = self.metrics.copy()
        metrics["last_scrape_ts"] = scraped_at
        return metrics


METRICS = MetricStore()


def sample_forever() -> None:
    """Keep tegrastats running so sub-scrape inference bursts are retained."""
    while True:
        try:
            proc = subprocess.Popen(
                ["/host/usr/bin/tegrastats", "--interval", "250"],
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
            )
            if proc.stdout is not None:
                for line in proc.stdout:
                    METRICS.record(line)
            proc.wait(timeout=1)
        except (OSError, subprocess.SubprocessError):
            pass
        sleep(1)


def render_metrics(metrics: dict) -> str:
    """Render the current metric snapshot in Prometheus text format."""
    out = []
    label = f'{{node="{NODE_NAME}"}}'
    for key, value in metrics.items():
        metric_type = "counter" if key in COUNTER_METRICS else "gauge"
        out.append(f"# TYPE jetson_{key} {metric_type}")
        out.append(f"jetson_{key}{label} {value}")
    return "\n".join(out) + "\n"


class Handler(http.server.BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path != "/metrics":
            self.send_response(404)
            self.end_headers()
            return
        body = render_metrics(METRICS.snapshot())
        self.send_response(200)
        self.send_header("Content-Type", "text/plain; version=0.0.4")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body.encode("utf-8"))

    def log_message(self, fmt, *args):
        return


if __name__ == "__main__":
    threading.Thread(target=sample_forever, daemon=True).start()
    with socketserver.TCPServer(("", PORT), Handler) as httpd:
        httpd.serve_forever()