require 'optparse'
require 'fileutils'
require 'open3'
require 'json'

# Configuration
# Default target is loosely based on user preference, but we use auto-density usually.
options = {
  threshold: 5.0, # Default 5 GB/hr
  dry_run: false,
  quality: 29 # Default CQ
}

OptionParser.new do |opts|
  opts.banner = "Usage: ruby video_optimizer.rb [directory] [options]"

  opts.on("-t", "--threshold N", Float, "Minimum GB per hour to trigger processing (default: 5.0)") do |v|
    options[:threshold] = v
  end

  opts.on("-n", "--dry-run", "Identify files but do not transcode") do
    options[:dry_run] = true
  end

  opts.on("-q", "--quality N", Integer, "Set Constant Quality (CQ/CRF) value (Overrides auto-density)") do |v|
    options[:cq] = v
    options[:mode] = :cq
  end

  opts.on("-b", "--bitrate RATE", "Set fixed bitrate (e.g. '2M', '2000k')") do |v|
    options[:bitrate] = v
    options[:mode] = :bitrate
  end

  opts.on("-s", "--target-size GB", Float, "Target file size in GB") do |v|
    options[:target_size] = v
    options[:mode] = :size
  end

  opts.on("-h", "--help", "Prints this help") do
    puts opts
    exit
  end
end.parse!

# Default mode
options[:mode] ||= :auto_density

directory = ARGV[0]
if directory.nil? || !Dir.exist?(directory)
  puts "Error: Please provide a valid directory."
  exit 1
end

# --- ENCODER LOGIC ---

def detect_encoders
  stdout, _, _ = Open3.capture3('ffmpeg', '-v', 'quiet', '-encoders')
  encoders = []
  encoders << :nvenc if stdout.include?("hevc_nvenc")
  encoders << :qsv if stdout.include?("hevc_qsv")
  encoders << :cpu if stdout.include?("libx265")
  encoders
end

def select_best_encoder(available_encoders)
  if available_encoders.include?(:nvenc)
    puts "  [Encoder] Selected: NVIDIA (hevc_nvenc)"
    return :nvenc
  elsif available_encoders.include?(:qsv)
    puts "  [Encoder] Selected: Intel QSV (hevc_qsv)"
    return :qsv
  elsif available_encoders.include?(:cpu)
    puts "  [Encoder] Selected: CPU (libx265)"
    return :cpu
  else
    puts "  [Error] No suitable HEVC encoder found!"
    exit 1
  end
end

def build_encoder_args(encoder, mode, target_val)
  args = []
  
  case encoder
  when :nvenc
    args << '-c:v' << 'hevc_nvenc' << '-preset' << 'p4'
    case mode
    when :cq
      args << '-rc' << 'constqp' << '-qp' << target_val.to_s
    when :bitrate, :size # Both use target bitrate
      args << '-rc' << 'vbr' << '-b:v' << target_val << '-maxrate' << target_val << '-bufsize' << (target_val.to_i * 2).to_s
    end
    
  when :qsv
    args << '-c:v' << 'hevc_qsv' << '-preset' << 'veryfast' << '-load_plugin' << 'hevc_hw'
    case mode
    when :cq
      # QSV uses -global_quality or -q:v for ICQ. roughly maps to CRF.
      args << '-global_quality' << target_val.to_s
    when :bitrate, :size
      # QSV VBR
      args << '-b:v' << target_val << '-maxrate' << target_val
    end

  when :cpu
    args << '-c:v' << 'libx265' << '-preset' << 'medium'
    case mode
    when :cq
      args << '-crf' << target_val.to_s
    when :bitrate, :size
      args << '-b:v' << target_val << '-maxrate' << target_val << '-bufsize' << (target_val.to_i * 2).to_s
    end
  end
  
  args
end

# --- METADATA & UTILS ---

def get_video_metadata(file_path)
  cmd = ['ffprobe', '-v', 'quiet', '-print_format', 'json', '-show_format', '-show_streams', '-select_streams', 'v:0', file_path]
  stdout, stderr, status = Open3.capture3(*cmd)
  return nil unless status.success?

  begin
    data = JSON.parse(stdout)
    format_info = data['format']
    stream_info = data['streams'][0]
    return nil unless format_info && stream_info

    size_bytes = format_info['size'].to_i
    duration_sec = format_info['duration'].to_f
    codec = stream_info['codec_name']
    
    fps_str = stream_info['avg_frame_rate']
    if fps_str && fps_str.include?('/')
      num, den = fps_str.split('/').map(&:to_f)
      fps = (den > 0) ? (num / den) : 30.0
    else
      fps = 30.0
    end

    return nil if size_bytes == 0 || duration_sec == 0

    { size_bytes: size_bytes, duration_sec: duration_sec, codec: codec, fps: fps }
  rescue => e
    puts "Failed to parse metadata: #{e.message}"
    nil
  end
end

def calculate_target_bitrate_from_size(duration_sec, target_gb)
  target_bytes = (target_gb * 1024 * 1024 * 1024).to_i
  audio_bitrate_bps = 192 * 1000
  estimated_audio_size = duration_sec * (audio_bitrate_bps / 8.0)
  target_video_size = target_bytes - estimated_audio_size
  return 500_000 if target_video_size <= 0
  (target_video_size * 8) / duration_sec
end

# --- PROCESS ---

AVAILABLE_ENCODERS = detect_encoders
BEST_ENCODER = select_best_encoder(AVAILABLE_ENCODERS)

def process_file(file_path, options)
  metadata = get_video_metadata(file_path)
  return unless metadata

  accepted_codecs = ['h264', 'avc1']
  return unless accepted_codecs.include?(metadata[:codec].to_s.downcase)

  size_gb = metadata[:size_bytes].to_f / (1024**3)
  hours = metadata[:duration_sec] / 3600.0
  gb_per_hour = size_gb / hours

  if gb_per_hour >= options[:threshold]
    puts "\n[MATCH] #{file_path}"
    puts "  Size: #{size_gb.round(2)} GB, Density: #{gb_per_hour.round(2)} GB/hr"

    if options[:dry_run]
      puts "  [Dry Run] Would transcode this file."
      return
    end

    # Determine Encode Mode & Target Value
    mode = :cq
    target_val = 29 # Default CQ

    if options[:mode] == :auto_density
      target_gb_hr = options[:threshold] / 2.0
      target_gb = hours * target_gb_hr
      bitrate = calculate_target_bitrate_from_size(metadata[:duration_sec], target_gb)
      
      mode = :bitrate # Use bitrate targeting for density match
      target_val = bitrate.to_i.to_s # e.g. "2000000"
      puts "  Auto-Density Target: #{target_gb.round(2)} GB (#{target_gb_hr} GB/hr) -> #{target_val.to_i/1000} kbps"
      
    elsif options[:mode] == :size
      bitrate = calculate_target_bitrate_from_size(metadata[:duration_sec], options[:target_size])
      mode = :bitrate
      target_val = bitrate.to_i.to_s
      puts "  Target Size: #{options[:target_size]} GB -> #{target_val.to_i/1000} kbps"

    elsif options[:mode] == :bitrate
      mode = :bitrate
      target_val = options[:bitrate] # User string e.g. "2M"
      puts "  Target Bitrate: #{target_val}"
      
    else # :cq
      mode = :cq
      target_val = options[:cq] || 29
      puts "  Target Quality (CQ/CRF): #{target_val}"
    end

    # Build flags
    encoder_args = build_encoder_args(BEST_ENCODER, mode, target_val)
    
    # Common Plex/Optimization Args
    gop_size = (metadata[:fps] * 2).round
    common_args = [
      '-pix_fmt', 'yuv420p',     # 8-bit HEVC (Main Profile)
      '-g', "#{gop_size}",       # Fixed GOP
      '-keyint_min', "#{gop_size}",
      '-movflags', '+faststart', # Streaming optimization
      '-c:a', 'copy'
    ]

    original_ext = File.extname(file_path)
    temp_file = file_path.sub(original_ext, ".tmp#{original_ext}")

    cmd = ['ffmpeg', '-y', '-hwaccel', 'auto', '-i', file_path] + encoder_args + common_args + [temp_file]

    puts "  Starting transcode with #{BEST_ENCODER}..."
    start_time = Time.now
    
    stdout, stderr, status = Open3.capture3(*cmd)
    
    if status.success?
      duration = Time.now - start_time
      puts "  Done in #{duration.round(1)}s."
      
      # Validation
      if File.exist?(temp_file) && File.size(temp_file) > 1000
        new_size_gb = File.size(temp_file).to_f / (1024**3)
        reduction = ((size_gb - new_size_gb) / size_gb * 100.0).round(1)
        
        puts "  New Size: #{new_size_gb.round(2)} GB (#{reduction}% reduction)."

        if reduction < 33.0
          puts "  [SKIP] Reduction too low (<33%). Keeping original."
          File.delete(temp_file)
        else
          # Verification
          new_meta = get_video_metadata(temp_file)
          if new_meta && (new_meta[:duration_sec] - metadata[:duration_sec]).abs < 30
            File.delete(file_path)
            File.rename(temp_file, file_path)
            puts "  [SUCCESS] File replaced."
          else
            puts "  [ERROR] Duration mismatch. Keeping original."
            File.delete(temp_file)
          end
        end
      else
        puts "  [ERROR] Output invalid."
      end
    else
      puts "  [ERROR] Transcode failed."
      puts stderr # Debug
      File.delete(temp_file) if File.exist?(temp_file)
    end
  end
end

puts "Scanning #{directory}..."
files = []
Dir.glob("#{directory}/**/*").each do |f| 
  next if File.directory?(f)
  next unless ['.mp4', '.mkv', '.avi', '.mov', '.m4v'].include?(File.extname(f).downcase)
  files << [f, File.size(f)] rescue nil
end

puts "Sorting #{files.count} files by size..."
files.sort_by! { |_, size| -size }

files.each { |f, _| process_file(f, options) }