off/devices.py

"""Device management: check status, power off, and budget tracking."""

import json
import os
import subprocess
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timedelta
from pathlib import Path

import config

BUDGET_S = config.BUDGET_S
STATE_FILE = Path(__file__).parent / ".device_state.json"
TICK_INTERVAL = config.TICK_INTERVAL
ALLOWED_WEEKDAY_START = config.WEEKDAY_START
ALLOWED_WEEKDAY_END = config.WEEKDAY_END
ALLOWED_WEEKEND_START = config.WEEKEND_START
ALLOWED_WEEKEND_END = config.WEEKEND_END


# ── Budget state (persisted to disk) ─────────────────────────────────

def _load_state() -> dict:
    if STATE_FILE.exists():
        try:
            with open(STATE_FILE) as f:
                return json.load(f)
        except (json.JSONDecodeError, OSError):
            return {}  # corrupt state file — start fresh instead of crashing
    return {}


def _save_state(state: dict):
    # Atomic write so a crash mid-write can't corrupt the state file
    tmp_file = STATE_FILE.with_name(STATE_FILE.name + ".tmp")
    with open(tmp_file, "w") as f:
        json.dump(state, f)
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp_file, STATE_FILE)





# ── Helpers ───────────────────────────────────────────────────────────

def _run(cmd: list[str], *, timeout: int = 10) -> subprocess.CompletedProcess | None:
    """Run a command, return CompletedProcess or None on error."""
    try:
        return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
    except (FileNotFoundError, subprocess.TimeoutExpired):
        return None


# ── TV ────────────────────────────────────────────────────────────────

def tv_check() -> dict:
    """Check if the TV is reachable and powered on.

    Returns an action dict describing the current state.
    """
    result = _run(["ping", "-c", "1", "-W", "1", "192.168.1.12"])
    if not result or result.returncode != 0:
        return {"online": False, "detail": "Offline"}

    _run(["adb", "connect", "192.168.1.12"])
    result = _run(["adb", "shell", "dumpsys", "power"])
    if result and "mWakefulness=Awake" in result.stdout:
        return {"online": True, "detail": "Online"}

    return {"online": False, "detail": "Offline"}


def tv_turnoff() -> dict | None:
    """Send power key to the TV if it's on. Returns action dict or None."""
    state = tv_check()
    if not state["online"]:
        return None

    _run(["adb", "shell", "input", "keyevent", "26"])
    return {"detail": "Screen turned off (power key sent)"}


# ── Gabi ──────────────────────────────────────────────────────────────

def gabi_check() -> bool:
    """Return True if user 'gabi' has an active login session."""
    result = _run(["loginctl", "list-users"])
    if result is None:
        return False
    # columns: UID USER LINGER STATE — username is the second field
    users = {line.split()[1] for line in result.stdout.splitlines() if len(line.split()) > 1}
    return "gabi" in users


def gabi_turnoff() -> dict | None:
    """Log out Gabi if she's logged in. Returns action dict or None."""
    if not gabi_check():
        return {"detail": "Not logged in — no action needed"}

    result = _run(["doas", "loginctl", "terminate-user", "gabi"])
    if result and result.returncode == 0:
        return {"detail": "Session terminated"}

    return {"detail": "Failed to terminate session"}


# ── Gaja ──────────────────────────────────────────────────────────────

# Override with the GAJA_PASSWORD env var if you don't want it in source
GAJA_PASSWORD = os.environ.get("GAJA_PASSWORD", "Nagaja")


def gaja_check() -> dict:
    """Check if Gaja's PC is online and she has an active desktop session.

    Returns an action dict describing the current state.
    """
    result = _run(["ping", "-c", "1", "-W", "1", "192.168.1.122"])
    if not result or result.returncode != 0:
        return {"online": False, "detail": "Offline"}

    # Check for a session with seat assigned (desktop session, not SSH)
    # -F /dev/null: skip system ssh config (broken includes make every ssh call fail)
    check = _run(
        ["sshpass", "-p", GAJA_PASSWORD, "ssh", "-F", "/dev/null",
         "-o", "ConnectTimeout=5",
         "-o", "StrictHostKeyChecking=no", "-F", "/dev/null",
         "gaja@192.168.1.122",
         "loginctl list-sessions --no-pager 2>/dev/null | grep gaja | grep seat"]
    )
    if check is not None and check.stdout.strip():
        return {"online": True, "detail": "Online"}

    return {"online": False, "detail": "Offline"}


def gaja_turnoff() -> dict | None:
    """Log out Gaja if she's logged in. Returns action dict or None."""
    state = gaja_check()
    if not state["online"]:
        return {"detail": "Not online — no action needed"}

    # -F /dev/null: skip system ssh config (broken includes make every ssh call fail)
    _run(
        ["sshpass", "-p", GAJA_PASSWORD, "ssh", "-F", "/dev/null",
         "-o", "ConnectTimeout=5",
         "-o", "StrictHostKeyChecking=no", "-F", "/dev/null",
         "gaja@192.168.1.122",
         "loginctl terminate-user gaja"],
        timeout=5
    )
    return {"detail": "Logged out"}


# ── Device registry ───────────────────────────────────────────────────

DEVICES = {
    "tv":      {"icon": "📺", "name": "TV",       "check": tv_check,  "turnoff": tv_turnoff},
    "gabi":    {"icon": "💻", "name": "Gabi's PC", "check": gabi_check, "turnoff": gabi_turnoff},
    "gaja":    {"icon": "💻", "name": "Gaja's PC", "check": gaja_check, "turnoff": gaja_turnoff},
}


def _get_budget(device: str) -> float:
    """Get remaining budget in seconds (float for precision)."""
    state = _load_state()
    value = state.get(device, {}).get("budget", float(BUDGET_S))
    return float(value)  # ensure always float


def _set_budget(device: str, budget_seconds: float):
    """Persist the remaining budget for a device."""
    state = _load_state()
    if device not in state:
        state[device] = {"budget": float(BUDGET_S)}
    state[device]["budget"] = budget_seconds
    _save_state(state)


def _reset_budget(device: str):
    _set_budget(device, float(BUDGET_S))


def _budget_to_minutes(budget_seconds: float) -> int:
    """Convert budget seconds to minutes for display (floor)."""
    return max(0, int(budget_seconds) // 60)


def is_curfew_allowed() -> bool:
    """Check if current time is within allowed hours (devices allowed).

    Allowed hours:
    Mon-Fri: WEEKDAY_START to WEEKDAY_END (minutes since midnight)
    Sat-Sun: WEEKEND_START to WEEKEND_END (minutes since midnight)
    """
    now = datetime.now()
    weekday = now.weekday()  # 0=Mon, 6=Sun
    current_minutes = now.hour * 60 + now.minute

    if weekday < 5:  # Mon-Fri
        start = ALLOWED_WEEKDAY_START
        end = ALLOWED_WEEKDAY_END
    else:  # Sat-Sun
        start = ALLOWED_WEEKEND_START
        end = ALLOWED_WEEKEND_END

    return start <= current_minutes < end


def _minutes_to_time(minutes: int) -> str:
    """Convert minutes since midnight to HH:MM format."""
    h = minutes // 60
    m = minutes % 60
    return f"{h:02d}:{m:02d}"


def _next_transition(now: datetime) -> tuple[int, int]:
    """Find the next allowed-hours boundary (start or end) strictly after now.

    Returns (minutes_until, boundary_minutes_since_midnight). Scans up to
    7 days ahead so weekday/weekend schedule switches are handled.
    """
    current_minutes = now.hour * 60 + now.minute
    best: tuple[int, int] | None = None
    for day_offset in range(7):
        day = now + timedelta(days=day_offset)
        if day.weekday() < 5:
            boundaries = (ALLOWED_WEEKDAY_START, ALLOWED_WEEKDAY_END)
        else:
            boundaries = (ALLOWED_WEEKEND_START, ALLOWED_WEEKEND_END)
        for boundary in boundaries:
            minutes_until = day_offset * (24 * 60) + boundary - current_minutes
            if minutes_until > 0 and (best is None or minutes_until < best[0]):
                best = (minutes_until, boundary)
    return best


def curfew_status() -> dict:
    """Get curfew status and time until next transition.

    Returns dict with:
      - in_curfew: bool (True = devices BLOCKED)
      - minutes_left: int (minutes until curfew ends or starts)
      - message: str (human-readable description)
      - next_time: str (clock time of next transition, e.g. '15:00')
    """
    now = datetime.now()
    weekday = now.weekday()  # 0=Mon, 6=Sun
    current_minutes = now.hour * 60 + now.minute

    if weekday < 5:  # Mon-Fri
        start = ALLOWED_WEEKDAY_START
        end = ALLOWED_WEEKDAY_END
    else:  # Sat-Sun
        start = ALLOWED_WEEKEND_START
        end = ALLOWED_WEEKEND_END

    in_curfew = not (start <= current_minutes < end)  # True when devices are BLOCKED

    # Find the next schedule boundary (handles Mon↔weekend switches correctly)
    minutes_left, boundary = _next_transition(now)
    next_time = _minutes_to_time(boundary)

    if in_curfew:
        message = f"Blocked until {next_time} ({minutes_left} min)"
    else:
        message = f"Allowed until {next_time} ({minutes_left} min)"

    return {
        "in_curfew": in_curfew,
        "minutes_left": minutes_left,
        "message": message,
        "next_time": next_time,
    }


def _shutdown_devices(device_ids: list[str]) -> dict[str, dict | None]:
    """Shut down a list of devices in parallel.

    Returns {dev_id: turnoff result or None on failure}.
    """
    results: dict[str, dict | None] = {}
    if not device_ids:
        return results
    with ThreadPoolExecutor(max_workers=len(device_ids)) as executor:
        futures = {
            executor.submit(DEVICES[dev_id]["turnoff"]): dev_id
            for dev_id in device_ids
        }
        for future, dev_id in futures.items():
            try:
                results[dev_id] = future.result()
            except Exception:
                results[dev_id] = None
    return results


# ── Budget timer (runs every 10s in background) ──────────────────────

_timer_running = False
_timer_lock = threading.Lock()
_last_check_times: dict[str, float] = {}  # in-memory only, not persisted


def _reset_budgets_if_new_day():
    """Reset all budgets if today's 7 AM daily reset was missed.

    The state file is only written while the server is running, so its
    mtime is the last time the server did anything. If that's before
    today's 7 AM (and it's now past 7 AM), the server was stopped across
    the daily reset — give every device a fresh budget.
    """
    if not STATE_FILE.exists():
        return
    now = datetime.now()
    today_reset = now.replace(hour=7, minute=0, second=0)
    mtime = datetime.fromtimestamp(STATE_FILE.stat().st_mtime)
    if mtime < today_reset <= now:
        for dev_id in DEVICES:
            _reset_budget(dev_id)


def _budget_tick():
    """Called every 10 seconds by the background timer."""
    # New day and we were offline at 7 AM? Reset all budgets to full.
    _reset_budgets_if_new_day()

    # Check all devices in parallel
    online_status = {}
    with ThreadPoolExecutor(max_workers=3) as executor:
        futures = {
            executor.submit(dev_info["check"]): dev_id
            for dev_id, dev_info in DEVICES.items()
        }
        for future in as_completed(futures):
            dev_id = futures[future]
            result = future.result()
            if isinstance(result, dict):
                online_status[dev_id] = result.get("online", False)
            else:
                online_status[dev_id] = bool(result)

    # Process results sequentially (budget operations need to be thread-safe)
    with _timer_lock:
        current_time = time.time()  # float timestamp for precise elapsed calculation
        
        # Initialize last check times on first run (server restart)
        if not _last_check_times:
            for dev_id in DEVICES:
                _last_check_times[dev_id] = current_time

        if not is_curfew_allowed():
            # Curfew in effect — turn off all online devices in parallel
            online_devices = [dev_id for dev_id, online in online_status.items() if online]
            _shutdown_devices(online_devices)
            # Update last check times even during curfew
            for dev_id in DEVICES:
                _last_check_times[dev_id] = current_time
            return

        devices_to_shutdown = []

        for dev_id, online in online_status.items():
            dev_info = DEVICES[dev_id]
            state = _load_state()
            current_budget = state.get(dev_id, {}).get("budget", float(BUDGET_S))
            last_check = _last_check_times.get(dev_id, current_time)

            # Expired budget + device back online → shut it down
            if current_budget <= 0 and online:
                devices_to_shutdown.append(dev_id)
                _set_budget(dev_id, 0.0)
                continue

            if online and current_budget > 0:
                # Calculate actual elapsed time since last check
                elapsed = current_time - last_check
                if elapsed < 1:  # ignore very small deltas (< 1s)
                    continue
                new_budget = current_budget - elapsed
                _set_budget(dev_id, new_budget)

                if new_budget <= 0:
                    devices_to_shutdown.append(dev_id)
            else:
                # Always update last_check even when offline (in-memory only)
                pass

            # Update last check time for next interval
            _last_check_times[dev_id] = current_time

        # Shutdown all devices that hit zero budget in parallel
        if devices_to_shutdown:
            _shutdown_devices(devices_to_shutdown)


def start_timer():
    """Start the background budget timer (runs every 10s)."""
    global _timer_running
    with _timer_lock:
        if _timer_running:
            return
        _timer_running = True

    def loop():
        while _timer_running:
            try:
                _budget_tick()
            except Exception:
                pass  # Don't crash the timer on device errors
            threading.Event().wait(TICK_INTERVAL)

    t = threading.Thread(target=loop, daemon=True)
    t.start()


def stop_timer():
    """Stop the background timer."""
    global _timer_running
    with _timer_lock:
        _timer_running = False


# ── Orchestrator ──────────────────────────────────────────────────────

def shutdown_all() -> list[dict]:
    """Indiscriminately turn off all devices in parallel.

    Returns ordered list of action dicts (registry order).
    """
    results = _shutdown_devices(list(DEVICES))
    actions = []
    for dev_id, dev_info in DEVICES.items():
        result = results.get(dev_id)
        if result:
            budget_minutes = _budget_to_minutes(_get_budget(dev_id))
            actions.append({
                "icon": dev_info["icon"],
                "title": dev_info["name"],
                "detail": f"{result['detail']} ({budget_minutes} min left)",
            })
    return actions


def _check_device(dev_id: str, dev_info: dict) -> dict:
    """Check a single device and return its status dict."""
    is_online = dev_info["check"]()
    if isinstance(is_online, dict):
        online = is_online.get("online", False)
        detail = is_online.get("detail", "")
    else:
        online = bool(is_online)
        detail = "Online" if online else "Offline"

    state = _load_state()
    budget_seconds = state.get(dev_id, {}).get("budget", float(BUDGET_S))
    budget_minutes = _budget_to_minutes(budget_seconds)

    # Budget warning
    if budget_minutes <= 0:
        icon = "🔴"
    elif budget_minutes < 5:
        icon = "🟠"
    else:
        icon = dev_info["icon"]

    return {
        "id": dev_id,
        "icon": icon,
        "title": dev_info["name"],
        "detail": f"{detail} · {budget_minutes} min left",
        "online": online,
    }


def status_all() -> tuple[list[dict], dict]:
    """Get current status + budget for all devices. Returns (ordered list, curfew)."""
    curfew = curfew_status()

    # Check all devices in parallel
    results = {}
    with ThreadPoolExecutor(max_workers=3) as executor:
        futures = {
            executor.submit(_check_device, dev_id, dev_info): dev_id
            for dev_id, dev_info in DEVICES.items()
        }
        for future in as_completed(futures):
            dev_id = futures[future]
            results[dev_id] = future.result()

    # Return in ordered list (DEVICES preserves insertion order)
    actions = [
        {
            "id": dev_id,
            "icon": results[dev_id]["icon"],
            "title": results[dev_id]["title"],
            "detail": results[dev_id]["detail"],
        }
        for dev_id in DEVICES
    ]
    return actions, curfew