#!/usr/bin/env bash
set -euo pipefail
shopt -s nullglob

REPO_ROOT="$(git rev-parse --show-toplevel)"
RULES_FILE="$REPO_ROOT/rules.nix"

if [[ ! -f "$RULES_FILE" ]]; then
  echo "ERROR: rules.nix not found at $RULES_FILE" >&2
  exit 1
fi

# --- Locate or clone nixpkgs ---
cleanup_nixpkgs=false
if [[ -z "${NIXPKGS_REPO:-}" ]]; then
  NIXPKGS_REPO=$(mktemp -d)
  cleanup_nixpkgs=true
  echo "Cloning nixpkgs (partial clone, may take a moment)..."
  git clone --filter=blob:none --single-branch --branch nixpkgs-unstable \
    https://github.com/NixOS/nixpkgs.git "$NIXPKGS_REPO" --quiet
fi

cleanup() {
  if "$cleanup_nixpkgs"; then
    rm -rf "$NIXPKGS_REPO"
  fi
}
trap cleanup EXIT

# Determine the ref to walk.  For a fresh clone we use HEAD (which is
# nixpkgs-unstable).  For a user-supplied repo we prefer the remote tracking
# branch so the result is independent of which branch is checked out locally.
if git -C "$NIXPKGS_REPO" rev-parse --verify origin/nixpkgs-unstable >/dev/null 2>&1; then
  NIXPKGS_REF="origin/nixpkgs-unstable"
else
  NIXPKGS_REF="HEAD"
fi

# --- Map package name to the nixpkgs file whose version we track ---
nixpkgs_path_for() {
  case "$1" in
    matrix-synapse)
      echo "pkgs/by-name/ma/matrix-synapse-unwrapped/package.nix"
      ;;
    *)
      echo "ERROR: no nixpkgs path mapping for package '$1'" >&2
      return 1
      ;;
  esac
}

# --- Extract version string from package.nix at a given commit ---
version_at_commit() {
  local commit="$1" pkg_path="$2"
  git -C "$NIXPKGS_REPO" show "${commit}:${pkg_path}" 2>/dev/null \
    | grep -m1 'version = ' \
    | sed 's/.*"\(.*\)".*/\1/'
}

# --- Generate flake.nix for a matrix-synapse pin ---
generate_flake_nix() {
  local pkg="$1" version="$2" commit="$3" outfile="$4"
  cat > "$outfile" << FLAKE_EOF
{
  description = "${pkg} ${version}";

  inputs.nixpkgs.url = "github:NixOS/nixpkgs/${commit}";

  outputs =
    { nixpkgs, ... }:
    let
      forAllSystems = nixpkgs.lib.genAttrs [
        "x86_64-linux"
        "aarch64-linux"
      ];
    in
    {
      packages = forAllSystems (system: {
        default = nixpkgs.legacyPackages.\${system}.matrix-synapse;
        matrix-synapse = nixpkgs.legacyPackages.\${system}.matrix-synapse;
        matrix-synapse-unwrapped = nixpkgs.legacyPackages.\${system}.matrix-synapse-unwrapped;
      });

      overlays.default = _final: prev:
        let
          pinned = nixpkgs.legacyPackages.\${prev.stdenv.hostPlatform.system};
        in
        {
          matrix-synapse = pinned.matrix-synapse;
          matrix-synapse-unwrapped = pinned.matrix-synapse-unwrapped;
        };
    };
}
FLAKE_EOF
}

# --- Read rules ---
RULES_JSON=$(nix eval --json --file "$RULES_FILE")
mapfile -t packages < <(echo "$RULES_JSON" | jq -r 'keys[]')

for pkg in "${packages[@]}"; do
  latest_n=$(echo "$RULES_JSON" | jq -r ".[\"$pkg\"].\"latest-revisions\"")
  mapfile -t pins < <(echo "$RULES_JSON" | jq -r ".[\"$pkg\"].pin // [] | .[]")
  pkg_path=$(nixpkgs_path_for "$pkg")

  echo "=== $pkg: latest $latest_n + ${#pins[@]} pin(s) ==="

  # Collect commits that touched the package file (newest first, first-parent only)
  mapfile -t commits < <(
    git -C "$NIXPKGS_REPO" log --first-parent --format="%H" "$NIXPKGS_REF" -- "$pkg_path" \
      | head -300
  )
  echo "  Scanning ${#commits[@]} commits..."

  # Build version -> commit map.  First occurrence (newest) wins.
  declare -A version_commit=()
  declare -a version_order=()

  for commit in "${commits[@]}"; do
    ver=$(version_at_commit "$commit" "$pkg_path") || true
    [[ -n "${ver:-}" ]] || continue
    if [[ -z "${version_commit[$ver]:-}" ]]; then
      version_commit["$ver"]="$commit"
      version_order+=("$ver")
    fi
  done

  echo "  Found ${#version_order[@]} unique versions"

  # Decide which versions to keep
  declare -A keep_versions=()

  for ((i = 0; i < latest_n && i < ${#version_order[@]}; i++)); do
    ver="${version_order[$i]}"
    keep_versions["$ver"]=1
    echo "  [latest]  $ver  ${version_commit[$ver]:0:12}"
  done

  for pin_ver in "${pins[@]}"; do
    if [[ -n "${version_commit[$pin_ver]:-}" ]]; then
      keep_versions["$pin_ver"]=1
      echo "  [pinned]  $pin_ver  ${version_commit[$pin_ver]:0:12}"
    else
      echo "  WARNING: pinned version $pin_ver not found in nixpkgs history" >&2
    fi
  done

  # Create or update pin directories
  for ver in "${!keep_versions[@]}"; do
    dir="$REPO_ROOT/${pkg}@${ver}"
    commit="${version_commit[$ver]}"
    flake_file="$dir/flake.nix"

    if [[ -f "$flake_file" ]]; then
      existing_commit=$(sed -n 's|.*github:NixOS/nixpkgs/\([a-f0-9]*\).*|\1|p' "$flake_file" | head -1)
      if [[ "$existing_commit" == "$commit" ]]; then
        echo "  $ver: up to date"
        continue
      fi
      echo "  $ver: updating commit"
    else
      echo "  $ver: creating"
    fi

    mkdir -p "$dir"
    generate_flake_nix "$pkg" "$ver" "$commit" "$flake_file"
    # flake.nix must be tracked by git before nix flake lock can see it
    git add "$flake_file"
    echo "  $ver: locking flake inputs..."
    nix flake lock "$dir"
    git add "$dir/flake.lock"
  done

  # Remove directories no longer required by rules
  for dir in "$REPO_ROOT/${pkg}@"*/; do
    [[ -d "$dir" ]] || continue
    dir_name=$(basename "$dir")
    ver="${dir_name#"${pkg}@"}"
    if [[ -z "${keep_versions[$ver]:-}" ]]; then
      echo "  Removing: $ver"
      rm -rf "$dir"
    fi
  done

  unset version_commit version_order keep_versions
done

echo "Done."