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

local http = require("socket.http")
local https = require("ssl.https")
local ltn12 = require("ltn12")
local json = require("cjson")
local socket = require("socket")

-- UncloseAI - Lua client for OpenAI-compatible APIs with streaming support
local UncloseAI = {}
UncloseAI.__index = UncloseAI

function UncloseAI.new(opts)
    opts = opts or {}

    local self = setmetatable({}, UncloseAI)
    self.models = {}
    self.tts_endpoints = {}
    self.api_key = opts.api_key
    self.timeout = opts.timeout or 30
    self.debug = opts.debug or false

    -- Discover endpoints from environment
    local model_endpoints = opts.model_endpoints or self:_discover_env_endpoints("MODEL_ENDPOINT")
    local tts_endpoints = opts.tts_endpoints or self:_discover_env_endpoints("TTS_ENDPOINT")

    if self.debug then
        print(string.format("[DEBUG] Initialized with %d endpoint(s)", #model_endpoints))
    end

    -- Discover models
    self:_discover_models(model_endpoints)
    self.tts_endpoints = tts_endpoints

    return self
end

function UncloseAI:_discover_env_endpoints(prefix)
    local endpoints = {}
    for i = 1, 9999 do
        local endpoint = os.getenv(prefix .. "_" .. tostring(i))
        if not endpoint then break end
        table.insert(endpoints, endpoint)
    end
    return endpoints
end

function UncloseAI:_discover_models(endpoints)
    for _, endpoint in ipairs(endpoints) do
        if self.debug then
            print("[DEBUG] Discovering from: " .. endpoint)
        end

        local success, err = pcall(function()
            local response_body = {}
            local res, code = https.request{
                url = endpoint .. "/models",
                method = "GET",
                sink = ltn12.sink.table(response_body)
            }

            if code == 200 then
                local response = json.decode(table.concat(response_body))
                if response.data then
                    for _, model in ipairs(response.data) do
                        table.insert(self.models, {
                            id = model.id,
                            endpoint = endpoint,
                            max_tokens = model.max_model_len or 8192
                        })
                        if self.debug then
                            print("[DEBUG] Discovered: " .. model.id)
                        end
                    end
                end
            end
        end)

        if not success and self.debug then
            print("[DEBUG] Error: " .. tostring(err))
        end
    end
end

function UncloseAI:list_models()
    local result = {}
    for _, model in ipairs(self.models) do
        table.insert(result, {
            id = model.id,
            endpoint = model.endpoint,
            max_tokens = model.max_tokens
        })
    end
    return result
end

function UncloseAI:_resolve_model(model)
    if #self.models == 0 then
        error("No models available")
    end

    if not model then
        return self.models[1]
    end

    for _, m in ipairs(self.models) do
        if m.id == model then
            return m
        end
    end

    error("Model '" .. model .. "' not found")
end

function UncloseAI:chat(messages, opts)
    opts = opts or {}
    local model_info = self:_resolve_model(opts.model)
    local max_tokens = opts.max_tokens or 100
    local temperature = opts.temperature or 0.7

    local payload = {
        model = model_info.id,
        messages = messages,
        max_tokens = max_tokens,
        temperature = temperature,
        stream = false
    }

    local request_body = json.encode(payload)
    local response_body = {}

    local headers = {
        ["Content-Type"] = "application/json",
        ["Content-Length"] = tostring(#request_body)
    }

    if self.api_key then
        headers["Authorization"] = "Bearer " .. self.api_key
    end

    local res, code = https.request{
        url = model_info.endpoint .. "/chat/completions",
        method = "POST",
        headers = headers,
        source = ltn12.source.string(request_body),
        sink = ltn12.sink.table(response_body)
    }

    if code == 200 then
        return json.decode(table.concat(response_body))
    else
        error("Request failed with code: " .. tostring(code))
    end
end

function UncloseAI:chat_stream(messages, opts, callback)
    opts = opts or {}
    local model_info = self:_resolve_model(opts.model)
    local max_tokens = opts.max_tokens or 500
    local temperature = opts.temperature or 0.7

    local payload = {
        model = model_info.id,
        messages = messages,
        max_tokens = max_tokens,
        temperature = temperature,
        stream = true
    }

    local request_body = json.encode(payload)

    -- Parse URL
    local protocol, host, port, path = model_info.endpoint:match("^(https?)://([^:/]+):?(%d*)(.*)$")
    port = port and tonumber(port) or (protocol == "https" and 443 or 80)
    path = (path == "" and "/v1" or path) .. "/chat/completions"

    -- Create socket connection
    local sock = socket.tcp()
    sock:settimeout(self.timeout)

    local success, err = pcall(function()
        assert(sock:connect(host, port))

        -- For HTTPS, wrap socket with SSL
        if protocol == "https" then
            local ssl = require("ssl")
            sock = assert(ssl.wrap(sock, {mode = "client", protocol = "tlsv1_2"}))
            assert(sock:dohandshake())
        end

        -- Send HTTP request
        local headers = {
            "POST " .. path .. " HTTP/1.1",
            "Host: " .. host,
            "Content-Type: application/json",
            "Content-Length: " .. tostring(#request_body),
            "Connection: close"
        }

        if self.api_key then
            table.insert(headers, "Authorization: Bearer " .. self.api_key)
        end

        local request = table.concat(headers, "\r\n") .. "\r\n\r\n" .. request_body
        assert(sock:send(request))

        -- Read response headers
        local line = sock:receive("*l")
        while line and line ~= "" do
            line = sock:receive("*l")
        end

        -- Read streaming response
        local buffer = ""
        while true do
            local chunk, err = sock:receive(1024)
            if not chunk then break end

            buffer = buffer .. chunk
            local lines = {}

            for line in buffer:gmatch("([^\n]*)\n") do
                table.insert(lines, line)
            end

            -- Keep incomplete line in buffer
            local last_newline = buffer:find("\n[^\n]*$")
            if last_newline then
                buffer = buffer:sub(last_newline + 1)
            end

            -- Process complete lines
            for i = 1, #lines - 1 do
                local line = lines[i]:gsub("\r", "")
                if line:match("^data: ") then
                    local data = line:sub(7)
                    if data == "[DONE]" then
                        return
                    end

                    local success, chunk_data = pcall(json.decode, data)
                    if success and chunk_data.choices and chunk_data.choices[1] then
                        local delta = chunk_data.choices[1].delta
                        if delta and delta.content then
                            callback(delta.content)
                        end
                    end
                end
            end
        end
    end)

    sock:close()

    if not success and self.debug then
        print("[DEBUG] Stream error: " .. tostring(err))
    end
end

function UncloseAI:tts(text, opts)
    opts = opts or {}

    if #self.tts_endpoints == 0 then
        error("No TTS endpoints available")
    end

    local endpoint = self.tts_endpoints[1]
    local voice = opts.voice or "alloy"
    local model = opts.model or "tts-1"
    local response_format = opts.response_format or "mp3"

    local payload = {
        model = model,
        voice = voice,
        input = text,
        response_format = response_format
    }

    local request_body = json.encode(payload)
    local response_body = {}

    local headers = {
        ["Content-Type"] = "application/json",
        ["Content-Length"] = tostring(#request_body)
    }

    if self.api_key then
        headers["Authorization"] = "Bearer " .. self.api_key
    end

    local res, code = https.request{
        url = endpoint .. "/audio/speech",
        method = "POST",
        headers = headers,
        source = ltn12.source.string(request_body),
        sink = ltn12.sink.table(response_body)
    }

    if code == 200 then
        return table.concat(response_body)
    else
        error("Request failed with code: " .. tostring(code))
    end
end

-- Demo when run as script
if not pcall(debug.getlocal, 4, 1) then
    print("=== UncloseAI Lua Client (with Streaming) ===\n")

    local client = UncloseAI.new({debug = true})

    local models = client:list_models()
    if #models == 0 then
        print("ERROR: No models discovered. Set environment variables:")
        print("  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.")
        os.exit(1)
    end

    print(string.format("\nDiscovered %d model(s):", #models))
    for _, model in ipairs(models) do
        print(string.format("  - %s (max_tokens: %d)", model.id, model.max_tokens))
    end
    print()

    -- Non-streaming chat
    print("=== Non-Streaming Chat ===")
    local success, response = pcall(function()
        return client:chat({
            {role = "system", content = "You are a helpful AI assistant."},
            {role = "user", content = "Explain quantum computing in one sentence."}
        }, {max_tokens = 100})
    end)

    if success then
        local content = response.choices[1].message.content
        print("Response: " .. content .. "\n")
    else
        print("Error: " .. tostring(response) .. "\n")
    end

    -- Streaming chat
    print("=== Streaming Chat ===")
    local model_id = #models > 1 and models[2].id or nil
    print("Model: " .. (model_id or models[1].id))
    io.write("Response: ")
    io.flush()

    local success, err = pcall(function()
        client:chat_stream({
            {role = "system", content = "You are a coding assistant."},
            {role = "user", content = "Write a Lua function to check if a number is prime"}
        }, {model = model_id, max_tokens = 200}, function(content)
            io.write(content)
            io.flush()
        end)
    end)

    if not success then
        print("\nError: " .. tostring(err))
    end
    print("\n")

    -- TTS
    if #client.tts_endpoints > 0 then
        print("=== TTS Speech Generation ===")
        local success, audio_data = pcall(function()
            return client:tts("Hello from UncloseAI Lua client! This demonstrates streaming support.")
        end)

        if success then
            local file = io.open("speech.mp3", "wb")
            file:write(audio_data)
            file:close()
            print(string.format("[OK] Speech file created: speech.mp3 (%d bytes)\n", #audio_data))
        else
            print("[ERROR] TTS Error: " .. tostring(audio_data) .. "\n")
        end
    end

    print("=== Examples Complete ===")
end

return UncloseAI