#!/opt/cpu-freq-guard/venv/bin/python
"""CPU 动态调频与带迟滞的温度守护，仅使用标准库。"""
import argparse
import fcntl
import json
import logging
import signal
import threading
import time
from pathlib import Path

LOG = logging.getLogger("cpu-freq-guard")


def read(path):
    return path.read_text().strip()


def write(path, value, verify=True, ceiling=False):
    if read(path) != str(value):
        path.write_text(str(value) + "\n")
    if not verify:
        return
    # 厂商内核异步更新 policy，等待有限时间再判定失败。
    deadline = time.monotonic() + 1.0
    while not (int(read(path)) <= int(value) if ceiling else read(path) == str(value)):
        if time.monotonic() >= deadline:
            raise RuntimeError(f"写入校验失败：{path}，请求 {value}，实际 {read(path)}")
        time.sleep(0.02)


def atomic_json(path, data):
    temp = path.with_suffix(".tmp")
    temp.write_text(json.dumps(data, indent=2, ensure_ascii=False) + "\n")
    temp.chmod(0o600)
    temp.replace(path)


def load_config(path):
    data = json.loads(path.read_text())
    required = {"governor", "hot_millic", "cool_millic", "poll_seconds", "cool_samples", "cap_khz"}
    if set(data) != required:
        raise ValueError("配置字段不完整或含未知字段")
    if data["governor"] != "schedutil":
        raise ValueError("本方案仅支持已验证的 schedutil")
    for key in required - {"governor"}:
        if type(data[key]) is not int:
            raise ValueError(f"{key} 必须是整数")
    if not 0 < data["cool_millic"] < data["hot_millic"] <= 100000:
        raise ValueError("温度必须满足 0 < 恢复阈值 < 限频阈值 <= 100°C")
    if not 1 <= data["poll_seconds"] <= 60 or not 1 <= data["cool_samples"] <= 60:
        raise ValueError("轮询秒数与恢复确认次数必须在 1~60 之间")
    if data["cap_khz"] <= 0:
        raise ValueError("限频值必须为正数")
    return data


class Guard:
    def __init__(self, cpu_root, thermal_root, state_dir, config):
        self.cpu_root = Path(cpu_root)
        self.thermal_root = Path(thermal_root)
        self.state_dir = Path(state_dir)
        self.config = config
        self.original = None
        self.mode = None
        self.cool_count = 0
        self.last_temp = None
        self.expected_zones = set(self.thermal_root.glob("thermal_zone[0-9]*"))

    def policies(self):
        return sorted(self.cpu_root.glob("policy[0-9]*"))

    def snapshot(self):
        policies = self.policies()
        if not policies:
            raise RuntimeError("未发现 CPU 调频策略，拒绝启动")
        result = {}
        for policy in policies:
            result[policy.name] = {
                "governor": read(policy / "scaling_governor"),
                "min": int(read(policy / "scaling_min_freq")),
                "max": int(read(policy / "scaling_max_freq")),
                "hardware_max": int(read(policy / "cpuinfo_max_freq")),
                "cpus": read(policy / "affected_cpus"),
            }
        return result

    def prepare(self):
        self.state_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
        current = self.snapshot()
        baseline = self.state_dir / "original.json"
        if baseline.exists():
            self.original = json.loads(baseline.read_text())
        else:
            self.original = current
            # 首次更改前持久保存原始值；后续重启绝不覆盖。
            atomic_json(baseline, current)
        if set(current) != set(self.original):
            raise RuntimeError("CPU 策略集合已变化，请人工核对原始快照")
        for policy in self.policies():
            old = self.original[policy.name]
            if old["cpus"] != read(policy / "affected_cpus") or old["hardware_max"] != int(read(policy / "cpuinfo_max_freq")):
                raise RuntimeError("CPU 拓扑或硬件频率已变化，请人工核对快照")
            if self.config["governor"] not in read(policy / "scaling_available_governors").split():
                raise RuntimeError(f"{policy.name} 不支持 schedutil")
            frequencies = [int(x) for x in read(policy / "scaling_available_frequencies").split()]
            low = int(read(policy / "cpuinfo_min_freq"))
            if old["max"] > old["hardware_max"] or old["max"] < low or not low <= old["min"] <= old["max"]:
                raise RuntimeError("原始频率范围无效")
            if not any(low <= f <= min(self.config["cap_khz"], old["max"]) for f in frequencies):
                raise RuntimeError(f"{policy.name} 没有适用的保护档位")

    def cap(self, policy):
        ceiling = min(self.config["cap_khz"], self.original[policy.name]["max"])
        return max(int(x) for x in read(policy / "scaling_available_frequencies").split() if int(x) <= ceiling)

    def apply(self, limited):
        for policy in self.policies():
            old = self.original[policy.name]
            # 先切策略，再放开最低频率，最后写上限，避免 min > max。
            write(policy / "scaling_governor", self.config["governor"])
            low = int(read(policy / "cpuinfo_min_freq"))
            # scaling_min_freq 是聚合后的有效 QoS 最低值，其他内核请求可短时抬高。
            # 只设置本方案的最低请求，不把合法的外部 boost 判成写入失败。
            write(policy / "scaling_min_freq", low, verify=False)
            write(policy / "scaling_max_freq", self.cap(policy) if limited else old["max"], ceiling=True)

    def temperature(self):
        zones = sorted(self.thermal_root.glob("thermal_zone[0-9]*"))
        if not zones or set(zones) != self.expected_zones:
            return None
        values = []
        for zone in zones:
            try:
                value = int(read(zone / "temp"))
                if not 0 <= value <= 150000:
                    return None
                values.append(value)
            except (OSError, ValueError):
                return None
        return max(values)

    def step(self):
        temp = self.temperature()
        self.last_temp = temp
        if temp is None:
            mode = "sensor-fault"
            self.cool_count = 0
        elif temp >= self.config["hot_millic"] or (self.mode is None and temp > self.config["cool_millic"]):
            mode = "hot"
            self.cool_count = 0
        elif self.mode in ("hot", "sensor-fault"):
            self.cool_count = self.cool_count + 1 if temp <= self.config["cool_millic"] else 0
            mode = "normal" if self.cool_count >= self.config["cool_samples"] else self.mode
        else:
            mode = "normal"
        self.apply(mode != "normal")
        if mode != self.mode:
            LOG.warning("温控状态 %s -> %s，最高温度=%s m°C", self.mode, mode, temp)
        self.mode = mode
        status = {
            "mode": mode, "temperature_millic": temp, "updated_unix": time.time(),
            "cool_count": self.cool_count, "config": self.config,
            "policies": {p.name: {"min_khz": int(read(p / "scaling_min_freq")), "max_khz": int(read(p / "scaling_max_freq")), "governor": read(p / "scaling_governor")} for p in self.policies()},
        }
        # 状态放在 tmpfs，每次轮询不写 eMMC。
        atomic_json(self.state_dir / "status.json", status)
        return mode

    def restore(self):
        original = json.loads((self.state_dir / "original.json").read_text())
        current = self.snapshot()
        if set(current) != set(original):
            raise RuntimeError("策略集合改变，拒绝自动回滚")
        for policy in self.policies():
            old = original[policy.name]
            if current[policy.name]["cpus"] != old["cpus"] or current[policy.name]["hardware_max"] != old["hardware_max"]:
                raise RuntimeError("硬件拓扑改变，拒绝自动回滚")
        for policy in self.policies():
            old = original[policy.name]
            write(policy / "scaling_min_freq", int(read(policy / "cpuinfo_min_freq")), verify=False)
            write(policy / "scaling_max_freq", old["max"])
            write(policy / "scaling_governor", old["governor"])
            write(policy / "scaling_min_freq", old["min"])
            write(policy / "scaling_max_freq", old["max"])
        LOG.warning("已恢复部署前 CPU governor/min/max")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", default="/etc/default/cpu-freq-guard.json")
    parser.add_argument("--restore-original", action="store_true")
    args = parser.parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    state_dir = Path("/var/lib/cpu-freq-guard")
    runtime_dir = Path("/run/cpu-freq-guard")
    state_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
    runtime_dir.mkdir(parents=True, exist_ok=True, mode=0o700)
    # 守护与回滚共用文件锁，避免两个进程争写 sysfs。
    with (runtime_dir / "lock").open("a") as lock:
        try:
            fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError as exc:
            raise RuntimeError("守护进程仍在运行，请先停止服务") from exc
        guard = Guard("/sys/devices/system/cpu/cpufreq", "/sys/class/thermal", state_dir, load_config(Path(args.config)))
        if args.restore_original:
            guard.restore()
            return
        guard.prepare()
        # 实机使用持久基线；运行状态单独写入 /run。
        persistent_dir = guard.state_dir
        guard.state_dir = runtime_dir
        stop = threading.Event()
        signal.signal(signal.SIGTERM, lambda *_: stop.set())
        signal.signal(signal.SIGINT, lambda *_: stop.set())
        try:
            while not stop.is_set():
                guard.step()
                stop.wait(guard.config["poll_seconds"])
        finally:
            guard.state_dir = persistent_dir
            # 停止时保留当前安全设置，不自动恢复到锁定满频。
            LOG.warning("守护已停止，保留当前频率限制；回滚须显式执行 --restore-original")


if __name__ == "__main__":
    try:
        main()
    except (OSError, ValueError, RuntimeError) as error:
        LOG.error("CPU 调频守护失败：%s", error)
        raise SystemExit(1) from error