#!/usr/bin/env python3
"""Patch /home/guru/gururmm/agent/src/updater/mod.rs with two fixes:
1. Replace restart_service() Windows block with detached-cmd approach.
2. Fix hardcoded "gururmm-agent" service name in create_windows_rollback_watchdog.
"""

import sys

PATH = "/home/guru/gururmm/agent/src/updater/mod.rs"

with open(PATH, "r") as f:
    src = f.read()

original = src

# ---------------------------------------------------------------------------
# Fix 1: restart_service()
# ---------------------------------------------------------------------------

OLD_RESTART = """\
    /// Restart the agent service
    async fn restart_service(&self) -> Result<()> {
        #[cfg(unix)]
        {
            info!("Exiting for service restart by systemd");
            std::process::exit(0);
        }

        #[cfg(windows)]
        {
            use crate::service::windows::SERVICE_NAME;
            // Restart Windows service
            tokio::process::Command::new("sc.exe")
                .args(["stop", SERVICE_NAME])
                .status()
                .await?;

            tokio::time::sleep(std::time::Duration::from_secs(2)).await;

            tokio::process::Command::new("sc.exe")
                .args(["start", SERVICE_NAME])
                .status()
                .await?;
        }

        // Give the new process a moment to start
        tokio::time::sleep(std::time::Duration::from_secs(1)).await;

        // Exit this process - the new version should be running now
        std::process::exit(0);
    }"""

NEW_RESTART = """\
    /// Restart the agent service
    async fn restart_service(&self) -> Result<()> {
        #[cfg(unix)]
        {
            info!("Exiting for service restart by systemd");
            std::process::exit(0);
        }

        #[cfg(windows)]
        {
            use crate::service::windows::SERVICE_NAME;

            // Spawn a detached cmd.exe that waits for this process to exit,
            // then starts the service. We cannot await sc start here because
            // calling sc stop on our own service causes the SCM to send
            // SERVICE_CONTROL_STOP back to this process, which tears down the
            // tokio runtime and cancels any pending awaits before sc start runs.
            let restart_cmd = format!(
                "timeout /t 5 /nobreak >nul 2>&1 & sc.exe start {}",
                SERVICE_NAME
            );
            match std::process::Command::new("cmd")
                .args(["/c", &restart_cmd])
                .creation_flags(0x08000000) // CREATE_NO_WINDOW
                .spawn()
            {
                Ok(_) => info!("Restart helper spawned, exiting for service stop/start"),
                Err(e) => error!("Failed to spawn restart helper: {} — service will stay stopped", e),
            }

            // Exit this process. SCM detects the exit and marks the service
            // Stopped. The detached cmd above will sc start after 5 seconds.
            std::process::exit(0);
        }

        // Unreachable on Windows (process::exit above), but needed for
        // Unix path and to satisfy the return type.
        #[allow(unreachable_code)]
        Ok(())
    }"""

if OLD_RESTART not in src:
    print("ERROR: restart_service old block not found in file", file=sys.stderr)
    idx = src.find("async fn restart_service")
    if idx != -1:
        print(f"restart_service found at offset {idx}:", file=sys.stderr)
        print(repr(src[idx:idx+1000]), file=sys.stderr)
    sys.exit(1)

src = src.replace(OLD_RESTART, NEW_RESTART, 1)
print("[OK] Fix 1 applied: restart_service() replaced with detached-cmd approach")

# ---------------------------------------------------------------------------
# Fix 2: create_windows_rollback_watchdog
# ---------------------------------------------------------------------------

# 2a: Replace three hardcoded service name strings in the PS script template
replacements = [
    ('$service = Get-Service -Name "gururmm-agent" -ErrorAction SilentlyContinue',
     '$service = Get-Service -Name $ServiceName -ErrorAction SilentlyContinue',
     'Get-Service line'),
    ('Stop-Service -Name "gururmm-agent" -Force -ErrorAction SilentlyContinue',
     'Stop-Service -Name $ServiceName -Force -ErrorAction SilentlyContinue',
     'Stop-Service line'),
    ('Start-Service -Name "gururmm-agent"',
     'Start-Service -Name $ServiceName',
     'Start-Service line'),
]

for old, new, label in replacements:
    if old not in src:
        print(f"ERROR: {label} not found", file=sys.stderr)
        sys.exit(1)
    src = src.replace(old, new, 1)
    print(f"[OK] Fix 2a: replaced {label}")

# 2b: Add $ServiceName PS variable after $Timeout assignment in the script template
OLD_TIMEOUT_LINE = '$Timeout = {timeout}\n\nStart-Sleep -Seconds $Timeout'
NEW_TIMEOUT_LINE = '$Timeout = {timeout}\n$ServiceName = "{service_name}"\n\nStart-Sleep -Seconds $Timeout'

if OLD_TIMEOUT_LINE not in src:
    print("ERROR: $Timeout line context not found", file=sys.stderr)
    sys.exit(1)
src = src.replace(OLD_TIMEOUT_LINE, NEW_TIMEOUT_LINE, 1)
print("[OK] Fix 2b: $ServiceName PS variable added to script template")

# 2c: Add service_name = SERVICE_NAME to format! args
# The actual bytes in the file for the format args closing (confirmed via repr):
#   replace('\', "\\"),\n            timeout = timeout\n        );
# In Python string literals (raw bytes via repr):
#   replace(\'\\\\', "\\\\\\\\"),\n            timeout = timeout\n        );
OLD_FORMAT_TAIL = "            timeout = timeout\n        );"
NEW_FORMAT_TAIL = "            timeout = timeout,\n            service_name = SERVICE_NAME\n        );"

# There are two format! calls in this file (Unix watchdog and Windows watchdog).
# The Unix one ends with: script = script_path.display()\n        );
# The Windows one ends with: timeout = timeout\n        );
# So OLD_FORMAT_TAIL is unique to the Windows block. Verify:
count = src.count(OLD_FORMAT_TAIL)
if count != 1:
    print(f"ERROR: expected 1 match for format tail, got {count}", file=sys.stderr)
    sys.exit(1)
src = src.replace(OLD_FORMAT_TAIL, NEW_FORMAT_TAIL, 1)
print("[OK] Fix 2c: service_name added to format! args")

# 2d: Add SERVICE_NAME import at top of the Windows watchdog function body
# The function opens with these lines (no prior use statement):
OLD_WATCHDOG_OPEN = (
    "    async fn create_windows_rollback_watchdog(&self) -> Result<()> {\n"
    "        let backup_path = self.config.backup_path();\n"
)
NEW_WATCHDOG_OPEN = (
    "    async fn create_windows_rollback_watchdog(&self) -> Result<()> {\n"
    "        use crate::service::windows::SERVICE_NAME;\n"
    "        let backup_path = self.config.backup_path();\n"
)

if OLD_WATCHDOG_OPEN not in src:
    print("ERROR: watchdog function opening not found", file=sys.stderr)
    sys.exit(1)
src = src.replace(OLD_WATCHDOG_OPEN, NEW_WATCHDOG_OPEN, 1)
print("[OK] Fix 2d: SERVICE_NAME import added to watchdog function")

# ---------------------------------------------------------------------------
# Write result
# ---------------------------------------------------------------------------
with open(PATH, "w") as f:
    f.write(src)

lines_changed = sum(1 for a, b in zip(original.splitlines(), src.splitlines()) if a != b)
lines_added = len(src.splitlines()) - len(original.splitlines())
print(f"\nFile written. ~{lines_changed} lines changed, {lines_added:+d} net lines.")