#!/usr/bin/env ruby
#
# Ollama CLI tool for interacting with Ollama models via command line.
#
# This script provides a command-line interface to communicate with Ollama
# models, supporting chat sessions, prompt templating, and various
# configuration options. It can be used for both interactive and automated
# model interactions.
#
# Features:
#   - Chat sessions with model responses
#   - Prompt templating with variable substitution
#   - Streaming and thinking modes
#   - Configuration via environment variables or command-line options
#   - JSON-based chat history persistence
#
# Environment Variables:
#   - OLLAMA_URL: Base URL for Ollama server (default: 'http://localhost')
#   - OLLAMA_HOST: Hostname for Ollama server (used when OLLAMA_URL not set)
#   - OLLAMA_CLIENT: JSON configuration for client
#   - OLLAMA_MODEL: Default model to use (default: 'llama3.1')
#   - OLLAMA_MODEL_OPTIONS: JSON model options
#   - OLLAMA_SYSTEM: Default system prompt
#   - OLLAMA_PROMPT: Default user prompt

require 'ollama'
include Ollama
require 'tins'
include Tins::GO
require 'tins/xt/secure_write'
require 'tins/xt/full'
require 'json'
require 'tmpdir'

module Ollama::Handlers
  # A handler that starts a chat session and prints the initial response.
  #
  # @example
  #   chat(model: 'llama3.1', messages: message, &ChatStart)
  class ChatStart
    include Ollama::Handlers::Concern

    # The initialize method sets up a new Markdown handler instance with the
    # specified output.
    #
    # @param output [IO] the output stream to write markdown content to,
    #   defaults to $stdout
    def initialize(output: $stdout)
      super
      @output.sync = true
      @content     = ''
    end

    # The content attribute reader returns the textual content of the message.
    #
    # @return [ String ] the content of the message
    attr_reader :content

    # The call method processes a response by appending its content to an
    # internal buffer and output stream. It also appends a newline and flushes
    # the output when the response indicates completion.
    #
    # @param response [Object] the response object to process
    # @return [self] returns itself to allow for method chaining
    def call(response)
      if content = response.response
        @content << content
        @output << content
      end
      response.done and @output.puts
      self
    end
  end
end

# Returns the contents of a file or string, or a default value if neither is provided.
#
# @param [String] path_or_content The path to a file or a string containing
#                 the content.
#
# @param [String] default The default value to return if no valid input is
#                 given. Defaults to nil.
#
# @return [String] The contents of the file, the string, or the default value.
#
# @example Get the contents of a file
#   get_file_argument('path/to/file')
#
# @example Use a string as content
#   get_file_argument('string content')
#
# @example Return a default value if no valid input is given
#   get_file_argument(nil, default: 'default content')
def get_file_argument(path_or_content, default: nil)
  if path_or_content.present? && path_or_content.size < 2 ** 15 &&
      File.basename(path_or_content).size < 2 ** 8 &&
      File.exist?(path_or_content)
    then
    File.read(path_or_content)
  elsif path_or_content.present?
    path_or_content
  else
    default
  end
end

# The usage method displays the command-line usage information for the
# ollama_cli executable.
#
# This method prints out a formatted help message that describes all available
# options and their expected arguments for using the ollama_cli tool.
def usage
  puts <<~EOT
    Usage: #{File.basename($0)} [OPTIONS]

      -u URL         the ollama base url, $OLLAMA_URL
      -c CLIENT      the ollama client config (JSON), $OLLAMA_CLIENT
      -m MODEL       the ollama model to chat with, $OLLAMA_MODEL
      -M OPTIONS     the ollama model options (JSON), $OLLAMA_MODEL_OPTIONS
      -s SYSTEM      the system prompt as plain text, $OLLAMA_SYSTEM
      -p PROMPT      the user prompt as plain text, $OLLAMA_PROMPT
                     if it contains %{stdin} it is substituted by stdin input
      -P VARIABLE    sets prompt var %{foo} to "bar" if VARIABLE is foo=bar
      -H HANDLER     the handler to use for the response, defaults to ChatStart
      -I IMAGE       image is sent to the (visual) model (can be used many times)
      -S             use streaming for generation
      -T             use thinking for generation
      -d             enable debug mode
      -i             display ollama server version information
      -h             this help

  EOT
  exit 0
end

opts = go 'u:m:M:s:p:P:H:c:I:STdih', defaults: { ?H => 'ChatStart' }

opts[?h] and usage

base_url = opts[?u] || ENV['OLLAMA_URL'] || 'http://%s' % ENV.fetch('OLLAMA_HOST')
api_key  = ENV['OLLAMA_API_KEY']
client_config = Client::Config[
  { base_url:, api_key:, debug: !!opts[?d] } |
  JSON(get_file_argument(opts[?c], default: ENV['OLLAMA_CLIENT']).full? || '{}')
]
model         = (opts[?m] || ENV['OLLAMA_MODEL']).full? || 'llama3.1'
model_options = JSON(get_file_argument(opts[?M], default: ENV['OLLAMA_MODEL_OPTIONS']))
options = if model_options.is_a?(Hash)
            Ollama::Options.from_hash(model_options)
          else
            Ollama::Options.new
          end
system = get_file_argument(opts[?s], default: ENV['OLLAMA_SYSTEM'])
prompt = get_file_argument(opts[?p], default: ENV['OLLAMA_PROMPT'])
images = opts[?I].to_a.map { Ollama::Image.for_filename(_1) }

ollama = Client.configure_with(client_config)

case
when opts[?i]
  puts ollama.version.version
  exit
when prompt.nil?
  prompt = STDIN.read
else
  vars = prompt.named_placeholders
  stdin = (STDIN.read if vars.include?(:stdin)).to_s
  values = opts[?P].to_a.inject({ stdin: }) { |h, pair|
    n, v = pair.split(?=, 2)
    h.merge(n.to_sym => v)
  }
  default = -> v {
    msg = "prompt var #{v.inspect} not set"
    STDERR.puts msg
    "[#{msg}]"
  }
  prompt = prompt.named_placeholders_interpolate(values, default:)
end

if opts[?d]
  puts <<~EOT
    base_url      = #{base_url.inspect}
    model         = #{model.inspect}
    system        = #{system.inspect}
    prompt        = #{prompt.inspect}
    options       = #{options.to_json}
    client_config = #{client_config.to_json}
  EOT
end

handler = Ollama::Handlers.const_get(opts[?H])
handler = case
          when handler == Ollama::Handlers::ChatStart
            handler.new
          when handler == Ollama::Handlers::Markdown
            handler.new(stream: !!opts[?S])
          else
            handler
          end

ollama.generate(
  model:,
  system:,
  prompt:,
  options:,
  images: images,
  stream: !!opts[?S],
  think: !!opts[?T],
  &handler
)

if handler.is_a?(Ollama::Handlers::ChatStart)
  filename = File.join(Dir.tmpdir, 'chat_start_%u.json' % $$)
  File.secure_write(filename) do |out|
    JSON.dump(
      [
        Message.new(role: 'user', content: prompt),
        Message.new(role: 'assistant', content: handler.content),
      ],
      out
    )
  end
  if STDERR.tty?
    STDERR.puts "\nContinue the chat with:\n  ollama_chat -c '%s'" % filename
  end
end
