screen-locker/screen_locker/_shutdown.py

"""Shutdown schedule adjustment mixin for the screen locker."""

from __future__ import annotations

import calendar
from datetime import datetime, timedelta, timezone
import json
import logging
import subprocess
from typing import TYPE_CHECKING

from screen_locker._constants import (
    ADJUST_SHUTDOWN_SCRIPT,
    ALARM_DAYS,
    RTCWAKE_BIN,
    SHUTDOWN_CONFIG_FILE,
    SICK_DAY_STATE_FILE,
    WAKE_AFTER_HOURS,
)

if TYPE_CHECKING:
    from pathlib import Path

_logger = logging.getLogger(__name__)

_SHUTDOWN_CONFIG_KEYS = ("MON_WED_HOUR", "THU_SUN_HOUR", "MORNING_END_HOUR")


def read_shutdown_config(path: Path) -> tuple[int, int, int] | None:
    """Read shutdown config from *path*. Returns (mw_hour, ts_hour, me_hour) or None.

    Reading needs no privilege (only writing does, via
    ``adjust_shutdown_schedule.sh``) — safe to call from a read-only status view.
    """
    if not path.exists():
        _logger.warning("Config not found: %s", path)
        return None
    parsed: dict[str, int] = {}
    with path.open() as f:
        for line in f:
            stripped = line.strip()
            for key in _SHUTDOWN_CONFIG_KEYS:
                if stripped.startswith(f"{key}="):
                    parsed[key] = int(stripped.split("=")[1])
    if len(parsed) < len(_SHUTDOWN_CONFIG_KEYS):
        _logger.warning("Shutdown config missing required values")
        return None
    return (
        parsed["MON_WED_HOUR"],
        parsed["THU_SUN_HOUR"],
        parsed["MORNING_END_HOUR"],
    )


class ShutdownMixin:
    """Mixin providing shutdown schedule adjustment functionality."""

    def _apply_earlier_shutdown(self, today: str) -> bool:
        """Read config, save state, and write earlier shutdown hours."""
        config_values = self._read_shutdown_config()
        if config_values is None:
            return False
        mon_wed_hour, thu_sun_hour, morning_end_hour = config_values
        if not self._save_sick_day_state(today, mon_wed_hour, thu_sun_hour):
            _logger.error("Failed to save state - aborting adjustment")
            return False
        new_mon_wed = max(18, mon_wed_hour - 1)
        new_thu_sun = max(18, thu_sun_hour - 1)
        return self._write_shutdown_config(
            new_mon_wed,
            new_thu_sun,
            morning_end_hour,
        )

    def _adjust_shutdown_time_earlier(self) -> bool:
        """Adjust shutdown schedule 1.5 hours earlier (stricter).

        This can only be used once per day. Original values are saved and
        automatically restored when checked the next day.

        Returns True if successful, False otherwise.
        """
        today = datetime.now(tz=timezone.utc).strftime("%Y-%m-%d")
        self._restore_original_config_if_needed()
        if self._sick_mode_used_today():
            _logger.warning("Sick mode already used today")
            return False
        try:
            return self._apply_earlier_shutdown(today)
        except (OSError, ValueError) as e:
            _logger.warning("Failed to adjust shutdown time: %s", e)
            return False

    def _adjust_shutdown_time_later(self) -> bool:
        """Adjust shutdown schedule 2 hours later as workout reward.

        Returns True if successful, False otherwise.
        """
        try:
            config_values = self._read_shutdown_config()
            if config_values is None:
                return False
            mon_wed_hour, thu_sun_hour, morning_end_hour = config_values
            new_mon_wed = min(23, mon_wed_hour + 2)
            new_thu_sun = min(23, thu_sun_hour + 2)
            return self._write_shutdown_config(
                new_mon_wed,
                new_thu_sun,
                morning_end_hour,
                restore=True,
            )
        except (OSError, ValueError) as e:
            _logger.warning("Failed to adjust shutdown time for workout: %s", e)
            return False

    def _adjust_shutdown_time_by(self, extra_hours: int) -> bool:
        """Adjust shutdown hours by *extra_hours*, capped at 24 (midnight).

        Used for extra-workout bonuses beyond the weekly minimum.  A cap of 24
        works because ``day-specific-shutdown-check.sh`` fires at 00:00 and
        catches it via the morning-window condition (0 <= 300 minutes).

        Returns True if successful, False otherwise.
        """
        try:
            config_values = self._read_shutdown_config()
            if config_values is None:
                return False
            mw, ts, morning = config_values
            return self._write_shutdown_config(
                min(24, mw + extra_hours),
                min(24, ts + extra_hours),
                morning,
                restore=True,
            )
        except (OSError, ValueError) as e:
            _logger.warning(
                "Failed to adjust shutdown time by %d h: %s", extra_hours, e
            )
            return False

    def _sick_mode_used_today(self) -> bool:
        """Check if sick mode was already used today."""
        if not SICK_DAY_STATE_FILE.exists():
            return False

        try:
            with SICK_DAY_STATE_FILE.open() as f:
                state = json.load(f)
            today = datetime.now(tz=timezone.utc).strftime("%Y-%m-%d")
            return state.get("date") == today
        except (OSError, json.JSONDecodeError) as exc:
            _logger.warning(
                "Could not read sick-day state from %s: %s — treating sick mode "
                "as UNUSED today, which may allow a second sick day",
                SICK_DAY_STATE_FILE,
                exc,
            )
            return False

    def _save_sick_day_state(
        self,
        date: str,
        orig_mon_wed: int,
        orig_thu_sun: int,
    ) -> bool:
        """Save sick day state with original config values.

        Returns True if saved successfully, False otherwise.
        """
        state = {
            "date": date,
            "original_mon_wed_hour": orig_mon_wed,
            "original_thu_sun_hour": orig_thu_sun,
        }
        try:
            with SICK_DAY_STATE_FILE.open("w") as f:
                json.dump(state, f, indent=2)
        except OSError as e:
            _logger.warning("Failed to save sick day state: %s", e)
            return False

        _logger.info("Saved sick day state for %s", date)
        return True

    def _load_sick_day_state(self) -> tuple[str, int, int] | None:
        """Load sick day state file.

        Returns (date, orig_mon_wed_hour, orig_thu_sun_hour) or None.
        """
        with SICK_DAY_STATE_FILE.open() as f:
            state = json.load(f)
        date = state.get("date")
        orig_mw = state.get("original_mon_wed_hour")
        orig_ts = state.get("original_thu_sun_hour")
        if date is None or orig_mw is None or orig_ts is None:
            return None
        return (str(date), int(orig_mw), int(orig_ts))

    def _write_restored_config(
        self,
        orig_mw: int,
        orig_ts: int,
        state_date: str,
    ) -> None:
        """Write restored config values and clean up state file."""
        config_values = self._read_shutdown_config()
        if config_values:
            _, _, morning_end = config_values
            _logger.info(
                "Restoring original shutdown config from %s",
                state_date,
            )
            self._write_shutdown_config(
                orig_mw,
                orig_ts,
                morning_end,
                restore=True,
            )
        SICK_DAY_STATE_FILE.unlink()
        _logger.info("Removed stale sick day state from %s", state_date)

    def _restore_original_config_if_needed(self) -> None:
        """Restore original config if sick day state is from a previous day."""
        if not SICK_DAY_STATE_FILE.exists():
            return
        try:
            loaded = self._load_sick_day_state()
            if loaded is None:
                return
            state_date, orig_mw, orig_ts = loaded
            today = datetime.now(tz=timezone.utc).strftime("%Y-%m-%d")
            if state_date != today:
                self._write_restored_config(orig_mw, orig_ts, state_date)
        except (OSError, json.JSONDecodeError) as e:
            _logger.warning("Error checking sick day state: %s", e)

    def _read_shutdown_config(self) -> tuple[int, int, int] | None:
        """Read shutdown config. Returns (mw_hour, ts_hour, me_hour) or None."""
        return read_shutdown_config(SHUTDOWN_CONFIG_FILE)

    def _build_shutdown_cmd(
        self,
        mon_wed: int,
        thu_sun: int,
        morning: int,
        *,
        restore: bool,
    ) -> list[str]:
        """Build the shutdown adjustment command."""
        cmd = ["/usr/bin/sudo", str(ADJUST_SHUTDOWN_SCRIPT)]
        if restore:
            cmd.append("--restore")
        cmd.extend([str(mon_wed), str(thu_sun), str(morning)])
        return cmd

    def _write_shutdown_config(
        self,
        mon_wed_hour: int,
        thu_sun_hour: int,
        morning_end_hour: int,
        *,
        restore: bool = False,
    ) -> bool:
        """Write new shutdown config values using helper script.

        Args:
            mon_wed_hour: Shutdown hour for Monday-Wednesday.
            thu_sun_hour: Shutdown hour for Thursday-Sunday.
            morning_end_hour: Morning end hour.
            restore: If True, allows restoring to later times.

        Returns True if successful, False otherwise.
        """
        if not ADJUST_SHUTDOWN_SCRIPT.exists():
            _logger.warning(
                "Script not found: %s",
                ADJUST_SHUTDOWN_SCRIPT,
            )
            return False
        cmd = self._build_shutdown_cmd(
            mon_wed_hour,
            thu_sun_hour,
            morning_end_hour,
            restore=restore,
        )
        return self._run_shutdown_cmd(cmd, mon_wed_hour, thu_sun_hour)

    def _run_shutdown_cmd(
        self,
        cmd: list[str],
        mon_wed_hour: int,
        thu_sun_hour: int,
    ) -> bool:
        """Execute the shutdown adjustment command."""
        try:
            result = subprocess.run(
                cmd,
                check=True,
                capture_output=True,
                text=True,
            )
        except subprocess.SubprocessError as e:
            _logger.warning("Failed to adjust shutdown config: %s", e)
            return False
        _logger.info(
            "Adjusted shutdown: Mon-Wed=%d, Thu-Sun=%d. %s",
            mon_wed_hour,
            thu_sun_hour,
            result.stdout.strip(),
        )
        return True

    # ------------------------------------------------------------------
    # rtcwake integration for weekend wake alarm
    # ------------------------------------------------------------------

    @staticmethod
    def _is_tomorrow_alarm_day() -> bool:
        """Check if tomorrow is an alarm day."""
        tomorrow = datetime.now(tz=timezone.utc) + timedelta(days=1)
        return tomorrow.weekday() in ALARM_DAYS

    @staticmethod
    def _compute_wake_timestamp() -> int:
        """Compute the UTC epoch timestamp for the next wake alarm.

        Returns:
            Epoch seconds WAKE_AFTER_HOURS from now.
        """
        wake_time = datetime.now(tz=timezone.utc) + timedelta(
            hours=WAKE_AFTER_HOURS,
        )
        return calendar.timegm(wake_time.utctimetuple())

    @staticmethod
    def _schedule_rtcwake() -> bool:
        """Set rtcwake to power on the PC after WAKE_AFTER_HOURS.

        Uses ``rtcwake -m disk`` to hibernate immediately while programming
        the RTC to restore power at wake_epoch.  Hibernate is completely
        silent and dark (state written to swap file), making it suitable
        when the PC is in a bedroom.

        Returns:
            True if rtcwake was set successfully, False otherwise.
        """
        wake_epoch = ShutdownMixin._compute_wake_timestamp()
        cmd = [
            "/usr/bin/sudo",
            RTCWAKE_BIN,
            "-m",
            "disk",
            "-t",
            str(wake_epoch),
        ]
        try:
            subprocess.run(
                cmd,
                check=True,
                capture_output=True,
                text=True,
            )
        except subprocess.SubprocessError as exc:
            _logger.warning("Failed to set rtcwake: %s", exc)
            return False
        _logger.info(
            "rtcwake set: PC will wake at epoch %d",
            wake_epoch,
        )
        return True

    def schedule_wake_if_needed(self) -> bool:
        """Schedule rtcwake if tomorrow is an alarm day.

        Call this at shutdown time.

        Returns:
            True if wake was scheduled, False if not needed or failed.
        """
        if not self._is_tomorrow_alarm_day():
            _logger.info("Tomorrow is not an alarm day — skipping rtcwake")
            return False
        return self._schedule_rtcwake()