#!/usr/bin/env nu
# Build a golden BuildKit runner snapshot from Debian 12 ARM64 base.
# Idempotent: re-running with the same date tag is a no-op when the snapshot already exists.
#
# Requires: hcloud CLI authenticated (HCLOUD_TOKEN set), SSH agent with orchestrator key.
# SSH key coupling: the orchestrator pubkey baked into the golden image is the one injected
# via --ssh-key at server creation. Key rotation requires a fresh golden image build — document
# in bootstrap playbook and verify fingerprint before declaring an image good.

use std log

const GOLDEN_IMAGE_PREFIX = "buildkit-runner-golden"
const CLOUD_INIT_REL_PATH = ["..","cloud-init","runner.yaml"]
const VERIFY_TOOLS = ["buildctl","sccache","nu"]

def main [
    --location: string = "fsn1"
    --server_type: string = "cax21"
    --ssh_key: string = ""
    --date_tag: string = ""
    --cloud_init: string = ""
]: nothing -> nothing {
    let tag = if ($date_tag | is-empty) {
        date now | format date "%Y-%m-%d"
    } else {
        $date_tag
    }

    let snapshot_name = $"($GOLDEN_IMAGE_PREFIX):($tag)"
    let seed_name = $"buildkit-runner-seed-($tag)"

    # Resolve ssh_key name to filesystem path for SSH commands.
    # Hetzner key names use dashes; look in ~/.ssh/ for the matching private key.
    let ssh_key_path = if ($ssh_key | is-not-empty) {
        let candidate = ($"~/.ssh/($ssh_key)" | path expand)
        if ($candidate | path exists) { $candidate } else { "" }
    } else { "" }

    let cloud_init_path = if ($cloud_init | is-empty) {
        $env.FILE_PWD | path join ...$CLOUD_INIT_REL_PATH | path expand
    } else {
        $cloud_init | path expand
    }

    if not ($cloud_init_path | path exists) {
        error make { msg: $"cloud-init file not found: ($cloud_init_path)" }
    }

    log info $"Target snapshot: ($snapshot_name)"

    # Idempotency check — skip if today's snapshot already exists
    let existing_out = do {
        ^hcloud image list --type snapshot --selector $"app=($GOLDEN_IMAGE_PREFIX)" --output json
    } | complete
    if $existing_out.exit_code != 0 {
        error make { msg: $"hcloud image list failed:\n($existing_out.stderr)" }
    }
    let images = $existing_out.stdout | from json
    let already_exists = $images | where { ($in | get -o labels | default {} | get -o buildkit_tag | default "") == $tag }
    if ($already_exists | length) > 0 {
        log info $"Snapshot ($snapshot_name) already exists — no-op."
        return
    }

    # Create seed server
    log info $"Creating seed server ($seed_name) at ($location) type ($server_type)..."
    let create_args = (
        ["server","create","--name",$seed_name,"--image","debian-12","--type",$server_type,
         "--location",$location,"--user-data-from-file",$cloud_init_path,
         "--label",$"app=($GOLDEN_IMAGE_PREFIX)"]
        | if ($ssh_key | is-not-empty) { append ["--ssh-key",$ssh_key] } else { $in }
    )
    let create_out = do { ^hcloud ...$create_args } | complete
    if $create_out.exit_code != 0 {
        error make { msg: $"hcloud server create failed:\n($create_out.stderr)" }
    }

    # Retrieve server details by name
    let describe_out = do { ^hcloud server describe $seed_name --output json } | complete
    if $describe_out.exit_code != 0 {
        error make { msg: $"hcloud server describe failed:\n($describe_out.stderr)" }
    }
    let server = $describe_out.stdout | from json
    let server_id = $server.id
    let server_ip = $server.public_net.ipv4.ip

    log info $"Seed server id=($server_id) ip=($server_ip) — waiting for running state..."
    _wait_for_status $server_id "running" 300

    log info "Waiting for SSH and cloud-init completion..."
    _wait_for_ssh $server_ip 420 $ssh_key_path

    log info "Verifying tools..."
    _verify_tools $server_ip $ssh_key_path

    # Snapshot the seed server
    log info $"Creating snapshot '($snapshot_name)'..."
    let snap_out = do {
        ^hcloud server create-image $server_id --type snapshot --description $snapshot_name --label $"app=($GOLDEN_IMAGE_PREFIX)" --label $"buildkit_tag=($tag)"
    } | complete
    if $snap_out.exit_code != 0 {
        _cleanup_server $server_id
        error make { msg: $"Snapshot creation failed:\n($snap_out.stderr)" }
    }

    _cleanup_server $server_id

    let snap_list = do {
        ^hcloud image list --type snapshot --selector $"buildkit_tag=($tag)" --output json
    } | complete
    let snap_id = if $snap_list.exit_code == 0 {
        $snap_list.stdout | from json | get -o 0 | get -o id | default "unknown"
    } else { "unknown" }
    log info $"Golden image ($snapshot_name) ready. id=($snap_id)"
}

def _wait_for_status [server_id: int, target_status: string, timeout_sec: int]: nothing -> nothing {
    let deadline = (date now) + ($timeout_sec * 1sec)
    loop {
        let r = do { ^hcloud server describe $server_id --output json } | complete
        if $r.exit_code == 0 {
            let status = $r.stdout | from json | get status
            if $status == $target_status { return }
        }
        if (date now) > $deadline {
            error make { msg: $"Timeout: server ($server_id) did not reach status '($target_status)'" }
        }
        sleep 10sec
    }
}

def _wait_for_ssh [host: string, timeout_sec: int, ssh_key: string = ""]: nothing -> nothing {
    let deadline = (date now) + ($timeout_sec * 1sec)
    let key_args = if ($ssh_key | is-not-empty) { ["-i" $ssh_key] } else { [] }
    loop {
        let r = do {
            ^ssh ...$key_args -o StrictHostKeyChecking=no -o ConnectTimeout=5 -o BatchMode=yes $"root@($host)" "
                if [ -f /var/lib/cloud/instance/boot-finished ]; then
                    grep -iq error /var/log/cloud-init.log && echo error || echo done
                else
                    echo waiting
                fi
            "
        } | complete
        if $r.exit_code == 0 {
            let out = ($r.stdout | str trim)
            if $out == "done" { return }
            if $out == "error" {
                error make { msg: $"cloud-init reported errors on ($host) — check /var/log/cloud-init.log" }
            }
        }
        if (date now) > $deadline {
            error make { msg: $"Timeout waiting for SSH + cloud-init on ($host)" }
        }
        sleep 15sec
    }
}

def _verify_tools [host: string, ssh_key: string = ""]: nothing -> nothing {
    let key_args = if ($ssh_key | is-not-empty) { ["-i" $ssh_key] } else { [] }
    for tool in $VERIFY_TOOLS {
        let r = do {
            ^ssh ...$key_args -o StrictHostKeyChecking=no -o BatchMode=yes $"root@($host)" $"command -v ($tool) && ($tool) --version 2>&1 | head -1"
        } | complete
        if $r.exit_code != 0 {
            error make { msg: $"Tool '($tool)' not found on runner — cloud-init may have failed.\nSSH output: ($r.stdout)\n($r.stderr)" }
        }
        log info $"  ($tool): ($r.stdout | str trim)"
    }
}

def _cleanup_server [server_id: int]: nothing -> nothing {
    log info $"Destroying seed server ($server_id)..."
    let r = do { ^hcloud server delete $server_id } | complete
    if $r.exit_code != 0 {
        log warning $"Server delete returned non-zero (server may need manual cleanup): ($r.stderr)"
    }
}