#!/usr/bin/env ruby
#
# A utility script that displays information about running Ollama models in a
# formatted table.
#
# This script queries the Ollama API's /api/ps endpoint and presents running
# models with detailed information including:
#
# * Model name and ID
# * Memory usage (size and processor allocation)
# * Context window length
# * Parameter count and quantization level
# * Time until expiration
#
# == Usage
#
#   ollama_ps
#
# The script uses the OLLAMA_URL environment variable for the API endpoint,
# defaulting to 'http://' + OLLAMA_HOST if not set.
#
# == Output
#
# Displays a formatted table with columns:
# * NAME - Model name
# * ID - Truncated model digest
# * SIZE - Human-readable model size
# * PROCESSOR - CPU/GPU allocation percentage
# * CONTEXT - Context window size
# * PARAMS - Parameter count (e.g., 30.5B, 23M)
# * QUANT - Quantization level (e.g., Q4_K_M, F16)
# * UNTIL - Time until model expiration

require 'ollama'
include Ollama
require 'terminal-table'
require 'term/ansicolor'
require 'tins/xt'
include Tins::GO

# The usage method displays the command-line usage information for the
# application.
#
# This method prints a formatted help message to the standard output, showing
# the available command-line options and their descriptions, and then exits the
# program. It is typically called when the user requests help or when invalid
# arguments are provided.
def usage
  puts <<~EOT
    Usage: #{File.basename($0)} [OPTIONS]

      -f FORMAT      output format: json, yaml, or table (default)
      -i INTERVAL    output the ps table every INTERVAL seconds
      -h             this help

  EOT
  exit 0
end

# The base_url method returns the Ollama API base URL.
#
# It first checks for the OLLAMA_URL environment variable.
# If not set, it defaults to 'http://' followed by the OLLAMA_HOST environment
# variable.
#
# @return [ String ] the base URL for the Ollama API
def base_url
  ENV['OLLAMA_URL'] || 'http://%s' % ENV.fetch('OLLAMA_HOST')
end

# The format_bytes method formats a number of bytes into a human-readable
# string with units.
#
# @param bytes [ Integer ] the number of bytes to format
#
# @return [ String ] a formatted string representing the byte value with
#   appropriate units
def format_bytes(bytes)
  Tins::Unit.format(bytes, format: '%.2f %U', unit: ?B)
end

# The format_expiry method converts an expiration time into a formatted
# duration string.
#
# This method takes an expiration timestamp, calculates the time difference
# from the current time, and formats it into a human-readable duration string
# showing days, hours, minutes, and seconds.
#
# @param expires_at [ String ] the expiration timestamp to be formatted
#
# @return [ String ] a formatted duration string representing the time until
#   expiration
def format_expiry(expires_at)
  Tins::Duration.new(Time.parse(expires_at) - Time.now).format('%S%d+%h:%m:%s')
end

# The format_processor method computes a CPU/GPU usage string based on memory
# sizes.
#
# This method takes two memory size parameters and calculates the percentage of
# GPU and CPU usage, returning a formatted string that indicates the
# distribution of resources. It handles special cases where all memory is on
# GPU or CPU, and calculates the percentage for mixed scenarios.
#
# @param size_vram [ Integer ] the amount of VRAM allocated
# @param size_total [ Integer ] the total amount of memory available
#
# @return [ String ] a formatted string indicating CPU/GPU usage percentage
#                     or special cases like "100% GPU" or "100% CPU"
def format_processor(size_vram, size_total)
  if size_vram == size_total
    "100% GPU"
  elsif size_vram == 0
    "100% CPU"
  else
    cpu_percent = ((size_total - size_vram) * 100.0 / size_total).round(0)
    gpu_percent = 100 - cpu_percent
    "#{cpu_percent}%/#{gpu_percent}% CPU/GPU"
  end
end

# The fetch_ps_models method retrieves information about running models from an
# Ollama server.
#
# This method creates a new Ollama client instance using the provided base URL,
# executes a ps command to fetch details about currently running models, and
# returns the array of model information if any models are running.
#
# @return [ Array<Hash>, nil ] an array of model information hashes if models
#   are running, or nil if no models are currently running
def fetch_ps_models
  ollama = Client.new(base_url:)
  result = ollama.ps
  models = result.models
  models.empty? and return
  models
end

# The interpret_models method processes an array of model objects into a
# standardized hash format.
#
# This method transforms model data into a consistent structure with specific
# keys including name, id, size, processor, context, params, quant, and until.
#
# @param models [ Array ] an array of model objects to be processed
#
# @return [ Array<Hash> ] an array of hashes containing standardized model information
def interpret_models(models)
  names = %w[ name id size processor context params quant until ]
  models.map do |model|
    hash = {}
    names.zip(
      [
        model.name,
        model.digest[0, 12],
        format_bytes(model.size),
        format_processor(model.size_vram, model.size),
        format_bytes(model.context_length),
        model.details&.parameter_size || 'n/a',
        model.details&.quantization_level || 'n/a',
        format_expiry(model.expires_at),
      ]
    ) do |n, v|
      hash[n] = v
    end
    hash
  end
end

# The ps_table method displays a formatted table of model information.
#
# This method takes an array of model hashes and presents them in a formatted
# table with aligned columns, using Unicode box-drawing characters for borders.
# It formats the size, parameters, and quantization columns to be right-aligned
# for better readability.
#
# @param models [ Array<Hash> ] an array of hashes containing model information
def ps_table(models)
  table = Terminal::Table.new
  table.style = {
    all_separators: true,
    border:         :unicode_round,
  }
  headings = models.first.keys.map(&:upcase)
  table.headings = headings
  models.each { |model| table << model.values }
  table.align_column(headings.index("SIZE"), :right)
  table.align_column(headings.index("PARAMS"), :right)
  table.align_column(headings.index("QUANT"), :right)
  puts table
end

opts = go 'f:i:h'

opts[?h] and usage

case opts[?f]
when 'json'
  models = fetch_ps_models or exit
  puts JSON.pretty_generate(interpret_models(models))
when 'yaml'
  models = fetch_ps_models or exit
  YAML.dump(interpret_models(models), STDOUT)
else
  duration = opts[?i].to_i.clamp(0..)
  begin
    print Term::ANSIColor.clear_screen, Term::ANSIColor.move_home if duration > 0
    if models = fetch_ps_models.full?
      ps_table(interpret_models(models))
    end
    duration == 0 and break
    sleep duration
  rescue => e
    warn "Caught #{e.class}: #{e}"
    sleep duration
  end while duration
end
