#!/usr/bin/env nu
# NixOS Flake Generator for Workspace Servers
# Generates flake.nix, configuration.nix, and hardware-configuration.nix from Nickel workspace configs
#
# Usage:
#   nu provisioning/scripts/generate-flakes.nu workspaces/librecloud_hetzner
#   nu provisioning/scripts/generate-flakes.nu --dry-run workspaces/workspace_librecloud
#
# Features:
#   - Reads taskservs from servers.ncl
#   - Generates complete flake.nix WITH dynamic taskserv inputs
#   - Generates configuration.nix WITH taskserv NixOS modules
#   - Phase separation: export → map → generate → write
#   - Result pattern {ok, err} (NO try-catch)
#   - Relative path normalization for flake inputs

use std log

# Helper: Search up directory tree for provisioning root
def search-up-for-provisioning [dir: string] {
    if ($"($dir)/provisioning/schemas" | path exists) {
        return $dir
    }

    let parent = ($dir | path dirname)
    if $parent == $dir {
        return ""
    }

    search-up-for-provisioning $parent
}

# Helper: Detect provisioning project root
def get-provisioning-root [] {
    # 1. Check PROVISIONING environment variable
    if ($env.PROVISIONING? != null) {
        return $env.PROVISIONING
    }

    # 2. Check if we're in provisioning/ directory or above
    let cwd = (pwd)

    if ($"($cwd)/provisioning/schemas" | path exists) {
        return $cwd
    }

    # 3. Search up the directory tree
    let found = (search-up-for-provisioning $cwd)
    if $found != "" {
        return $found
    }

    # 4. Fallback to current directory
    $cwd
}

# Helper: Result type - success
def ok [value: any] {
    {ok: $value, err: null}
}

# Helper: Result type - error
def err [message: string] {
    {ok: null, err: $message}
}

# Helper: Check if result is ok
def is-ok [result: record] {
    $result.err == null
}

# Helper: Extract value from ok result
def unwrap-ok [result: record] {
    $result.ok
}

# Helper: Extract error from err result
def unwrap-err [result: record] {
    $result.err
}

# Phase 0: Map taskserv names to extension paths
def map-taskserv-to-path [taskserv_name: string] {
    match $taskserv_name {
        "etcd" => "cluster/etcd",
        "kubernetes" => "cluster/kubernetes",
        "coredns" => "cluster/coredns",
        "k8s_nodejoin" => "cluster/k8s_nodejoin",
        "containerd" => "container_runtime/containerd",
        "crio" => "container_runtime/crio",
        "podman" => "container_runtime/podman",
        "crun" => "container_runtime/crun",
        "youki" => "container_runtime/youki",
        "runc" => "container_runtime/runc",
        "cilium" => "networking/cilium",
        "coredns_dns" => "networking/coredns",
        "resolv" => "networking/resolv",
        "proxy" => "networking/proxy",
        "rook_ceph" => "storage/rook_ceph",
        "external_nfs" => "storage/external_nfs",
        "postgres" => "databases/postgres",
        "redis" => "databases/redis",
        "os" => "infrastructure/os",
        "webhook" => "infrastructure/webhook",
        "provisioning" => "infrastructure/provisioning",
        "kubectl" => "infrastructure/kubectl",
        "prometheus" => "development/prometheus",
        "grafana" => "development/grafana",
        "loki" => "development/loki",
        "gitea" => "development/gitea",
        "oras" => "development/oras",
        _ => null,
    }
}

# Phase 1: Export servers.ncl via Nickel
def export-servers-config [workspace_path: string, provisioning_root: string] {
    # Guard: Workspace directory exists
    if not ($workspace_path | path exists) {
        return (err $"Workspace directory not found: ($workspace_path)")
    }

    # Guard: Find servers.ncl (could be in infra/, infra/default/, infra/main/, infra/region/)
    let servers_file = (
        if ($"($workspace_path)/infra/main/servers.ncl" | path exists) {
            $"($workspace_path)/infra/main/servers.ncl"
        } else if ($"($workspace_path)/infra/servers.ncl" | path exists) {
            $"($workspace_path)/infra/servers.ncl"
        } else if ($"($workspace_path)/infra/default/servers.ncl" | path exists) {
            $"($workspace_path)/infra/default/servers.ncl"
        } else {
            ""
        }
    )

    if ($servers_file | is-empty) {
        return (err $"servers.ncl not found in workspace: ($workspace_path)\nLooking for: ($workspace_path)/infra/main/servers.ncl, ($workspace_path)/infra/servers.ncl, or ($workspace_path)/infra/*/servers.ncl")
    }

    # Export servers.ncl to JSON via nickel
    let export_result = (
        do {
            with-env { NICKEL_IMPORT_PATH: $provisioning_root } {
                ^nickel export --format json $servers_file
            }
        } | complete
    )

    # Guard: nickel export succeeded
    if $export_result.exit_code != 0 {
        let error_msg = (
            if ($export_result.stderr | str contains "command not found") {
                "nickel command not found - install Nickel to use this script"
            } else {
                $"nickel export failed: ($export_result.stderr)"
            }
        )
        return (err $error_msg)
    }

    # Parse JSON output
    let servers_json = (
        try {
            $export_result.stdout | from json
        } catch {
            return (err "Failed to parse nickel export output as JSON")
        }
    )

    # Guard: JSON has servers field
    let has_servers = (($servers_json | get servers?) != null)
    if not $has_servers {
        return (err "Exported config missing 'servers' field")
    }

    ok $servers_json
}

# Phase 2: Extract enabled taskservs from config record
def extract-enabled-taskservs [taskservs_record: any] {
    if ($taskservs_record == null) {
        return (ok [])
    }

    if not (($taskservs_record | describe) =~ "record") {
        return (err $"taskservs must be a record, got (($taskservs_record | describe))")
    }

    let enabled = (
        $taskservs_record
        | items { |k, v| {key: $k, value: $v} }
        | where { |kv| $kv.value.enable? == true or $kv.value.enabled? == true }
        | get key
        | sort
    )

    ok $enabled
}

# Helper: Build single flake input record
def build-flake-input [name: string] {
    let path = (map-taskserv-to-path $name)
    if ($path == null) {
        {
            input_name: $"taskserv-($name)",
            input_line: $'    # WARNING: taskserv ($name) has unknown mapping - skipped',
            skip: true
        }
    } else {
        let relative_path = $"../../../extensions/taskservs/($path)"
        {
            input_name: $"taskserv-($name)",
            input_line: $'    taskserv-($name).url = "path:($relative_path)";',
            skip: false
        }
    }
}

# Phase 2b: Generate flake inputs for taskservs
def generate-flake-inputs [enabled_taskservs: list, provisioning_root: string] {
    let inputs = (
        $enabled_taskservs
        | each { |name| (build-flake-input $name) }
    )

    let grouped = ($inputs | group-by { |x| $x.input_name })
    let unique_inputs = (
        $grouped
        | items { |k, v| $v | get 0 }
    )

    ok $unique_inputs
}

# Phase 2c: Generate flake.nix for a server
def generate-flake-nix [hostname: string, enabled_taskservs: list, system: string, flake_inputs: any] {
    let taskserv_input_names = (
        $flake_inputs
        | where { |inp| $inp.skip == false }
        | each { |inp| $inp.input_name }
    )

    let input_lines = (
        $flake_inputs
        | each { |inp| $inp.input_line }
    )

    let outputs_params = (
        ["self" "nixpkgs" "flake-utils"]
        | append $taskserv_input_names
        | str join ", "
    )

    let module_imports = (
        $flake_inputs
        | where { |inp| $inp.skip == false }
        | each { |inp| $"          inputs.($inp.input_name).nixosModules.default" }
    )

    let part1 = [
        "{"
        $'  description = "NixOS Flake for ($hostname)";'
        ""
        "  inputs = {"
        '    nixpkgs.url = "github:NixOS/nixpkgs/nixos-24.11";'
        '    flake-utils.url = "github:numtide/flake-utils";'
    ]

    let part2 = if ($input_lines | length) > 0 { $input_lines } else { [] }

    let part3 = [
        "  };"
        ""
        $"  outputs = { ($outputs_params) }@inputs:"
        "    flake-utils.lib.eachDefaultSystem (system: {"
        "      devShells.default = import ./shell.nix { inherit pkgs; };"
        "    }) // {"
        $"      nixosConfigurations.($hostname) = nixpkgs.lib.nixosSystem {"
        $'        system = "($system)";'
        "        specialArgs = { inherit inputs; };"
        "        modules = ["
        "          ./configuration.nix"
        "          ./hardware-configuration.nix"
    ]

    let part4 = if ($module_imports | length) > 0 { $module_imports } else { [] }

    let part5 = [
        "        ];"
        "      };"
        "    };"
        "}"
    ]

    ($part1 | append $part2 | append $part3 | append $part4 | append $part5) | str join "\n"
}

# Phase 2d: Generate configuration.nix with taskserv modules
def generate-configuration-nix [hostname: string, taskservs_record: any, private_ip: any, system: string] {
    let taskserv_config_lines = (
        if ($taskservs_record == null) {
            []
        } else {
            $taskservs_record
            | items { |k, v| {key: $k, value: $v} }
            | where { |kv| $kv.value.enable? == true or $kv.value.enabled? == true }
            | each { |kv|
                let name = $kv.key
                let config = $kv.value
                let enabled_str = "    enable = true;"

                # Build taskserv-specific configuration
                let ts_config = (
                    match $name {
                        "etcd" => [
                            "  provisioning.taskservs.etcd = {"
                            $enabled_str
                            '    listen_client_urls = "http://0.0.0.0:2379";'
                            $'    advertise_client_urls = "http://($hostname).internal:2379";'
                            '    listen_peer_urls = "http://0.0.0.0:2380";'
                            $'    initial_advertise_peer_urls = "http://($hostname).internal:2380";'
                            $'    initial_cluster = "($hostname)=http://($hostname).internal:2380";'
                            '    cluster_token = "provisioned-cluster";'
                            "  };"
                        ],
                        "kubernetes" => {
                            let role = ($config.role? | default "worker")
                            [
                                "  provisioning.taskservs.kubernetes = {"
                                $enabled_str
                                $'    role = "($role)";'
                                (if ($config.role? == "control-plane") { '    apiServerAdvertiseAddress = "";' } else { "" })
                                "  };"
                            ]
                            | where { |x| $x != "" }
                        },
                        "cilium" => [
                            "  provisioning.taskservs.cilium = {"
                            $enabled_str
                            '    ipam = "kubernetes";'
                            "  };"
                        ],
                        "rook_ceph" => [
                            "  provisioning.taskservs.rook_ceph = {"
                            $enabled_str
                            $'    cluster_name = "($config.cluster_name? | default "ceph")";'
                            "  };"
                        ],
                        "containerd" => [
                            "  provisioning.taskservs.containerd = {"
                            $enabled_str
                            "  };"
                        ],
                        "coredns" => [
                            "  provisioning.taskservs.coredns = {"
                            $enabled_str
                            "  };"
                        ],
                        "resolv" => [
                            "  provisioning.taskservs.resolv = {"
                            $enabled_str
                            "  };"
                        ],
                        "prometheus" => [
                            "  provisioning.taskservs.prometheus = {"
                            $enabled_str
                            $'    retention_days = ($config.retention_days? | default 30);'
                            "  };"
                        ],
                        "grafana" => [
                            "  provisioning.taskservs.grafana = {"
                            $enabled_str
                            "  };"
                        ],
                        "loki" => [
                            "  provisioning.taskservs.loki = {"
                            $enabled_str
                            "  };"
                        ],
                        _ => [
                            $"  provisioning.taskservs.($name) = {"
                            $enabled_str
                            "  };"
                        ]
                    }
                )

                $ts_config | str join "\n"
            }
        }
    )

    let firewall_ports = [
        "  # Firewall ports for enabled services"
        "  networking.firewall.allowedTCPPorts = ["
        "    22       # SSH"
        "    80       # HTTP"
        "    443      # HTTPS"
    ] | append (
        if ($taskservs_record.etcd?.enable? == true or $taskservs_record.etcd?.enabled? == true) {
            ["    2379     # etcd client", "    2380     # etcd peer"]
        } else {
            []
        }
    ) | append (
        if ($taskservs_record.kubernetes?.enable? == true or $taskservs_record.kubernetes?.enabled? == true) {
            ["    6443     # Kubernetes API"]
        } else {
            []
        }
    ) | append (
        if ($taskservs_record.prometheus?.enable? == true or $taskservs_record.prometheus?.enabled? == true) {
            ["    9090     # Prometheus"]
        } else {
            []
        }
    ) | append (
        if ($taskservs_record.grafana?.enable? == true or $taskservs_record.grafana?.enabled? == true) {
            ["    3000     # Grafana"]
        } else {
            []
        }
    ) | append ["  ];"]

    let base = [
        "{ config, pkgs, inputs, ... }:"
        ""
        "{"
        "  # System hostname"
        $'  networking.hostName = "($hostname)";'
        '  networking.domain = "internal";'
        ""
        "  # Time and locale"
        '  time.timeZone = "UTC";'
        ""
        "  # Nix settings"
        "  nix.settings = {"
        "    auto-optimise-store = true;"
        '    trusted-users = [ "root" "nixos" ];'
        "    experimental-features = [ \"flakes\" \"nix-command\" ];"
        "  };"
        ""
        "  # SSH"
        "  services.openssh = {"
        "    enable = true;"
        "    settings = {"
        "      PasswordAuthentication = false;"
        "      PubkeyAuthentication = true;"
        "      PermitRootLogin = \"prohibit-password\";"
        "    };"
        "  };"
        ""
        "  # System packages"
        "  environment.systemPackages = with pkgs; ["
        "    curl"
        "    wget"
        "    git"
        "    vim"
        "    htop"
        "    jq"
        "    nushell"
        "    nix-output-monitor"
        "  ];"
        ""
        "  # Kernel parameters for Kubernetes"
        "  boot.kernel.sysctl = {"
        '    "net.ipv4.ip_forward" = 1;'
        '    "net.ipv6.conf.all.forwarding" = 1;'
        '    "net.bridge.bridge-nf-call-iptables" = 1;'
        '    "net.bridge.bridge-nf-call-ip6tables" = 1;'
        "  };"
        ""
        "  # Required kernel modules"
        "  boot.kernelModules = [ \"overlay\" \"br_netfilter\" ];"
        ""
    ]

    let with_taskservs = (
        if ($taskserv_config_lines | length) > 0 {
            $base | append ["  # Taskserv configurations"] | append $taskserv_config_lines | append [""]
        } else {
            $base
        }
    )

    let with_firewall = (
        $with_taskservs
        | append $firewall_ports
        | append [""]
        | append ['  # System state version - do not change']
        | append ['  system.stateVersion = "24.11";']
        | append ["}"]
    )

    $with_firewall | str join "\n"
}

# Phase 2e: Generate hardware-configuration.nix with system-specific settings
def generate-hardware-configuration-nix [hostname: string, system: string, server_type: any] {
    let is_aarch64 = ($system =~ "aarch64")
    let kernel_module = (if $is_aarch64 { "kvm-arm" } else { "kvm-intel" })

    let base = [
        $"# Hardware configuration for $hostname"
        "# Generated by provisioning system - customize as needed"
        "# System: $system"
        "{ config, lib, pkgs, modulesPath, ... }:"
        ""
        "{"
        "  imports = ["
        '    (modulesPath + "/profiles/qemu-guest.nix")'
        "  ];"
        ""
        "  boot = {"
        "    loader.grub = {"
        "      enable = true;"
        '      device = "/dev/sda";'
        "    };"
        "    initrd.availableKernelModules = ["
        '      "ata_piix"'
        '      "uhci_pci"'
        '      "virtio_pci"'
        '      "sr_mod"'
        '      "virtio_blk"'
        "    ];"
        $'    kernelModules = [ "($kernel_module)" ];'
        "  };"
        ""
        '  fileSystems."/" = {'
        '    device = "/dev/sda1";'
        '    fsType = "ext4";'
        "  };"
        ""
        "  swapDevices = ["
        '    { device = "/dev/sda2"; }'
        "  ];"
        ""
        "  networking.usePredictableInterfaceNames = true;"
        ""
        "  nix.settings.max-jobs = lib.mkDefault 2;"
        "  nix.settings.cores = lib.mkDefault 2;"
    ]

    let with_closing = ($base | append ["}"])

    $with_closing | str join "\n"
}

# Phase 3: Write files to filesystem
def write-flake-files [
    output_dir: string,
    hostname: string,
    flake_nix: string,
    configuration_nix: string,
    hardware_configuration_nix: string,
    --dry-run
] {
    # Guard: Output directory exists
    if not ($output_dir | path exists) {
        if $dry_run {
            log info $"[DRY-RUN] Would create directory: ($output_dir)"
        } else {
            mkdir $output_dir
            if not ($output_dir | path exists) {
                return (err $"Failed to create directory: ($output_dir)")
            }
        }
    }

    # Write flake.nix
    let flake_path = $"($output_dir)/flake.nix"
    if $dry_run {
        log info $"[DRY-RUN] Would write: ($flake_path)"
    } else {
        $flake_nix | save --force $flake_path
        if not ($flake_path | path exists) {
            return (err $"Failed to write: ($flake_path)")
        }
    }

    # Write configuration.nix
    let config_path = $"($output_dir)/configuration.nix"
    if $dry_run {
        log info $"[DRY-RUN] Would write: ($config_path)"
    } else {
        $configuration_nix | save --force $config_path
        if not ($config_path | path exists) {
            return (err $"Failed to write: ($config_path)")
        }
    }

    # Write hardware-configuration.nix
    let hardware_path = $"($output_dir)/hardware-configuration.nix"
    if $dry_run {
        log info $"[DRY-RUN] Would write: ($hardware_path)"
    } else {
        $hardware_configuration_nix | save --force $hardware_path
        if not ($hardware_path | path exists) {
            return (err $"Failed to write: ($hardware_path)")
        }
    }

    ok true
}

# Main: Process a single server
def process-server [
    server: record,
    workspace_path: string,
    provisioning_root: string,
    --dry-run
] {
    # Guard: hostname field exists
    let hostname = $server.hostname? | default ""
    if ($hostname | is-empty) {
        return (err "Server missing hostname field")
    }

    # Guard: os_type must be nixos
    let os_type = $server.os_type? | default ""
    if ($os_type != "nixos") {
        log debug $"Skipping ($hostname) - os_type is ($os_type), not nixos"
        return (ok true)
    }

    # Guard: enabled field exists and is true (optional check)
    let enabled = $server.enabled? | default true
    if not $enabled {
        log debug $"Skipping ($hostname) - enabled is false"
        return (ok true)
    }

    # Determine output directory from nixos.flake_path or generate standard path
    let output_dir = (
        if not ($server.nixos.flake_path? | is-empty) {
            $server.nixos.flake_path
        } else {
            let workspace_name = ($workspace_path | path basename)
            $"($workspace_path)/../nixos/($hostname)"
        }
    )

    # Get system architecture
    let system = ($server.nixos.system? | default "aarch64-linux")

    log info $"Processing ($hostname)..."

    # Phase 2: Extract enabled taskservs
    log info "  Extracting taskservs..."
    let taskservs_result = (extract-enabled-taskservs $server.taskservs?)
    if not (is-ok $taskservs_result) {
        return (err $"Failed to extract taskservs: (unwrap-err $taskservs_result)")
    }

    let enabled_taskservs = (unwrap-ok $taskservs_result)
    let ts_list = ($enabled_taskservs | str join ", ")
    log info $"  Enabled taskservs: ($ts_list)"

    # Phase 2b: Generate flake inputs
    log info "  Generating flake inputs..."
    let flake_inputs_result = (generate-flake-inputs $enabled_taskservs $provisioning_root)
    if not (is-ok $flake_inputs_result) {
        return (err $"Failed to generate flake inputs: (unwrap-err $flake_inputs_result)")
    }

    let flake_inputs = (unwrap-ok $flake_inputs_result)

    # Phase 2c: Generate Nix files
    let flake_nix = (generate-flake-nix $hostname $enabled_taskservs $system $flake_inputs)
    let input_count = ($flake_inputs | length)
    log info $"  Generated flake.nix with ($input_count) inputs"

    let configuration_nix = (generate-configuration-nix $hostname $server.taskservs? $server.networking.private_ip? $system)
    log info "  Generated configuration.nix"

    let hardware_configuration_nix = (generate-hardware-configuration-nix $hostname $system $server.server_type?)
    log info "  Generated hardware-configuration.nix"

    # Phase 3: Write files
    let write_result = (
        write-flake-files $output_dir $hostname $flake_nix $configuration_nix $hardware_configuration_nix --dry-run=$dry_run
    )
    let write_ok = (is-ok $write_result)
    log info $"  Write result ok: ($write_ok)"

    if (is-ok $write_result) {
        let ts_count = ($enabled_taskservs | length)
        log info $"✅ Generated flakes for ($hostname) with ($ts_count) taskservs"
        ok true
    } else {
        log error $"Failed to generate flakes for ($hostname): (unwrap-err $write_result)"
        err (unwrap-err $write_result)
    }
}

# Main entry point
def main [
    workspace_path?: string,  # Path to workspace (e.g., workspaces/librecloud_hetzner)
    --dry-run,  # Dry-run mode (show what would be done)
] {
    let dry_run = $dry_run
    # Detect provisioning root
    let prov_root = (get-provisioning-root)
    log debug $"Provisioning root: ($prov_root)"

    # Guard: workspace_path provided
    if ($workspace_path | is-empty) {
        log error "Usage: generate-flakes.nu <workspace_path> [--dry-run]"
        log error "Example: generate-flakes.nu workspaces/librecloud_hetzner"
        return 1
    }

    # Resolve workspace path (relative or absolute)
    # Workspaces are at repo root/workspaces, not provisioning/workspaces
    let ws_path = (
        if ($workspace_path | str starts-with "/") {
            $workspace_path
        } else {
            # Check if workspace exists relative to repo root
            let repo_root = ($prov_root | path dirname)
            let candidate = $"($repo_root)/($workspace_path)"
            if ($candidate | path exists) {
                $candidate
            } else {
                # Fallback to provisioning root
                $"($prov_root)/($workspace_path)"
            }
        }
    )

    log info "═══════════════════════════════════════════════════════════════════"
    log info "NixOS Flake Generator"
    log info "═══════════════════════════════════════════════════════════════════"
    log info $"Workspace: ($ws_path)"
    log info $"Dry-run: ($dry_run)"

    # Phase 1: Export servers config from Nickel
    log info ""
    log info "Phase 1: Exporting servers.ncl from Nickel..."
    let export_result = (export-servers-config $ws_path $prov_root)

    if not (is-ok $export_result) {
        log error $"Export failed: (unwrap-err $export_result)"
        return 1
    }

    let servers_config = (unwrap-ok $export_result)
    let servers = $servers_config.servers
    let server_count = ($servers | length)

    log info $"Exported ($server_count) servers from Nickel"

    # Phase 2-3: Process each server
    log info ""
    log info "Phase 2-3: Generating flakes for each server..."
    mut generation_errors = []

    for server in $servers {
        let result = (
            if $dry_run {
                process-server $server $ws_path $prov_root --dry-run
            } else {
                process-server $server $ws_path $prov_root
            }
        )
        if not (is-ok $result) {
            $generation_errors | append (unwrap-err $result) | let generation_errors
        }
    }

    # Final report
    log info ""
    log info "═══════════════════════════════════════════════════════════════════"

    let error_count = ($generation_errors | length)
    if $error_count == 0 {
        log info "All flakes generated successfully!"
        return 0
    } else {
        log error $"Failed to generate ($error_count) flakes"
        $generation_errors | each { |e| log error $"   • ($e)" }
        return 1
    }
}

# Script is executed as:
#   nu generate-flakes.nu <workspace_path> [--dry-run]
# Nushell automatically passes CLI arguments to main