#!/usr/bin/env python3
"""
GPU temperature monitor.
SSHes into VM 109, checks nvidia-smi, sends email if critical.
"""
import smtplib
import subprocess
import time
import logging
import os
from dataclasses import dataclass
from email.mime.text import MIMEText
from dotenv import load_dotenv

load_dotenv()

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    handlers=[
        logging.FileHandler("/var/log/gpu-monitor.log"),
        logging.StreamHandler(),
    ]
)
log = logging.getLogger(__name__)

# ── config ────────────────────────────────────────────────────────────────────

GPU_HOST     = os.getenv("GPU_HOST",     "192.168.20.47")
GPU_USER     = os.getenv("GPU_USER",     "user")
GPU_PASS     = os.getenv("GPU_PASS",     "")

SMTP_HOST    = os.getenv("SMTP_HOST",    "smtp.gmail.com")
SMTP_PORT    = int(os.getenv("SMTP_PORT", "587"))
SMTP_USER    = os.getenv("SMTP_USER",    "")
SMTP_PASS    = os.getenv("SMTP_PASS",    "")
ALERT_TO     = os.getenv("ALERT_TO",     "")

WARN_TEMP    = int(os.getenv("WARN_TEMP",  "80"))
CRIT_TEMP    = int(os.getenv("CRIT_TEMP",  "88"))
CHECK_EVERY  = int(os.getenv("CHECK_EVERY", "60"))   # seconds

# ── GPU query ─────────────────────────────────────────────────────────────────

@dataclass
class GpuStats:
    temp: int
    fan: str
    power: str
    util: str
    mem_used: str
    mem_total: str

def get_gpu_stats() -> GpuStats | None:
    cmd = [
        "sshpass", "-p", GPU_PASS,
        "ssh", "-o", "StrictHostKeyChecking=no",
        "-o", "ConnectTimeout=10",
        f"{GPU_USER}@{GPU_HOST}",
        "nvidia-smi --query-gpu=temperature.gpu,fan.speed,power.draw,"
        "utilization.gpu,memory.used,memory.total --format=csv,noheader"
    ]
    try:
        out = subprocess.check_output(cmd, stderr=subprocess.DEVNULL, timeout=15)
        parts = [p.strip() for p in out.decode().strip().split(",")]
        return GpuStats(
            temp=int(parts[0]),
            fan=parts[1],
            power=parts[2],
            util=parts[3],
            mem_used=parts[4],
            mem_total=parts[5],
        )
    except Exception as e:
        log.error(f"Failed to query GPU: {e}")
        return None

# ── email ─────────────────────────────────────────────────────────────────────

def send_alert(subject: str, body: str):
    if not SMTP_USER or not ALERT_TO:
        log.warning("Email not configured, skipping alert")
        return
    try:
        msg = MIMEText(body, "plain", "utf-8")
        msg["Subject"] = subject
        msg["From"]    = SMTP_USER
        msg["To"]      = ALERT_TO
        with smtplib.SMTP(SMTP_HOST, SMTP_PORT, timeout=15) as s:
            s.starttls()
            s.login(SMTP_USER, SMTP_PASS)
            s.send_message(msg)
        log.info(f"Alert sent: {subject}")
    except Exception as e:
        log.error(f"Failed to send email: {e}")

# ── main loop ─────────────────────────────────────────────────────────────────

def main():
    log.info(f"GPU monitor started — checking {GPU_HOST} every {CHECK_EVERY}s")
    log.info(f"Thresholds: warn={WARN_TEMP}°C  crit={CRIT_TEMP}°C")

    last_alert_temp = 0   # avoid spamming — only re-alert if temp changes by 3°C

    while True:
        stats = get_gpu_stats()

        if stats is None:
            time.sleep(CHECK_EVERY)
            continue

        log.info(
            f"GPU {stats.temp}°C | fan {stats.fan} | {stats.power} | "
            f"util {stats.util} | mem {stats.mem_used}/{stats.mem_total}"
        )

        if stats.temp >= CRIT_TEMP:
            if abs(stats.temp - last_alert_temp) >= 3:
                send_alert(
                    f"🔥 КРИТИЧНО: GPU {stats.temp}°C на {GPU_HOST}",
                    f"GPU перегрев!\n\n"
                    f"Температура : {stats.temp}°C  (критично от {CRIT_TEMP}°C)\n"
                    f"Вентилятор  : {stats.fan}\n"
                    f"Потребление : {stats.power}\n"
                    f"Загрузка GPU: {stats.util}\n"
                    f"Память      : {stats.mem_used} / {stats.mem_total}\n\n"
                    f"Проверь охлаждение немедленно."
                )
                last_alert_temp = stats.temp

        elif stats.temp >= WARN_TEMP:
            if abs(stats.temp - last_alert_temp) >= 3:
                send_alert(
                    f"⚠️ ВНИМАНИЕ: GPU {stats.temp}°C на {GPU_HOST}",
                    f"GPU горячий.\n\n"
                    f"Температура : {stats.temp}°C  (предупреждение от {WARN_TEMP}°C)\n"
                    f"Вентилятор  : {stats.fan}\n"
                    f"Потребление : {stats.power}\n"
                    f"Загрузка GPU: {stats.util}\n"
                    f"Память      : {stats.mem_used} / {stats.mem_total}\n"
                )
                last_alert_temp = stats.temp
        else:
            # reset cooldown when temp is back to normal
            if last_alert_temp > 0:
                log.info(f"GPU temp back to normal ({stats.temp}°C)")
                last_alert_temp = 0

        time.sleep(CHECK_EVERY)

if __name__ == "__main__":
    main()