# VM State Recovery (Phase 2)
#
# Recovers VM state after host reboot and restarts permanent VMs.
# Rule 1: Single purpose, Rule 5: Atomic operations

# Selective imports (ADR-025 Phase 3 Layer 2).
use platform/vm/vm_persistence.nu [get-vm-persistence-info list-permanent-vms]
use platform/vm/lifecycle.nu [vm-info vm-start]

export def "recover-vms-on-boot" []: record {
    """
    Recover and restart VMs after host reboot.

    Executed during system startup to restore permanent VMs.
    """

    print "🔄 VM State Recovery: Starting permanent VM restoration..."

    # Get all permanent VMs
    let permanent_vms = (list-permanent-vms)

    if ($permanent_vms | length) == 0 {
        return {
            success: true
            message: "No permanent VMs to recover"
        }
    }

    # Sort by start order
    let sorted_vms = (
        $permanent_vms
        | sort-by {|vm| ($vm.start_order // 100)}
    )

    # Start each VM in order
    let results = (
        $sorted_vms
        | map {|vm|
            start-permanent-vm-on-boot $vm
        }
    )

    let successful = ($results | where success == true | length)
    let failed = ($results | where success == false | length)

    print $"✓ VM Recovery Complete: ($successful) started, ($failed) failed"

    {
        success: ($failed == 0)
        total: ($sorted_vms | length)
        started: $successful
        failed: $failed
        details: $results
    }
}

def start-permanent-vm-on-boot [vm_info: record]: record {
    """Start a permanent VM on boot with retry logic"""

    let vm_name = $vm_info.vm_name
    let start_retries = ($vm_info.start_retries // 3)

    print $"  Starting VM: ($vm_name)..."

    # Wait for VM state to stabilize
    let wait_delay = ($vm_info.start_delay_seconds // 0)
    if $wait_delay > 0 {
        print $"    Waiting ($wait_delay) seconds before start..."
        sleep ($wait_delay)s
    }

    # Try to start VM with retries
    let start_result = (
        (0..$start_retries)
        | reduce {|attempt, result_so_far|
            if $result_so_far.success {
                return $result_so_far
            }

            # Attempt to start VM (no try-catch, guard pattern)
            let try_result = (vm-start $vm_name)

            if $try_result.success {
                {success: true, attempt: ($attempt + 1)}
            } else if $attempt < $start_retries {
                print $"    Retry $($attempt + 1)/$start_retries in 10 seconds..."
                sleep 10s
                $result_so_far
            } else {
                $try_result
            }
        }
    )

    if not $start_result.success {
        print $"  ✗ Failed to start ($vm_name): ($start_result.error)"
        return {
            success: false
            vm_name: $vm_name
            error: $start_result.error
        }
    }

    print $"  ✓ Started ($vm_name)"

    {
        success: true
        vm_name: $vm_name
        message: "VM started successfully"
    }
}

export def "save-vm-state-snapshot" [
    vm_name: string               # VM name
]: record {
    """Save a snapshot of VM state for recovery"""

    let vm_info = (vm-info $vm_name)

    if ("error" in $vm_info) {
        return {
            success: false
            error: $vm_info.error
        }
    }

    let snapshot = {
        vm_name: $vm_name
        snapshot_time: (date now)
        vm_state: $vm_info.status
        ip_address: ($vm_info.ip_address // "")
        cpu_cores: $vm_info.cpu_cores
        memory_mb: $vm_info.memory_mb
    }

    let snapshot_file = (get-snapshot-file $vm_name)

    # Save snapshot (no try-catch)
    let save_result = (do { bash -c $"cat > ($snapshot_file) << 'EOF'\n($snapshot | to json)\nEOF" } | complete)
    if $save_result.exit_code != 0 {
        return {
            success: false
            error: $"Failed to save state snapshot: ($save_result.stderr)"
        }
    }

    {
        success: true
        vm_name: $vm_name
        message: "State snapshot saved"
    }
}

export def "restore-vm-state-snapshot" [
    vm_name: string               # VM name
]: record {
    """Restore VM from saved state snapshot"""

    let snapshot_file = (get-snapshot-file $vm_name)

    if not ($snapshot_file | path exists) {
        return {
            success: false
            error: "No snapshot found for VM"
        }
    }

    # Load snapshot (no try-catch)
    let snap_result = (do { open $snapshot_file | from json } | complete)
    if $snap_result.exit_code != 0 {
        return {
            success: false
            error: $"Failed to load snapshot: ($snap_result.stderr)"
        }
    }

    let snapshot = ($snap_result.stdout)

    # Only restore if it was running
    if $snapshot.vm_state != "running" {
        return {
            success: true
            message: "VM was not running at snapshot time"
        }
    }

    # Start the VM (no try-catch)
    vm-start $vm_name
}

export def "register-vm-autostart" [
    vm_name: string               # VM name
    --start-order: int = 100      # Priority order (lower starts first)
    --start-delay: int = 0        # Seconds to delay before start
    --wait-for-ssh: bool = true   # Wait for SSH before continuing
    --max-retries: int = 3        # Max start attempts
]: record {
    """Register VM for automatic startup on boot"""

    let persist_info = (get-vm-persistence-info $vm_name)

    if not ("mode" in $persist_info) {
        return {
            success: false
            error: "VM not found in persistence"
        }
    }

    let updated = (
        $persist_info
        | upsert auto_start true
        | upsert start_order $start_order
        | upsert start_delay_seconds $start_delay
        | upsert wait_for_ssh $wait_for_ssh
        | upsert max_start_retries $max_retries
    )

    let persist_file = (get-persistence-file $vm_name)

    # Save autostart configuration (no try-catch)
    let save_result = (do { bash -c $"cat > ($persist_file) << 'EOF'\n($updated | to json)\nEOF" } | complete)
    if $save_result.exit_code != 0 {
        return {
            success: false
            error: $"Failed to save autostart configuration: ($save_result.stderr)"
        }
    }

    {
        success: true
        vm_name: $vm_name
        start_order: $start_order
        message: "VM registered for autostart"
    }
}

export def "get-vms-pending-recovery" []: table {
    """Get list of VMs that need recovery after crash"""

    let permanent_vms = (list-permanent-vms)

    $permanent_vms
    | filter {|vm|
        ($vm.auto_start // false) == true
    }
    | map {|vm|
        {
            vm_name: $vm.vm_name
            auto_start: true
            start_order: ($vm.start_order // 100)
            start_delay_seconds: ($vm.start_delay_seconds // 0)
        }
    }
    | sort-by {|vm| $vm.start_order}
}

export def "wait-for-vm-ssh" [
    vm_name: string               # VM name
    --timeout: int = 300          # Timeout in seconds
]: record {
    """Wait for VM SSH to become available"""

    let start_time = (date now)
    let timeout_duration = ($timeout)sec
    mut attempts = 0
    let max_attempts = ($timeout / 2) + 1  # Safety limit based on sleep 2s

    while { $attempts < $max_attempts } {
        let elapsed = ((date now) - $start_time)

        if $elapsed >= $timeout_duration {
            return {
                success: false
                error: $"SSH timeout after ($timeout_seconds) seconds"
            }
        }

        # Check SSH availability (no try-catch)
        let ssh_check = (do { vm-ssh $vm_name --command "echo ok" } | complete)

        if $ssh_check.exit_code == 0 {
            return {
                success: true
                message: "SSH ready"
            }
        }

        sleep 2s
        $attempts += 1
    }

    {
        success: false
        error: $"SSH timeout after ($timeout) seconds"
    }
}

export def "create-vm-recovery-script" []: record {
    """Create systemd service for VM recovery on boot"""

    let script_path = "/usr/local/bin/vm-recovery.sh"
    let service_path = "/etc/systemd/system/vm-recovery.service"

    # Create recovery script
    let script_content = '#!/bin/bash
echo "Starting VM recovery..."
nu -c "use platform/vm/state_recovery.nu *; recover-vms-on-boot"
echo "VM recovery complete"
'

    # Create recovery script (no try-catch)
    let create_result = (do { bash -c $"cat > ($script_path) << 'SCRIPT'\n($script_content)\nSCRIPT" } | complete)
    if $create_result.exit_code != 0 {
        return {
            success: false
            error: $"Failed to create recovery script: ($create_result.stderr)"
        }
    }

    let chmod_result = (do { bash -c $"chmod +x ($script_path)" } | complete)
    if $chmod_result.exit_code != 0 {
        return {
            success: false
            error: $"Failed to set script permissions: ($chmod_result.stderr)"
        }
    }

    # Create systemd service
    let service_content = '[Unit]
Description=VM Recovery Service
After=network.target libvirtd.service
Wants=libvirtd.service

[Service]
Type=oneshot
ExecStart=/usr/local/bin/vm-recovery.sh
RemainAfterExit=yes
StandardOutput=journal
StandardError=journal

[Install]
WantedBy=multi-user.target
'

    # Create systemd service (no try-catch)
    let service_write_result = (do { bash -c $"cat > ($service_path) << 'SERVICE'\n($service_content)\nSERVICE" } | complete)
    if $service_write_result.exit_code != 0 {
        return {
            success: false
            error: $"Failed to write systemd service file: ($service_write_result.stderr)"
        }
    }

    let daemon_reload_result = (do { bash -c "systemctl daemon-reload || true" } | complete)
    if $daemon_reload_result.exit_code != 0 {
        return {
            success: false
            error: $"Failed to reload systemd: ($daemon_reload_result.stderr)"
        }
    }

    let enable_result = (do { bash -c "systemctl enable vm-recovery.service || true" } | complete)
    if $enable_result.exit_code != 0 {
        return {
            success: false
            error: $"Failed to enable systemd service: ($enable_result.stderr)"
        }
    }

    {
        success: true
        script_path: $script_path
        service_path: $service_path
        message: "VM recovery service created"
    }
}

def get-snapshot-file [vm_name: string]: string {
    """Get snapshot file path"""
    $"{{paths.workspace}}/vms/snapshots/($vm_name).json"
}

def get-persistence-file [vm_name: string]: string {
    """Get persistence file"""
    $"{{paths.workspace}}/vms/persistence/($vm_name).json"
}