#!/usr/bin/python3
"""At midnight, arm a 06:00 RTC alarm and request an orderly shutdown."""
import argparse
import datetime as dt
from pathlib import Path
import shutil
import subprocess
TZ = dt.timezone(dt.timedelta(hours=8), 'Asia/Shanghai')
ALARM = Path('/sys/class/rtc/rtc0/wakealarm')


def wake_time(now):
    target = now.astimezone(TZ).replace(hour=6, minute=0, second=0, microsecond=0)
    return target if target > now else target + dt.timedelta(days=1)


def power_cycle(now, check=False):
    target = wake_time(now)
    print(f'当前时间：{now.astimezone(TZ).isoformat(timespec="seconds")}', flush=True)
    print(f'计划唤醒：{target.isoformat(timespec="seconds")}', flush=True)
    rtcwake = shutil.which('rtcwake')
    systemctl = shutil.which('systemctl')
    if not rtcwake or not systemctl or not ALARM.exists():
        raise RuntimeError('未发现 RTC 唤醒接口或必要命令，本次不执行关机。')
    if check:
        print('检查完成；没有设置闹钟，也没有执行关机。', flush=True)
        return
    local = now.astimezone(TZ)
    if local.hour != 0 or local.minute >= 5:
        raise RuntimeError('此任务仅允许在北京时间 00:00–00:04 执行；本次不关机。')
    # This server stores its hardware clock in UTC (timedatectl LocalRTC=no).
    subprocess.run([rtcwake, '--utc', '--device', '/dev/rtc0', '--mode', 'no',
                    '--time', str(int(target.timestamp()))], check=True, timeout=20)
    alarm = int(ALARM.read_text().strip() or '0')
    if abs(alarm - target.timestamp()) > 2:
        raise RuntimeError('RTC 唤醒时间回读不一致，本次不执行关机。')
    print('RTC 闹钟已确认，开始正常关机。正在运行的计算任务将停止。', flush=True)
    subprocess.run([systemctl, '--no-block', 'poweroff'], check=True, timeout=20)


def main():
    parser = argparse.ArgumentParser(description='每天 00:00 关机，设置当天 06:00 RTC 唤醒。')
    parser.add_argument('--check', action='store_true', help='仅检查环境和时间，不设置闹钟、不关机')
    args = parser.parse_args()
    try:
        power_cycle(dt.datetime.now(TZ), args.check)
    except (RuntimeError, OSError, ValueError, subprocess.SubprocessError) as exc:
        parser.exit(1, f'夜间节电计划未完成：{exc}\n')


if __name__ == '__main__':
    main()