#!/usr/bin/env ruby
# frozen_string_literal: true

require 'net/http'
require 'uri'
require 'json'

# uncloseai. - Ruby client for OpenAI-compatible APIs with streaming support
# Ruby requires class names to be constants, so we use a class constant name
class Uncloseai
  attr_reader :models, :tts_endpoints

  def initialize(endpoints: nil, tts_endpoints: nil, api_key: nil, timeout: 30, debug: false)
    @api_key = api_key
    @timeout = timeout
    @debug = debug
    @models = []
    @tts_endpoints = []

    endpoints ||= discover_endpoints_from_env('MODEL_ENDPOINT')
    tts_endpoints ||= discover_endpoints_from_env('TTS_ENDPOINT')

    puts "[DEBUG] Initialized with #{endpoints.length} endpoint(s)" if @debug

    discover_models(endpoints)
    @tts_endpoints = tts_endpoints
  end

  def list_models
    @models
  end

  def chat(messages, model: nil, max_tokens: 100, temperature: 0.7, **kwargs)
    model_info = resolve_model(model)

    payload = {
      model: model_info[:id],
      messages: messages,
      max_tokens: max_tokens,
      temperature: temperature,
      **kwargs
    }

    response = http_request("#{model_info[:endpoint]}/chat/completions", :post, payload)
    JSON.parse(response)
  end

  def chat_stream(messages, model: nil, max_tokens: 500, temperature: 0.7, **kwargs)
    model_info = resolve_model(model)

    payload = {
      model: model_info[:id],
      messages: messages,
      max_tokens: max_tokens,
      temperature: temperature,
      stream: true,
      **kwargs
    }

    uri = URI.parse("#{model_info[:endpoint]}/chat/completions")
    http = Net::HTTP.new(uri.host, uri.port)
    http.use_ssl = uri.scheme == 'https'
    http.read_timeout = @timeout

    request = Net::HTTP::Post.new(uri.request_uri)
    request['Content-Type'] = 'application/json'
    request['Authorization'] = "Bearer #{@api_key}" if @api_key
    request.body = payload.to_json

    buffer = ''
    http.request(request) do |response|
      raise "HTTP #{response.code}" unless response.code.to_i == 200

      response.read_body do |chunk|
        buffer += chunk
        lines = buffer.split("\n")
        buffer = lines.pop || ''

        lines.each do |line|
          next unless line.start_with?('data: ')

          data = line[6..-1].strip
          break if data == '[DONE]'

          begin
            parsed = JSON.parse(data)
            yield parsed
          rescue JSON::ParserError => e
            puts "[DEBUG] Parse error: #{e.message}" if @debug
          end
        end
      end
    end
  end

  def tts(text, voice: 'alloy', model: 'tts-1')
    raise 'No TTS endpoints available' if @tts_endpoints.empty?

    payload = {
      model: model,
      voice: voice,
      input: text
    }

    http_request("#{@tts_endpoints[0]}/audio/speech", :post, payload)
  end

  private

  def discover_endpoints_from_env(prefix)
    endpoints = []
    (1..9999).each do |i|
      endpoint = ENV["#{prefix}_#{i}"]
      break unless endpoint
      endpoints << endpoint
    end
    endpoints
  end

  def discover_models(endpoints)
    endpoints.each do |endpoint|
      puts "[DEBUG] Discovering from: #{endpoint}" if @debug

      begin
        response = http_request("#{endpoint}/models", :get)
        data = JSON.parse(response)

        data['data'].each do |model|
          @models << {
            id: model['id'],
            endpoint: endpoint,
            max_tokens: model['max_model_len'] || 8192
          }
          puts "[DEBUG] Discovered: #{model['id']}" if @debug
        end
      rescue => e
        puts "[DEBUG] Error: #{e.message}" if @debug
      end
    end
  end

  def resolve_model(model)
    raise 'No models available' if @models.empty?
    return @models[0] if model.nil?

    found = @models.find { |m| m[:id] == model }
    raise "Model '#{model}' not found" unless found
    found
  end

  def http_request(url, method, payload = nil)
    uri = URI.parse(url)
    http = Net::HTTP.new(uri.host, uri.port)
    http.use_ssl = uri.scheme == 'https'
    http.read_timeout = @timeout

    request = case method
              when :get
                Net::HTTP::Get.new(uri.request_uri)
              when :post
                req = Net::HTTP::Post.new(uri.request_uri)
                req['Content-Type'] = 'application/json'
                req.body = payload.to_json if payload
                req
              end

    request['Authorization'] = "Bearer #{@api_key}" if @api_key

    response = http.request(request)
    raise "HTTP #{response.code}" unless response.code.to_i == 200
    response.body
  end
end

# Demo when run as script
if __FILE__ == $0
  puts "=== uncloseai. Ruby Client (with Streaming) ===\n"

  client = Uncloseai.new(debug: true)

  if client.models.empty?
    puts "ERROR: No models discovered. Set environment variables:"
    puts "  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc."
    exit 1
  end

  puts "\nDiscovered #{client.models.length} model(s):"
  client.models.each do |m|
    puts "  - #{m[:id]} (max_tokens: #{m[:max_tokens]})"
  end
  puts

  # Non-streaming chat
  puts "=== Non-Streaming Chat ==="
  response = client.chat([
    { role: 'system', content: 'You are a helpful AI assistant.' },
    { role: 'user', content: 'Explain quantum computing in one sentence.' }
  ])
  puts "Response: #{response['choices'][0]['message']['content']}\n\n"

  # Streaming chat
  puts "=== Streaming Chat ==="
  model_id = client.models.length > 1 ? client.models[1][:id] : nil
  puts "Model: #{model_id || client.models[0][:id]}"
  print "Response: "

  client.chat_stream([
    { role: 'system', content: 'You are a coding assistant.' },
    { role: 'user', content: 'Write a Ruby function to check if a number is prime' }
  ], model: model_id, max_tokens: 200) do |chunk|
    content = chunk.dig('choices', 0, 'delta', 'content')
    print content if content
  end

  puts "\n\n"

  # TTS
  if client.tts_endpoints.any?
    puts "=== TTS Speech Generation ==="
    audio_data = client.tts('Hello from uncloseai. Ruby client! This demonstrates streaming support.')
    File.binwrite('speech.mp3', audio_data)
    puts "[OK] Speech file created: speech.mp3 (#{audio_data.bytesize} bytes)\n\n"
  end

  puts "=== Examples Complete ==="
end