#!/usr/bin/env ruby

require 'json'
require 'net/http'
require 'socket'
require "openai"

def request(method, path, body=nil, base_url: "http://localhost:50080", auth: nil)
  body_json = JSON.generate(body) if body
  uri = URI("#{base_url}#{path}")
  # Use https if base_url is https
  response = Net::HTTP.start(uri.host, uri.port, use_ssl: base_url.start_with?('https://')) do |http|
    request = Net::HTTP.const_get(method.capitalize).new(uri)
    request.body = body_json
    request['Content-Type'] = 'application/json'
    request['Authorization'] = "Bearer #{auth}" if auth
    http.request(request)
  end

  if response.code.to_i >= 400
    raise "Request failed with status #{response.code}: #{response.body}"
  end

  if response.header['Content-Type'] == 'application/json'
    JSON.parse(response.body)
  else
    response.body
  end
end

def run_tests
  file_dir = File.dirname(__FILE__)

  cmd = "cargo run"
  
  # Run the command in a separate process
  openai_endpoint = "https://api.openai.com/v1"
  openai_api_key = ENV.fetch("OPENAI_API_KEY")
  pid = Process.spawn({
    "OPENAI_ENDPOINT" => openai_endpoint,
    "OPENAI_API_KEY" => openai_api_key,
  }, cmd)
  puts "Running command: #{cmd} (PID: #{pid})"
  
  # Wait for the process to start listening on port 50080
  connected = false
  while !connected do
    begin
      TCPSocket.open('localhost', 50080) { connected = true }
    rescue Errno::ECONNREFUSED
      sleep 0.3
    end
  end
  
  puts "Running tests..."

  puts "Test healthcheck"
  response = request(:get, '/health')
  raise "Expected healthcheck response, got #{response.inspect}" unless response == { "healthy" => true }

  puts "Test create project"
  response = request(:post, '/admin/v1/projects', {
    "name" => "Test Project",
    "description" => "Test Description",
  }, auth: "admin")
  id = response["id"]
  token = response["token"]
  raise "Expected create project response, got #{response.inspect}" unless token

  puts "Test regular completion"
  response = response = request(:post, '/v1/openai/v1/chat/completions', {
    "model" => "gpt-4o",
    "messages" => [{ "role" => "user", "content" => "Hello, how are you?" }],
  }, auth: token)
  raise "Expected response, got #{response.inspect}" unless response.dig("choices", 0, "message", "content")

  puts "Test usage metrics"
  response = request(:get, "/v1/metrics", auth: token)
  raise "Expected usage metrics, got #{response.inspect}" unless response[0]['prompt_token_count'] == 13

  puts "Test streaming completion"
  streaming_response = response = request(:post, '/v1/openai/v1/chat/completions', {
    "model" => "gpt-4o",
    "messages" => [{ "role" => "user", "content" => "Hello, how are you?" }],
    "stream" => true
  }, auth: token)
  raise "Expected SSE stream, got #{response.inspect}" unless response.split("\n")[-2] == "data: [DONE]"

  puts "Test streaming usage metrics"
  response = request(:get, "/v1/metrics", auth: token)
  raise "Expected 2 usage metrics, got #{response.length}: #{response.inspect}" unless response.length == 2
  raise "Expected usage metrics, got #{response.inspect}" unless response.all? { |m| m['prompt_token_count'] == 13 }

  puts "Test if proxied stream is compatible with OpenAI client when streaming"
  results = ""
  client = OpenAI::Client.new(
    uri_base: "http://localhost:50080/v1/openai",
    access_token: token,
    log_errors: true
  )
  response = client.chat(
    parameters: {
      model: "gpt-4o", # Required.
      messages: [{ role: "user", content: "Hello, how are you?"}], # Required.
      stream: proc do |chunk, _bytesize|
        content = chunk.dig("choices", 0, "delta", "content")
        results << content if content
      end
    }
  )
  raise "Expected results to come in, got \n\n#{results.inspect}" unless results.length > 0

  puts "Test if proxied stream is compatible with OpenAI client when not streaming"
  results = client.chat(
    parameters: {
      model: "gpt-4o", # Required.
      messages: [{ role: "user", content: "Hello, how are you?"}], # Required.
    }
  )
  raise "Expected results to come in, got \n\n#{results.inspect}" unless results.length > 0
ensure
  # Kill the process
  Process.kill('TERM', pid) if pid
end

run_tests

puts "\n\nAll tests passed!\n\n"