#!/usr/bin/env ruby
# Checks that Gemfile.lock and Gemfile.saas.lock are in sync for shared dependencies.
# Since Gemfile.saas evals Gemfile, shared gems should have identical versions.
#
# Usage:
#   bin/bundle-drift [check]     # check for drift (default subcommand)
#   bin/bundle-drift correct     # restore alignment (Gemfile.saas.lock is authoritative)
#   bin/bundle-drift forward     # push Gemfile.lock changes into Gemfile.saas.lock
#   bin/bundle-drift self-test   # run built-in tests
require "bundler"
require "fileutils"

GEMFILE_LOCK = "Gemfile.lock"
GEMFILE_SAAS_LOCK = "Gemfile.saas.lock"

class GemfileDriftChecker
  def initialize
    @oss_lockfile = parse_lockfile(GEMFILE_LOCK)
    @saas_lockfile = parse_lockfile(GEMFILE_SAAS_LOCK)
  end

  def check
    find_drift.tap do
      report it
    end
  end

  private
    def parse_lockfile(path)
      Bundler::LockfileParser.new(File.read(path))
    end

    def find_drift
      oss_specs, saas_specs = specs_hash(@oss_lockfile), specs_hash(@saas_lockfile)
      oss_deps, saas_deps = deps_hash(@oss_lockfile), deps_hash(@saas_lockfile)
      oss_sdeps, saas_sdeps = spec_deps_hash(@oss_lockfile), spec_deps_hash(@saas_lockfile)
      shared_gems = oss_specs.keys & saas_specs.keys
      shared_deps = oss_deps.keys & saas_deps.keys

      shared_gems.filter_map do |name|
        oss_version, saas_version = oss_specs[name], saas_specs[name]
        version_drift = oss_version != saas_version

        oss_dep, saas_dep = oss_deps[name], saas_deps[name]
        dep_drift = shared_deps.include?(name) && oss_dep != saas_dep

        oss_sd, saas_sd = oss_sdeps[name] || {}, saas_sdeps[name] || {}
        subdep_changes = (oss_sd.keys | saas_sd.keys).filter_map do |dep_name|
          if oss_sd[dep_name] != saas_sd[dep_name]
            { name: dep_name, oss: oss_sd[dep_name], saas: saas_sd[dep_name] }
          end
        end

        if version_drift || dep_drift || subdep_changes.any?
          { name: name, oss: oss_version, saas: saas_version, oss_dep: oss_dep, saas_dep: saas_dep, subdep_changes: subdep_changes }
        end
      end.sort_by { |d| d[:name] }
    end

    def specs_hash(lockfile)
      lockfile.specs.to_h { |spec| [ spec.name, spec.version.to_s ] }
    end

    def deps_hash(lockfile)
      lockfile.dependencies.transform_values { |dep| dep.requirement.to_s }
    end

    def spec_deps_hash(lockfile)
      lockfile.specs.to_h do |spec|
        [spec.name, spec.dependencies.map { |d| [d.name, d.requirement.to_s] }.sort.to_h]
      end
    end

    def report(drift)
      if drift.empty?
        puts "✓ Gemfile.lock and Gemfile.saas.lock are in sync"
      else
        puts "✗ Gemfile lock files have drifted!\n\n"

        spec_drift = drift.select { |d| d[:oss] != d[:saas] }
        dep_drift = drift.select { |d| d[:oss_dep] != d[:saas_dep] }

        if spec_drift.any?
          puts "  Spec versions:"
          spec_drift.each do |d|
            puts "    #{d[:name]}: #{d[:oss]} (Gemfile.lock) vs #{d[:saas]} (Gemfile.saas.lock)"
          end
        end

        if dep_drift.any?
          puts "  Dependency requirements:"
          dep_drift.each do |d|
            puts "    #{d[:name]}: #{d[:oss_dep]} (Gemfile.lock) vs #{d[:saas_dep]} (Gemfile.saas.lock)"
          end
        end

        subdep_drift = drift.select { |d| d[:subdep_changes]&.any? }
        if subdep_drift.any?
          puts "  Spec sub-dependency requirements:"
          subdep_drift.each do |d|
            d[:subdep_changes].each do |sd|
              puts "    #{d[:name]} → #{sd[:name]}: #{sd[:oss] || "(absent)"} (Gemfile.lock) vs #{sd[:saas] || "(absent)"} (Gemfile.saas.lock)"
            end
          end
        end

        puts "\nRun 'bin/bundle-drift correct' to restore alignment."
      end
    end
end

class GemfileDriftCorrector
  def correct
    drift = GemfileDriftChecker.new.check
    return puts "\nNothing to correct." if drift.empty?

    puts "\nRestoring alignment (Gemfile.saas.lock is authoritative)...\n\n"

    original_content = File.read(GEMFILE_LOCK)

    FileUtils.cp(GEMFILE_SAAS_LOCK, GEMFILE_LOCK)

    puts "▸ Re-locking Gemfile (seeded from Gemfile.saas.lock)"
    unless system("BUNDLE_GEMFILE=Gemfile bundle lock")
      File.write(GEMFILE_LOCK, original_content)
      abort("Failed to lock Gemfile. Restored original.")
    end

    puts "\n▸ Verifying alignment"
    new_drift = GemfileDriftChecker.new.check

    if new_drift.empty?
      puts "\n✓ Lock files are now in sync"
      show_diff(original_content, File.read(GEMFILE_LOCK))
    else
      puts "\n✗ Lock files still have drift after correction."
      puts "  Bundler couldn't resolve to matching versions."
      puts "  Restoring original Gemfile.lock."
      File.write(GEMFILE_LOCK, original_content)
      exit 1
    end
  end

  private
    def show_diff(original, corrected)
      require "tempfile"

      Tempfile.create("gemfile-lock-original") do |f|
        f.write(original)
        f.flush

        diff = `diff -u #{f.path} #{GEMFILE_LOCK} 2>/dev/null`
        unless diff.empty?
          puts "\nChanges made to Gemfile.lock:"
          puts diff
        end
      end
    end
end

class GemfileDriftForwarder
  def forward
    drift = GemfileDriftChecker.new.check
    return puts "\nNothing to forward." if drift.empty?

    puts "\nForwarding Gemfile.lock changes into Gemfile.saas.lock...\n\n"

    original_content = File.read(GEMFILE_SAAS_LOCK)
    patched_content = original_content.dup

    drift.each do |d|
      if d[:oss] != d[:saas]
        puts "  #{d[:name]}: spec #{d[:saas]} → #{d[:oss]}"
        patched_content.gsub!(/#{Regexp.escape(d[:name])} \(#{Regexp.escape(d[:saas])}([^)]*)\)/) do
          "#{d[:name]} (#{d[:oss]}#{$1})"
        end
      end

      if d[:oss_dep] && d[:saas_dep] && d[:oss_dep] != d[:saas_dep]
        puts "  #{d[:name]}: dep #{d[:saas_dep]} → #{d[:oss_dep]}"
        patched_content = patch_dependency_requirement(patched_content, d[:name], d[:saas_dep], d[:oss_dep])
      end

      if d[:subdep_changes]&.any?
        d[:subdep_changes].each do |sd|
          if sd[:oss] && sd[:saas]
            puts "  #{d[:name]} → #{sd[:name]}: subdep #{sd[:saas]} → #{sd[:oss]}"
            patched_content = patch_spec_subdependency(patched_content, d[:name], sd[:name], sd[:saas], sd[:oss])
          end
        end
      end
    end

    File.write(GEMFILE_SAAS_LOCK, patched_content)

    puts "\n▸ Verifying alignment"
    new_drift = GemfileDriftChecker.new.check

    if new_drift.empty?
      puts "\n✓ Lock files are now in sync"
      show_diff(original_content, patched_content)
    else
      puts "\n✗ Lock files still have drift after forwarding."
      puts "  Restoring original Gemfile.saas.lock."
      File.write(GEMFILE_SAAS_LOCK, original_content)
      exit 1
    end
  end

  private
    def patch_spec_subdependency(content, parent_name, dep_name, old_req, new_req)
      in_parent_spec = false
      content.lines.map do |line|
        if line.match?(/\A {4}\S/)
          in_parent_spec = line.match?(/\A {4}#{Regexp.escape(parent_name)} \(/)
        elsif !line.match?(/\A {6}/)
          in_parent_spec = false
        end

        if in_parent_spec && line.match?(/\A {6}#{Regexp.escape(dep_name)} \(#{Regexp.escape(old_req)}\)/)
          line.sub("#{dep_name} (#{old_req})", "#{dep_name} (#{new_req})")
        else
          line
        end
      end.join
    end

    def patch_dependency_requirement(content, name, old_req, new_req)
      in_deps = false
      content.lines.map do |line|
        if line.strip == "DEPENDENCIES"
          in_deps = true
        elsif in_deps && line.match?(/\A\S/)
          in_deps = false
        end

        if in_deps
          line.sub(/#{Regexp.escape(name)} \(#{Regexp.escape(old_req)}\)/, "#{name} (#{new_req})")
        else
          line
        end
      end.join
    end

    def show_diff(original, patched)
      require "tempfile"

      Tempfile.create("gemfile-saas-lock-original") do |f|
        f.write(original)
        f.flush

        diff = `diff -u #{f.path} #{GEMFILE_SAAS_LOCK} 2>/dev/null`
        unless diff.empty?
          puts "\nChanges made to Gemfile.saas.lock:"
          puts diff
        end
      end
    end
end

case command = ARGV[0] || "check"
when "check"
  exit 1 unless GemfileDriftChecker.new.check.empty?
when "correct"
  GemfileDriftCorrector.new.correct
when "forward"
  GemfileDriftForwarder.new.forward
when "self-test"
  require "minitest/autorun"
  require "tmpdir"

  LOCKFILE_TEMPLATE = <<~LOCKFILE
    GEM
      remote: https://rubygems.org/
      specs:
        shared_gem (%{shared_version})
        oss_only_gem (1.0.0)
        saas_only_gem (2.0.0)

    PLATFORMS
      ruby

    DEPENDENCIES
    %{dependencies}

    BUNDLED WITH
      2.6.2
  LOCKFILE

  LOCKFILE_WITH_SUBDEPS_TEMPLATE = <<~LOCKFILE
    GEM
      remote: https://rubygems.org/
      specs:
        child_gem (%{child_version})
        parent_gem (%{parent_version})
          child_gem (%{child_req})

    PLATFORMS
      ruby

    DEPENDENCIES
      parent_gem (%{parent_dep})

    BUNDLED WITH
      2.6.2
  LOCKFILE

  class GemfileDriftCheckerTest < Minitest::Test
    def setup
      @original_dir = Dir.pwd
      @tmpdir = Dir.mktmpdir("bundle-drift-test")
      Dir.chdir(@tmpdir)
    end

    def teardown
      Dir.chdir(@original_dir)
      FileUtils.rm_rf(@tmpdir)
    end

    def test_detects_spec_version_drift
      write_lockfiles(
        oss_version: "1.0.0", oss_deps: "  shared_gem (~> 1.0)\n  oss_only_gem",
        saas_version: "2.0.0", saas_deps: "  shared_gem (~> 2.0)\n  saas_only_gem"
      )

      drift = GemfileDriftChecker.new.send(:find_drift)

      assert_equal 1, drift.length
      assert_equal "shared_gem", drift.first[:name]
    end

    def test_detects_dependency_requirement_drift_even_when_spec_versions_match
      write_lockfiles(
        oss_version: "1.4.0", oss_deps: "  shared_gem (~> 1.4)\n  oss_only_gem",
        saas_version: "1.4.0", saas_deps: "  shared_gem (~> 1.3)\n  saas_only_gem"
      )

      drift = GemfileDriftChecker.new.send(:find_drift)

      assert_equal 1, drift.length
      assert_equal "shared_gem", drift.first[:name]
    end

    def test_reports_no_drift_when_specs_and_dependencies_match
      write_lockfiles(
        oss_version: "1.4.0", oss_deps: "  shared_gem (~> 1.4)\n  oss_only_gem",
        saas_version: "1.4.0", saas_deps: "  shared_gem (~> 1.4)\n  saas_only_gem"
      )

      drift = GemfileDriftChecker.new.send(:find_drift)

      assert_empty drift
    end

    def test_forward_patches_dependency_requirements_in_addition_to_spec_versions
      write_lockfiles(
        oss_version: "1.4.0", oss_deps: "  shared_gem (~> 1.4)\n  oss_only_gem",
        saas_version: "1.3.2", saas_deps: "  shared_gem (~> 1.3)\n  saas_only_gem"
      )

      capture_io { GemfileDriftForwarder.new.forward }

      saas_content = File.read("Gemfile.saas.lock")
      assert_includes saas_content, "shared_gem (1.4.0)"
      assert_includes saas_content, "shared_gem (~> 1.4)"
      refute_includes saas_content, "shared_gem (~> 1.3)"
    end

    def test_detects_spec_subdependency_drift_even_when_versions_match
      write_lockfiles_with_subdeps(
        oss_child_req: "~> 3, >= 3.2.0",
        saas_child_req: "~> 3, >= 3.1.0"
      )

      drift = GemfileDriftChecker.new.send(:find_drift)

      assert_equal 1, drift.length
      assert_equal "parent_gem", drift.first[:name]
      assert_equal 1, drift.first[:subdep_changes].length
      assert_equal "child_gem", drift.first[:subdep_changes].first[:name]
    end

    def test_no_drift_when_spec_subdependencies_match
      write_lockfiles_with_subdeps(
        oss_child_req: "~> 3, >= 3.2.0",
        saas_child_req: "~> 3, >= 3.2.0"
      )

      drift = GemfileDriftChecker.new.send(:find_drift)

      assert_empty drift
    end

    def test_forward_patches_spec_subdependency_requirements
      write_lockfiles_with_subdeps(
        oss_child_req: "~> 3, >= 3.2.0",
        saas_child_req: "~> 3, >= 3.1.0"
      )

      capture_io { GemfileDriftForwarder.new.forward }

      saas_content = File.read("Gemfile.saas.lock")
      assert_includes saas_content, "child_gem (~> 3, >= 3.2.0)"
      refute_includes saas_content, "child_gem (~> 3, >= 3.1.0)"
    end

    private
      def write_lockfiles(oss_version:, oss_deps:, saas_version:, saas_deps:)
        File.write("Gemfile.lock", LOCKFILE_TEMPLATE % { shared_version: oss_version, dependencies: oss_deps })
        File.write("Gemfile.saas.lock", LOCKFILE_TEMPLATE % { shared_version: saas_version, dependencies: saas_deps })
      end

      def write_lockfiles_with_subdeps(oss_child_req:, saas_child_req:)
        common = { parent_version: "1.0.0", child_version: "3.2.0", parent_dep: "~> 1.0" }
        File.write("Gemfile.lock", LOCKFILE_WITH_SUBDEPS_TEMPLATE % common.merge(child_req: oss_child_req))
        File.write("Gemfile.saas.lock", LOCKFILE_WITH_SUBDEPS_TEMPLATE % common.merge(child_req: saas_child_req))
      end
  end
else
  abort "Usage: bin/bundle-drift [check|correct|forward|self-test]"
end