import os
import sys
import datetime
import subprocess
from ctypes import windll, c_long, byref

# ============================================================
#  CONFIGURATION  -  edit these values before pushing
# ============================================================

ACTION   = "ENABLE"     # ENABLE = schedule reboot | DISABLE = remove reboot

SCHEDULE = "THU"        # DAILY | ONCE | MON | TUE | WED | THU | FRI | SAT | SUN
TIME     = "01:00"      # 24-hour format HH:MM  (01:00 = 1 AM, 23:30 = 11:30 PM)
DATE     = "2026-12-25" # Reboot date when SCHEDULE = "ONCE".
                        # ALWAYS enter it as ISO  YYYY-MM-DD  (year-month-day).
                        # The script converts it to the machine's own regional
                        # format automatically, so there is NO MM/DD vs DD/MM
                        # guesswork for the operator.

TASK_NAME = "SPM_Scheduled_Reboot"   # Task Scheduler name.

# ============================================================
#  END CONFIGURATION  -  do not edit below this line
# ============================================================

VALID_DAYS = ["MON", "TUE", "WED", "THU", "FRI", "SAT", "SUN"]

PROGRAMDATA = os.environ.get("PROGRAMDATA", r"C:\ProgramData")
WORK_DIR    = os.path.join(PROGRAMDATA, "ScheduledReboot")
LOG_FILE    = os.path.join(WORK_DIR, "scheduled_reboot.log")
BAT_FILE    = os.path.join(WORK_DIR, "reboot_command.bat")


class disable_file_system_redirection:
    _disable = windll.kernel32.Wow64DisableWow64FsRedirection
    _revert  = windll.kernel32.Wow64RevertWow64FsRedirection

    def __enter__(self):
        self.old_value = c_long()
        self.success = self._disable(byref(self.old_value))

    def __exit__(self, type, value, traceback):
        if self.success:
            self._revert(self.old_value)


def ensure_dir():
    try:
        if not os.path.isdir(WORK_DIR):
            os.makedirs(WORK_DIR)
    except Exception:
        pass


def safe_flush():

    for stream in (sys.stdout, sys.stderr):
        try:
            stream.flush()
        except Exception:
            pass


def log(msg):
    ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    line = ts + "  " + str(msg)
    print line
    safe_flush()
    try:
        ensure_dir()
        f = open(LOG_FILE, "a")
        f.write(line + "\n")
        f.close()
    except Exception:
        pass


def run(cmd_list):

    with disable_file_system_redirection():
        proc = subprocess.Popen(cmd_list,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE)
        out, err = proc.communicate()
    return proc.returncode, (out or "").strip(), (err or "").strip()


def get_schtasks_date(d):

    ps = ("(Get-Date -Year %d -Month %d -Day %d)."
          "ToString((Get-Culture).DateTimeFormat.ShortDatePattern)") % (d.year, d.month, d.day)
    rc, out, err = run(["powershell", "-NoProfile", "-NonInteractive", "-Command", ps])
    if rc == 0 and out:
        return out.strip()
    # Fallback if PowerShell is somehow unavailable
    return d.strftime("%m/%d/%Y")


def write_reboot_batch():

    ensure_dir()
    f = open(BAT_FILE, "w")
    f.write("@echo off\r\n")
    f.write("shutdown.exe /r /f /t 0\r\n")
    f.close()
    log("Reboot command written to: " + BAT_FILE)


def task_exists():
    rc, out, err = run(["schtasks", "/Query", "/TN", TASK_NAME])
    return rc == 0


def valid_time(t):
    if len(t) == 5 and t[2] == ":":
        try:
            hh = int(t[0:2])
            mm = int(t[3:5])
            return (0 <= hh <= 23) and (0 <= mm <= 59)
        except ValueError:
            return False
    return False


def clear_battery_conditions():

    ps = ("Set-ScheduledTask -TaskName '" + TASK_NAME + "' "
          "-Settings (New-ScheduledTaskSettingsSet "
          "-AllowStartIfOnBatteries -DontStopIfGoingOnBatteries)")
    rc, out, err = run(["powershell", "-NoProfile", "-NonInteractive", "-Command", ps])
    if rc == 0:
        log("Battery conditions cleared (task will fire on battery power).")
    else:

        log("WARNING: could not clear battery conditions; task may not fire on battery.")
        if err:
            log("  " + err)


def enable_task():
    sched = SCHEDULE.strip().upper()
    t     = TIME.strip()

    if sched not in (["DAILY", "ONCE"] + VALID_DAYS):
        log("ERROR: SCHEDULE '" + SCHEDULE + "' is invalid. Use DAILY, ONCE, or MON-SUN.")
        return 2

    if not valid_time(t):
        log("ERROR: TIME '" + TIME + "' is invalid. Use 24-hour HH:MM (e.g. 01:00).")
        return 2

    once_sd = ""
    once_human = ""
    if sched == "ONCE":
        if not DATE.strip():
            log("ERROR: SCHEDULE is ONCE but DATE is empty. Set DATE to the reboot date (YYYY-MM-DD).")
            return 2
        try:
            target = datetime.datetime.strptime(DATE.strip(), "%Y-%m-%d").date()
        except ValueError:
            log("ERROR: DATE '" + DATE + "' is not valid ISO format. Use YYYY-MM-DD (e.g. 2026-12-25).")
            return 2

        once_sd = get_schtasks_date(target)
        once_human = target.strftime("%Y-%m-%d")
        log("ONCE reboot date resolved to " + once_human + " (local schtasks format: " + once_sd + ")")

    write_reboot_batch()

    cmd = ["schtasks", "/Create", "/RU", "SYSTEM", "/TN", TASK_NAME,
           "/TR", BAT_FILE, "/ST", t, "/F"]

    if sched == "DAILY":
        cmd += ["/SC", "DAILY"]
    elif sched == "ONCE":
        cmd += ["/SC", "ONCE", "/SD", once_sd]
    else:
        cmd += ["/SC", "WEEKLY", "/D", sched]

    rc, out, err = run(cmd)

    if rc == 0:
        clear_battery_conditions()
        if sched == "ONCE":
            when = "once on " + once_human
        elif sched == "DAILY":
            when = "every day"
        else:
            when = "every " + sched
        log("SUCCESS: Reboot task '" + TASK_NAME + "' scheduled to run " + when + " at " + t + ".")
        return 0

    log("ERROR: Failed to create task. schtasks output:")
    if out:
        log("  " + out)
    if err:
        log("  " + err)
    return 2


def disable_task():
    if not task_exists():
        log("Nothing to do: no scheduled reboot task named '" + TASK_NAME + "' was found.")
        return 1

    rc, out, err = run(["schtasks", "/Delete", "/TN", TASK_NAME, "/F"])

    if rc == 0:
        log("SUCCESS: Scheduled reboot task '" + TASK_NAME + "' removed.")
        try:
            if os.path.exists(BAT_FILE):
                os.remove(BAT_FILE)
                log("Removed reboot command file: " + BAT_FILE)
        except Exception:
            pass
        return 0

    log("ERROR: Failed to remove task. schtasks output:")
    if out:
        log("  " + out)
    if err:
        log("  " + err)
    return 2


def main():
    log("==== Scheduled Reboot procedure started ====")
    act = ACTION.strip().upper()
    log("ACTION = " + act + "  |  TASK = " + TASK_NAME)

    if act == "ENABLE":
        rc = enable_task()
    elif act == "DISABLE":
        rc = disable_task()
    else:
        log("ERROR: ACTION '" + ACTION + "' is invalid. Use ENABLE or DISABLE.")
        rc = 2

    result = {0: "SUCCESS", 1: "NOTHING TO DO", 2: "FAILED", 3: "ERROR"}.get(rc, "UNKNOWN")
    log("==== Scheduled Reboot procedure finished: " + result + " (code " + str(rc) + ") ====")
    return rc


if __name__ == "__main__":

    try:
        main()
    except Exception as e:
        log("UNHANDLED EXCEPTION: " + str(e))