# NixOS Module for PXE Boot Server
#
# This module provides a complete PXE boot infrastructure for bare-metal
# provisioning of Centra Cloud nodes.
#
# Services provided:
# - DHCP server (ISC DHCP)
# - TFTP server (for iPXE bootloaders)
# - HTTP server (nginx, for iPXE scripts and NixOS images)
#
# Usage:
#   1. Import this module in your NixOS configuration
#   2. Enable and configure the PXE server
#   3. Deploy to your provisioning server
#
# Example:
#   imports = [ ./baremetal/pxe-server/nixos-module.nix ];
#
#   services.centra-pxe-server = {
#     enable = true;
#     interface = "eth0";
#     serverAddress = "10.0.100.10";
#     subnet = "10.0.100.0/24";
#     dhcpRange = {
#       start = "10.0.100.100";
#       end = "10.0.100.200";
#     };
#   };

{ config, lib, pkgs, ... }:

with lib;

let
  cfg = config.services.centra-pxe-server;

  # DHCP configuration file
  dhcpdConf = pkgs.writeText "dhcpd.conf" ''
    # ISC DHCP Server Configuration for PXE Boot
    # Auto-generated by NixOS module

    option space pxelinux;
    option architecture-type code 93 = unsigned integer 16;

    default-lease-time ${toString cfg.dhcp.defaultLeaseTime};
    max-lease-time ${toString cfg.dhcp.maxLeaseTime};

    authoritative;
    log-facility local7;

    subnet ${cfg.dhcp.subnet} netmask ${cfg.dhcp.netmask} {
        range ${cfg.dhcp.range.start} ${cfg.dhcp.range.end};

        option routers ${cfg.dhcp.router};
        option subnet-mask ${cfg.dhcp.netmask};
        option broadcast-address ${cfg.dhcp.broadcast};
        option domain-name-servers ${concatStringsSep ", " cfg.dhcp.nameservers};
        option domain-name "${cfg.dhcp.domainName}";

        next-server ${cfg.serverAddress};

        if exists user-class and option user-class = "iPXE" {
            filename "http://${cfg.serverAddress}/boot/ipxe/boot.ipxe";
        } elsif option architecture-type = 00:00 {
            filename "undionly.kpxe";
        } elsif option architecture-type = 00:06 {
            filename "ipxe-i386.efi";
        } elsif option architecture-type = 00:07 {
            filename "ipxe.efi";
        } elsif option architecture-type = 00:09 {
            filename "ipxe.efi";
        } else {
            filename "undionly.kpxe";
        }
    }

    ${cfg.dhcp.extraConfig}
  '';

  # iPXE boot script
  bootIpxeScript = pkgs.writeText "boot.ipxe" ''
    #!ipxe

    set boot-server ${cfg.serverAddress}
    set boot-url http://''${boot-server}/boot
    set nixos-url ''${boot-url}/nixos
    set provisioning-server http://''${boot-server}

    echo Network Configuration:
    echo IP Address: ''${ip}
    echo MAC Address: ''${mac}
    echo

    isset ''${profile} || set profile unknown

    ${concatStringsSep "\n" (mapAttrsToList (mac: node:
      "iseq ''${mac} ${mac} && set profile ${node.profile} && set hostname ${node.hostname} && goto boot ||"
    ) cfg.nodes)}

    goto menu

    :menu
    clear menu
    menu Centra Cloud - Bare-Metal Provisioning
    item --gap -- ------------------------- Boot Profiles -------------------------
    item control-plane    1. Control Plane Node (All Services)
    item worker           2. Worker Node (Compute Services)
    item all-in-one       3. All-in-One Node (Testing/Homelab)
    item --gap -- ------------------------- Advanced Options -------------------------
    item shell            iPXE Shell (for debugging)
    item reboot           Reboot System
    item exit             Exit to BIOS
    choose --timeout 30000 --default control-plane selected || goto cancel
    goto ''${selected}

    :cancel
    echo Boot cancelled, rebooting in 5 seconds...
    sleep 5
    reboot

    :control-plane
    set profile control-plane
    echo Booting: Control Plane Node
    goto boot

    :worker
    set profile worker
    echo Booting: Worker Node
    goto boot

    :all-in-one
    set profile all-in-one
    echo Booting: All-in-One Node
    goto boot

    :boot
    isset ''${hostname} || set hostname centra-node-''${mac:hexhyp}

    echo Profile:      ''${profile}
    echo Hostname:     ''${hostname}
    echo MAC Address:  ''${mac}

    set kernel-params initrd=initrd ip=dhcp
    set kernel-params ''${kernel-params} centra.profile=''${profile}
    set kernel-params ''${kernel-params} centra.hostname=''${hostname}
    set kernel-params ''${kernel-params} centra.mac=''${mac}
    set kernel-params ''${kernel-params} centra.provisioning-server=''${provisioning-server}
    set kernel-params ''${kernel-params} console=tty0 console=ttyS0,115200n8

    kernel ''${nixos-url}/bzImage ''${kernel-params} || goto failed
    initrd ''${nixos-url}/initrd || goto failed
    boot || goto failed

    :failed
    echo Boot Failed!
    echo Failed to load kernel or initrd from ''${nixos-url}
    goto shell

    :shell
    echo Entering iPXE shell...
    shell

    :reboot
    reboot

    :exit
    exit
  '';

  # Nginx configuration
  nginxConf = pkgs.writeText "nginx.conf" ''
    user nginx;
    worker_processes auto;
    error_log /var/log/nginx/error.log warn;
    pid /var/run/nginx.pid;

    events {
        worker_connections 1024;
    }

    http {
        include ${pkgs.nginx}/conf/mime.types;
        default_type application/octet-stream;

        log_format main '$remote_addr - $remote_user [$time_local] "$request" '
                        '$status $body_bytes_sent "$http_referer" '
                        '"$http_user_agent"';

        access_log /var/log/nginx/access.log main;

        sendfile on;
        tcp_nopush on;
        keepalive_timeout 65;

        types {
            application/octet-stream kpxe;
            application/octet-stream efi;
            text/plain ipxe;
        }

        server {
            listen ${toString cfg.http.port};
            server_name _;

            root ${cfg.bootAssetsPath};

            location /boot/ {
                alias ${cfg.bootAssetsPath}/;
                autoindex on;
                autoindex_exact_size off;
                autoindex_localtime on;
            }

            location ~ \.ipxe$ {
                add_header Cache-Control "no-store, no-cache, must-revalidate";
                expires -1;
            }

            location /health {
                access_log off;
                return 200 "OK\n";
                add_header Content-Type text/plain;
            }

            location = / {
                return 200 "Centra Cloud PXE Boot Server\n";
                add_header Content-Type text/plain;
            }
        }
    }
  '';

in {
  options.services.centra-pxe-server = {
    enable = mkEnableOption "Centra Cloud PXE Boot Server";

    interface = mkOption {
      type = types.str;
      default = "eth0";
      description = "Network interface to listen on for DHCP requests";
    };

    serverAddress = mkOption {
      type = types.str;
      example = "10.0.100.10";
      description = "IP address of the PXE boot server";
    };

    bootAssetsPath = mkOption {
      type = types.path;
      default = "/var/lib/pxe-boot";
      description = "Path to boot assets directory";
    };

    dhcp = {
      subnet = mkOption {
        type = types.str;
        example = "10.0.100.0";
        description = "Network subnet for DHCP";
      };

      netmask = mkOption {
        type = types.str;
        default = "255.255.255.0";
        description = "Network netmask";
      };

      broadcast = mkOption {
        type = types.str;
        example = "10.0.100.255";
        description = "Broadcast address";
      };

      range = {
        start = mkOption {
          type = types.str;
          example = "10.0.100.100";
          description = "Start of DHCP range";
        };

        end = mkOption {
          type = types.str;
          example = "10.0.100.200";
          description = "End of DHCP range";
        };
      };

      router = mkOption {
        type = types.str;
        example = "10.0.100.1";
        description = "Default gateway";
      };

      nameservers = mkOption {
        type = types.listOf types.str;
        default = [ "8.8.8.8" "8.8.4.4" ];
        description = "DNS nameservers";
      };

      domainName = mkOption {
        type = types.str;
        default = "centra.local";
        description = "Domain name";
      };

      defaultLeaseTime = mkOption {
        type = types.int;
        default = 600;
        description = "Default DHCP lease time in seconds";
      };

      maxLeaseTime = mkOption {
        type = types.int;
        default = 7200;
        description = "Maximum DHCP lease time in seconds";
      };

      extraConfig = mkOption {
        type = types.lines;
        default = "";
        description = "Additional DHCP configuration";
      };
    };

    http = {
      port = mkOption {
        type = types.int;
        default = 80;
        description = "HTTP server port";
      };
    };

    tftp = {
      enable = mkOption {
        type = types.bool;
        default = true;
        description = "Enable TFTP server for bootloader files";
      };
    };

    nodes = mkOption {
      type = types.attrsOf (types.submodule {
        options = {
          profile = mkOption {
            type = types.enum [ "control-plane" "worker" "all-in-one" ];
            description = "Node profile";
          };

          hostname = mkOption {
            type = types.str;
            description = "Node hostname";
          };

          ipAddress = mkOption {
            type = types.str;
            description = "Fixed IP address (optional)";
            default = "";
          };
        };
      });
      default = {};
      example = literalExpression ''
        {
          "52:54:00:12:34:56" = {
            profile = "control-plane";
            hostname = "control-plane-01";
            ipAddress = "10.0.100.50";
          };
        }
      '';
      description = "MAC address to node configuration mapping";
    };
  };

  config = mkIf cfg.enable {
    # DHCP Server
    services.dhcpd4 = {
      enable = true;
      interfaces = [ cfg.interface ];
      configFile = dhcpdConf;
    };

    # TFTP Server
    services.atftpd = mkIf cfg.tftp.enable {
      enable = true;
      root = "${cfg.bootAssetsPath}/ipxe";
    };

    # HTTP Server (Nginx)
    services.nginx = {
      enable = true;
      package = pkgs.nginx;
      appendHttpConfig = ''
        server {
          listen ${toString cfg.http.port};
          server_name _;
          root ${cfg.bootAssetsPath};

          location /boot/ {
            alias ${cfg.bootAssetsPath}/;
            autoindex on;
            autoindex_exact_size off;
            autoindex_localtime on;
          }

          location ~ \.ipxe$ {
            add_header Cache-Control "no-store, no-cache, must-revalidate";
            expires -1;
          }

          location /health {
            access_log off;
            return 200 "OK\n";
            add_header Content-Type text/plain;
          }
        }
      '';
    };

    # Firewall Rules
    networking.firewall = {
      allowedUDPPorts = [
        67  # DHCP server
        68  # DHCP client
        69  # TFTP
      ];
      allowedTCPPorts = [
        cfg.http.port  # HTTP
      ];
    };

    # Create boot assets directory structure
    systemd.tmpfiles.rules = [
      "d ${cfg.bootAssetsPath} 0755 nginx nginx -"
      "d ${cfg.bootAssetsPath}/ipxe 0755 nginx nginx -"
      "d ${cfg.bootAssetsPath}/nixos 0755 nginx nginx -"
      "L+ ${cfg.bootAssetsPath}/ipxe/boot.ipxe - - - - ${bootIpxeScript}"
    ];

    # Systemd service dependencies
    systemd.services.dhcpd4.after = [ "network-online.target" ];
    systemd.services.dhcpd4.wants = [ "network-online.target" ];

    systemd.services.atftpd.after = [ "network-online.target" ];
    systemd.services.atftpd.wants = [ "network-online.target" ];

    # Environment packages for management
    environment.systemPackages = with pkgs; [
      dhcp
      tftp-hpa
      curl
      wget
      ipxe
    ];
  };
}