# PUBLIC DOMAIN - NO LICENSE, NO WARRANTY
# Copyright 2025 TimeHexOn & foxhop & russell@unturf
# https://www.permacomputer.com

require "http/client"
require "json"

# UncloseAI Crystal Library
# OpenAI-compatible API client with streaming support
# Compatible with vLLM, Ollama, and OpenAI-compatible endpoints

# Model information structure
struct ModelInfo
  property id : String
  property endpoint : String
  property max_tokens : Int32

  def initialize(@id, @endpoint, @max_tokens = 8192)
  end
end

# UncloseAI Client class
class UncloseAI
  getter models : Array(ModelInfo)
  getter tts_endpoints : Array(String)
  property timeout : Int32

  def initialize(@timeout : Int32 = 30)
    @models = [] of ModelInfo
    @tts_endpoints = [] of String
    discover_models
  end

  # Discover models from environment variables
  private def discover_models
    puts "Initializing UncloseAI client..."

    # Discover chat/code models
    i = 1
    while i <= 9999
      endpoint = ENV["MODEL_ENDPOINT_#{i}"]?
      break unless endpoint

      puts "Endpoint #{i}: #{endpoint}"
      discover_models_from_endpoint(endpoint)
      i += 1
    end

    # Discover TTS endpoints
    i = 1
    while i <= 9999
      endpoint = ENV["TTS_ENDPOINT_#{i}"]?
      break unless endpoint
      @tts_endpoints << endpoint
      i += 1
    end

    puts "Discovered #{@models.size} models, #{@tts_endpoints.size} TTS endpoints\n"
  end

  # Discover models from a specific endpoint
  private def discover_models_from_endpoint(endpoint : String)
    begin
      response = HTTP::Client.get("#{endpoint}/models")
      json = JSON.parse(response.body)

      if json["data"]?
        json["data"].as_a.each do |model|
          model_id = model["id"].as_s

          # Skip modelperm-* entries
          next if model_id.starts_with?("modelperm-")

          max_tokens = model["max_model_len"]?.try(&.as_i) || 8192
          @models << ModelInfo.new(model_id, endpoint, max_tokens)
        end
      end
    rescue ex
      # Silently skip failed endpoints
    end
  end

  # Non-streaming chat completion
  def chat(messages : Array(Hash(String, String)), model_idx : Int32 = 0,
           max_tokens : Int32 = 100, temperature : Float64 = 0.7) : String
    raise "Invalid model index" if model_idx >= @models.size

    model = @models[model_idx]
    url = "#{model.endpoint}/chat/completions"

    payload = {
      "model"       => model.id,
      "messages"    => messages,
      "stream"      => false,
      "max_tokens"  => max_tokens,
      "temperature" => temperature,
    }

    response = HTTP::Client.post(
      url,
      headers: HTTP::Headers{"Content-Type" => "application/json"},
      body: payload.to_json
    )

    json = JSON.parse(response.body)
    json["choices"][0]["message"]["content"].as_s
  end

  # Streaming chat completion - yields content chunks as they arrive
  def chat_stream(messages : Array(Hash(String, String)), model_idx : Int32 = 0,
                  max_tokens : Int32 = 500, temperature : Float64 = 0.7, &block : String ->)
    raise "Invalid model index" if model_idx >= @models.size

    model = @models[model_idx]
    url = "#{model.endpoint}/chat/completions"

    payload = {
      "model"       => model.id,
      "messages"    => messages,
      "stream"      => true,
      "max_tokens"  => max_tokens,
      "temperature" => temperature,
    }

    HTTP::Client.post(
      url,
      headers: HTTP::Headers{"Content-Type" => "application/json"},
      body: payload.to_json
    ) do |response|
      response.body_io.each_line do |line|
        next unless line.starts_with?("data: ")

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

        begin
          json = JSON.parse(data)
          if content = json.dig?("choices", 0, "delta", "content")
            yield content.as_s
          end
        rescue
          # Skip malformed JSON
        end
      end
    end
  end

  # Text-to-speech generation
  def tts(text : String, voice : String = "alloy", output_file : String = "/tmp/speech.mp3") : Bool
    return false if @tts_endpoints.empty?

    endpoint = @tts_endpoints[0]
    url = "#{endpoint}/audio/speech"

    payload = {
      "model" => "tts-1",
      "voice" => voice,
      "input" => text,
    }

    response = HTTP::Client.post(
      url,
      headers: HTTP::Headers{"Content-Type" => "application/json"},
      body: payload.to_json
    )

    File.write(output_file, response.body)
    true
  rescue
    false
  end
end

# Demo program showing library usage
if __FILE__ == PROGRAM_NAME
  puts "=== UncloseAI Crystal Client (with Streaming) ===\n"

  # Initialize client
  client = UncloseAI.new

  if client.models.empty?
    puts "ERROR: No models discovered"
    exit 1
  end

  # Non-streaming chat example
  puts "=== Non-Streaming Chat ==="
  puts "Model: #{client.models[0].id}"

  begin
    messages = [{"role" => "user", "content" => "Explain quantum computing in one sentence"}]
    response = client.chat(messages)
    puts "Response: #{response}\n"
  rescue ex
    puts "Error: #{ex.message}\n"
  end

  # Streaming chat example
  model_idx = client.models.size >= 2 ? 1 : 0
  puts "=== Streaming Chat ==="
  puts "Model: #{client.models[model_idx].id}"
  print "Response: "

  begin
    messages = [{"role" => "user", "content" => "Write a hello world program in Crystal"}]
    client.chat_stream(messages, model_idx) do |content|
      print content
      STDOUT.flush
    end
    puts "\n"
  rescue ex
    puts "\nError: #{ex.message}\n"
  end

  # TTS example
  if !client.tts_endpoints.empty?
    puts "=== TTS Speech Generation ==="
    puts "Model: tts-1"

    if client.tts("Hello from UncloseAI Crystal client!", "alloy", "/tmp/speech.mp3")
      puts "Audio saved to /tmp/speech.mp3"
    else
      puts "TTS failed"
    end
  end

  puts "\n=== Examples Complete ==="
end