#!/usr/bin/env python3

"""
Scaffold a new NixOS host and add it to flake.nix.
"""

from __future__ import annotations

import argparse
import re
import sys
from pathlib import Path
from textwrap import dedent



def _supports_color() -> bool:
    return sys.stderr.isatty()


_C = {
    "green": "\033[32m" if _supports_color() else "",
    "yellow": "\033[33m" if _supports_color() else "",
    "red": "\033[31m" if _supports_color() else "",
    "bold": "\033[1m" if _supports_color() else "",
    "reset": "\033[0m" if _supports_color() else "",
}


def log(msg: str) -> None:
    print(f"{_C['green']}[+]{_C['reset']} {msg}", file=sys.stderr)


def warn(msg: str) -> None:
    print(f"{_C['yellow']}[!]{_C['reset']} {msg}", file=sys.stderr)


def die(msg: str, code: int = 1) -> None:
    print(f"{_C['red']}[-]{_C['reset']} {msg}", file=sys.stderr)
    sys.exit(code)



# RFC 1123 hostname: 1-63 chars, [a-z0-9-], no leading/trailing dash.
HOST_RE = re.compile(r"^[a-z0-9]([a-z0-9-]{0,61}[a-z0-9])?$")


def valid_hostname(name: str) -> bool:
    return bool(HOST_RE.match(name))



DEFAULT_NIX_TEMPLATE = dedent("""\
    {{ pkgs, ... }}: {{
      imports = [
        ./hardware.nix
        ../../modules/core.nix
        ../../modules/plasma.nix
        ../../modules/security.nix
      ];

      networking.hostName = "{name}";

      boot.loader.systemd-boot.enable = true;
      boot.loader.efi.canTouchEfiVariables = true;
      boot.kernelPackages = pkgs.linuxPackages_latest;

      networking.networkmanager.enable = true;

      programs.zsh.enable = true;
      programs.nix-ld.enable = true;
      programs.nix-ld.libraries = with pkgs; [
        stdenv.cc.cc.lib
        zlib
      ];

      users.users.alex = {{
        isNormalUser = true;
        shell = pkgs.zsh;
        extraGroups = [ "wheel" "networkmanager" ];
      }};

      programs._1password.enable = true;
      programs._1password-gui = {
        enable = true;
        polkitPolicyOwners = [ "alex" ];
      };

      system.stateVersion = "26.05";
    }}
""")


PLACEHOLDER_HARDWARE = dedent("""\
    # PLACEHOLDER hardware.nix for host "{name}".
    #
    # Replace this file with the output of `nixos-generate-config --show-hardware-config`
    # run on the target machine. The justfile has a helper for this:
    #
    #     just capture-hardware {name}
    #
    # Until then, this file intentionally does nothing useful and the host
    # will not build.
    {{ lib, ... }}: {{
      # nixpkgs hostPlatform — change if not x86_64-linux.
      nixpkgs.hostPlatform = lib.mkDefault "x86_64-linux";

      # Intentionally empty. See note above.
    }}
""")




def add_to_flake(flake_path: Path, name: str, *, template_host: str | None = None) -> None:
    """
    Insert a new `<name> = nixpkgs.lib.nixosSystem { ... };` entry inside
    the `nixosConfigurations = { ... };` block of flake.nix.

    Strategy: locate the literal `nixosConfigurations = {` (allowing inline
    or next-line `{`), find the matching closing `}` by brace counting,
    then insert the new entry just before it. If `template_host` is set,
    clone that host's block; otherwise emit a sensible default.
    """
    src = flake_path.read_text()

    m = re.search(r"nixosConfigurations\s*=\s*\{", src)
    if not m:
        die(f"could not find `nixosConfigurations = {{` in {flake_path}")

    open_pos = m.end() - 1  # index of the '{'
    # Brace-match to find the closing '}' of this attrset.
    depth = 0
    close_pos = None
    in_string = False
    i = open_pos
    while i < len(src):
        c = src[i]
        # Cheap string skipping. flake.nix shouldn't have weird escapes inside
        # nixosConfigurations, but be polite about double-quoted strings.
        if c == '"' and (i == 0 or src[i - 1] != "\\"):
            in_string = not in_string
        elif not in_string:
            if c == "{":
                depth += 1
            elif c == "}":
                depth -= 1
                if depth == 0:
                    close_pos = i
                    break
        i += 1

    if close_pos is None:
        die("could not find the closing `}` of nixosConfigurations")

    block = src[open_pos + 1 : close_pos]

    # Already present? Bail.
    if re.search(rf"(^|\W){re.escape(name)}\s*=\s*nixpkgs\.lib\.nixosSystem", block):
        die(f"host `{name}` already exists in {flake_path}")

    # If we can clone an existing entry, do it (preserves user customizations).
    new_entry: str
    cloned = False
    if template_host:
        # Find `<template_host> = nixpkgs.lib.nixosSystem { ... };` inside block.
        tmpl_re = re.compile(
            rf"(^[ \t]*){re.escape(template_host)}\s*=\s*nixpkgs\.lib\.nixosSystem\s*\{{",
            re.MULTILINE,
        )
        tm = tmpl_re.search(block)
        if tm:
            tmpl_indent = tm.group(1)
            tmpl_start = tm.start()
            # Brace-match the template entry to find its terminating `};`.
            d = 0
            j = tm.end() - 1  # index of opening '{'
            in_str = False
            entry_end = None
            while j < len(block):
                cc = block[j]
                if cc == '"' and (j == 0 or block[j - 1] != "\\"):
                    in_str = not in_str
                elif not in_str:
                    if cc == "{":
                        d += 1
                    elif cc == "}":
                        d -= 1
                        if d == 0:
                            # Expect a trailing ';' after the closing '}'.
                            semi = block.find(";", j)
                            if semi != -1 and block[j + 1 : semi].strip() == "":
                                entry_end = semi + 1
                            else:
                                entry_end = j + 1
                            break
                j += 1
            if entry_end is None:
                die(f"could not parse template host `{template_host}` in flake.nix")

            entry_text = block[tmpl_start:entry_end]
            # Rename the binding and rewrite ./hosts/<template_host> paths.
            new_entry_body = re.sub(
                rf"^([ \t]*){re.escape(template_host)}(\s*=\s*nixpkgs\.lib\.nixosSystem)",
                rf"\1{name}\2",
                entry_text,
                count=1,
                flags=re.MULTILINE,
            )
            new_entry_body = new_entry_body.replace(
                f"./hosts/{template_host}", f"./hosts/{name}"
            )
            new_entry = "\n" + new_entry_body
            cloned = True

    if not cloned:
        # Best-effort default block. Indent two extra spaces if we can guess
        # the indentation by looking at existing content; otherwise 6 spaces.
        indent_match = re.search(r"^([ \t]+)\S", block, re.MULTILINE)
        indent = indent_match.group(1) if indent_match else "      "
        new_entry = (
            f"\n{indent}{name} = nixpkgs.lib.nixosSystem {{\n"
            f"{indent}  inherit system;\n"
            f"{indent}  specialArgs = {{ inherit inputs; }};\n"
            f"{indent}  modules = [\n"
            f"{indent}    ./hosts/{name}\n"
            f"{indent}    home-manager.nixosModules.home-manager\n"
            f"{indent}    {{\n"
            f"{indent}      home-manager = {{\n"
            f"{indent}        useGlobalPkgs = true;\n"
            f"{indent}        useUserPackages = true;\n"
            f"{indent}        users.alex = import ./home/alex;\n"
            f"{indent}        extraSpecialArgs = {{ inherit inputs; }};\n"
            f"{indent}        sharedModules = [ inputs.plasma-manager.homeModules.plasma-manager ];\n"
            f"{indent}      }};\n"
            f"{indent}    }}\n"
            f"{indent}  ];\n"
            f"{indent}}};\n"
        )

    # Insert new_entry just before the line containing the closing brace,
    # so we never produce `};};` on a single line.
    line_start = src.rfind("\n", 0, close_pos)
    if line_start == -1:
        line_start = 0
    if not new_entry.endswith("\n"):
        new_entry += "\n"
    new_src = src[: line_start + 1] + new_entry + src[line_start + 1 :]

    # Write atomically.
    tmp = flake_path.with_suffix(flake_path.suffix + ".tmp")
    tmp.write_text(new_src)
    tmp.replace(flake_path)


def host_in_flake(flake_path: Path, name: str) -> bool:
    """Cheap textual check, does flake.nix already define this host?"""
    src = flake_path.read_text()
    return bool(
        re.search(rf"(^|\W){re.escape(name)}\s*=\s*nixpkgs\.lib\.nixosSystem", src)
    )


def find_template_host(flake_path: Path) -> str | None:
    """Return the name of the first existing nixosConfigurations entry, if any."""
    src = flake_path.read_text()
    m = re.search(r"nixosConfigurations\s*=\s*\{", src)
    if not m:
        return None
    tail = src[m.end():]
    m2 = re.search(r"^[ \t]*([A-Za-z_][A-Za-z0-9_-]*)\s*=\s*nixpkgs\.lib\.nixosSystem",
                   tail, re.MULTILINE)
    return m2.group(1) if m2 else None



def main() -> int:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("name", help="hostname for the new host (lowercase, RFC 1123)")
    p.add_argument("--repo-root", default=".", help="path to the repo root (default: cwd)")
    p.add_argument("--dry-run", action="store_true", help="describe actions without writing")
    args = p.parse_args()

    if not valid_hostname(args.name):
        die(f"invalid hostname: {args.name!r} (must match {HOST_RE.pattern})")

    repo = Path(args.repo_root).resolve()
    flake = repo / "flake.nix"
    hosts_dir = repo / "hosts" / args.name

    if not flake.is_file():
        die(f"flake.nix not found at {flake}")
    if hosts_dir.exists():
        die(f"hosts/{args.name} already exists")
    if host_in_flake(flake, args.name):
        die(f"host `{args.name}` already defined in {flake} — refusing to touch anything")

    template_host = find_template_host(flake)
    log(f"repo:          {repo}")
    log(f"new host:      {args.name}")
    log(f"template host: {template_host or '<none, using default block>'}")

    if args.dry_run:
        warn("dry-run: nothing written")
        return 0

    # Scaffold filesystem.
    hosts_dir.mkdir(parents=True)
    (hosts_dir / "default.nix").write_text(DEFAULT_NIX_TEMPLATE.format(name=args.name))
    (hosts_dir / "hardware.nix").write_text(PLACEHOLDER_HARDWARE.format(name=args.name))
    log(f"created hosts/{args.name}/default.nix")
    log(f"created hosts/{args.name}/hardware.nix (placeholder)")

    # Edit flake.nix.
    add_to_flake(flake, args.name, template_host=template_host)
    log(f"added `{args.name}` entry to flake.nix")

    print()
    log(f"{_C['bold']}Next steps:{_C['reset']}")
    print(f"   1. On the new machine (booted from installer or already running):")
    print(f"        just capture-hardware {args.name}")
    print(f"      This overwrites hosts/{args.name}/hardware.nix with the real config.")
    print(f"   2. Review hosts/{args.name}/default.nix — tweak modules, users, etc.")
    print(f"   3. git add hosts/{args.name} flake.nix && git commit")
    print(f"   4. Build: just rebuild     (when running ON the new host)")
    return 0


if __name__ == "__main__":
    sys.exit(main())