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

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

using HTTP
using JSON3

struct ModelInfo
    id::String
    endpoint::String
    max_tokens::Int
end

# UncloseAI Client type
mutable struct UncloseAIClient
    models::Vector{ModelInfo}
    tts_endpoints::Vector{String}
    timeout::Int
end

# Initialize client with auto-discovery
function init_client(timeout::Int=30)
    println("Initializing UncloseAI client...")

    models = ModelInfo[]
    tts_endpoints = String[]

    # Discover chat/code models from MODEL_ENDPOINT_1..9999
    for i in 1:9999
        endpoint = get(ENV, "MODEL_ENDPOINT_$i", nothing)
        isnothing(endpoint) && break

        println("Endpoint $i: $endpoint")

        try
            response = HTTP.get("$endpoint/models", readtimeout=10)
            data = JSON3.read(String(response.body))

            for model in data.data
                # Filter out modelperm-* entries
                if !startswith(model.id, "modelperm-")
                    max_tokens = get(model, :max_model_len, 8192)
                    push!(models, ModelInfo(model.id, endpoint, max_tokens))
                end
            end
        catch e
            # Silently skip failed endpoints
        end
    end

    # Discover TTS endpoints from TTS_ENDPOINT_1..9999
    for i in 1:9999
        endpoint = get(ENV, "TTS_ENDPOINT_$i", nothing)
        isnothing(endpoint) && break
        push!(tts_endpoints, endpoint)
    end

    println("Discovered $(length(models)) models, $(length(tts_endpoints)) TTS endpoints\n")

    return UncloseAIClient(models, tts_endpoints, timeout)
end

# Non-streaming chat completion
function chat(client::UncloseAIClient, messages::Vector; model_idx::Int=0, max_tokens::Int=100, temperature::Float64=0.7)
    if model_idx >= length(client.models)
        return (error="Invalid model index",)
    end

    model = client.models[model_idx + 1]

    request_data = Dict(
        "model" => model.id,
        "messages" => messages,
        "max_tokens" => max_tokens,
        "temperature" => temperature,
        "stream" => false
    )

    try
        response = HTTP.post(
            "$(model.endpoint)/chat/completions",
            ["Content-Type" => "application/json"],
            JSON3.write(request_data),
            readtimeout=client.timeout
        )

        data = JSON3.read(String(response.body))
        return (content=data.choices[1].message.content,)
    catch e
        return (error=string(e),)
    end
end

# Streaming chat completion - returns a Channel for async iteration
function chat_stream(client::UncloseAIClient, messages::Vector; model_idx::Int=0, max_tokens::Int=500, temperature::Float64=0.7)
    if model_idx >= length(client.models)
        error("Invalid model index")
    end

    model = client.models[model_idx + 1]

    request_data = Dict(
        "model" => model.id,
        "messages" => messages,
        "max_tokens" => max_tokens,
        "temperature" => temperature,
        "stream" => true
    )

    Channel() do channel
        try
            HTTP.open("POST", "$(model.endpoint)/chat/completions",
                     ["Content-Type" => "application/json"]) do http
                write(http, JSON3.write(request_data))
                closewrite(http)

                buffer = ""
                while !eof(http)
                    chunk = String(readavailable(http))
                    buffer *= chunk

                    while contains(buffer, "\n")
                        line_end = findfirst("\n", buffer)
                        line = buffer[1:line_end[1]-1]
                        buffer = buffer[line_end[1]+1:end]

                        if startswith(line, "data: ")
                            data_str = line[7:end]
                            if data_str == "[DONE]"
                                break
                            end

                            try
                                data = JSON3.read(data_str)
                                if haskey(data, :choices) && length(data.choices) > 0
                                    delta = data.choices[1].delta
                                    if haskey(delta, :content)
                                        put!(channel, delta.content)
                                    end
                                end
                            catch
                                # Skip malformed JSON
                            end
                        end
                    end
                end
            end
        catch e
            put!(channel, "Error: $e")
        end
    end
end

# Text-to-speech generation
function tts(client::UncloseAIClient, text::String; voice::String="alloy", output_file::String="/tmp/speech.mp3")
    if isempty(client.tts_endpoints)
        return (error="No TTS endpoints available",)
    end

    endpoint = client.tts_endpoints[1]

    request_data = Dict(
        "model" => "tts-1",
        "voice" => voice,
        "input" => text
    )

    try
        response = HTTP.post(
            "$endpoint/audio/speech",
            ["Content-Type" => "application/json"],
            JSON3.write(request_data),
            readtimeout=client.timeout
        )

        write(output_file, response.body)
        return (file=output_file,)
    catch e
        return (error=string(e),)
    end
end

# Demo program showing library usage
println("=== UncloseAI Julia Client (with Streaming) ===\n")

# Initialize client
client = init_client(30)

if isempty(client.models)
    println("ERROR: No models discovered")
    exit(1)
end

# Non-streaming chat example
println("=== Non-Streaming Chat ===")
println("Model: $(client.models[1].id)")

messages = [Dict("role" => "user", "content" => "Explain quantum computing in one sentence")]
result = chat(client, messages)

if haskey(result, :content)
    println("Response: $(result.content)\n")
else
    println("Error: $(result.error)\n")
end

# Streaming chat example
model_idx = length(client.models) >= 2 ? 1 : 0

println("=== Streaming Chat ===")
println("Model: $(client.models[model_idx + 1].id)")
print("Response: ")

stream_messages = [Dict("role" => "user", "content" => "Write a hello world program in Julia")]
for content in chat_stream(client, stream_messages, model_idx=model_idx)
    print(content)
end
println("\n")

# TTS example
if !isempty(client.tts_endpoints)
    println("=== TTS Speech Generation ===")
    println("Model: tts-1")

    tts_result = tts(client, "Hello from UncloseAI Julia client!", output_file="/tmp/speech.mp3")
    if haskey(tts_result, :file)
        println("Audio saved to $(tts_result.file)")
    else
        println("TTS failed: $(tts_result.error)")
    end
end

println("\n=== Examples Complete ===")
