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

import java.net.URI
import java.net.http.HttpClient
import java.net.http.HttpRequest
import java.net.http.HttpResponse
import java.io.File
import java.io.BufferedReader
import java.io.InputStreamReader
import org.json.JSONObject
import org.json.JSONArray
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.*

data class ModelInfo(
    val id: String,
    val endpoint: String,
    val maxTokens: Int
)

data class ChatMessage(
    val role: String,
    val content: String
)

class UncloseAI(
    private val modelEndpoints: List<String>? = null,
    private val ttsEndpoints: List<String>? = null,
    private val apiKey: String? = null,
    private val timeout: Long = 30000,
    private val debug: Boolean = false
) {
    private val models = mutableListOf<ModelInfo>()
    private val ttsEndpointList = mutableListOf<String>()
    private val client = HttpClient.newHttpClient()
    private var initialized = false

    init {
        val endpoints = modelEndpoints ?: discoverEndpointsFromEnv("MODEL_ENDPOINT")
        val ttsEnds = ttsEndpoints ?: discoverEndpointsFromEnv("TTS_ENDPOINT")

        if (debug) {
            println("[DEBUG] Initialized with ${endpoints.size} endpoint(s)")
        }

        discoverModels(endpoints)
        ttsEndpointList.addAll(ttsEnds)
        initialized = true
    }

    private fun discoverEndpointsFromEnv(prefix: String): List<String> {
        val endpoints = mutableListOf<String>()
        for (i in 1..9999) {
            val endpoint = System.getenv("${prefix}_$i") ?: break
            endpoints.add(endpoint)
        }
        return endpoints
    }

    private fun discoverModels(endpoints: List<String>) {
        for (endpoint in endpoints) {
            if (debug) {
                println("[DEBUG] Discovering from: $endpoint")
            }

            try {
                val request = HttpRequest.newBuilder()
                    .uri(URI.create("$endpoint/models"))
                    .GET()
                    .build()

                val response = client.send(request, HttpResponse.BodyHandlers.ofString())
                val jsonResponse = JSONObject(response.body())
                val data = jsonResponse.getJSONArray("data")

                for (i in 0 until data.length()) {
                    val model = data.getJSONObject(i)
                    val modelId = model.getString("id")

                    if (modelId.startsWith("modelperm-") || modelId.startsWith("chatcmpl-")) {
                        continue
                    }

                    val maxTokens = model.optInt("max_model_len", 8192)
                    models.add(ModelInfo(modelId, endpoint, maxTokens))

                    if (debug) {
                        println("[DEBUG] Discovered: $modelId")
                    }
                }
            } catch (e: Exception) {
                if (debug) {
                    println("[DEBUG] Error: ${e.message}")
                }
            }
        }
    }

    fun listModels(): List<ModelInfo> = models.toList()

    fun chat(
        messages: List<ChatMessage>,
        model: String? = null,
        maxTokens: Int = 100,
        temperature: Double = 0.7
    ): JSONObject {
        val modelInfo = resolveModel(model)

        val messagesArray = JSONArray()
        for (msg in messages) {
            messagesArray.put(
                JSONObject()
                    .put("role", msg.role)
                    .put("content", msg.content)
            )
        }

        val payload = JSONObject()
            .put("model", modelInfo.id)
            .put("messages", messagesArray)
            .put("max_tokens", maxTokens)
            .put("temperature", temperature)
            .put("stream", false)

        val requestBuilder = HttpRequest.newBuilder()
            .uri(URI.create("${modelInfo.endpoint}/chat/completions"))
            .header("Content-Type", "application/json")
            .POST(HttpRequest.BodyPublishers.ofString(payload.toString()))

        if (apiKey != null) {
            requestBuilder.header("Authorization", "Bearer $apiKey")
        }

        val response = client.send(requestBuilder.build(), HttpResponse.BodyHandlers.ofString())
        return JSONObject(response.body())
    }

    fun chatStream(
        messages: List<ChatMessage>,
        model: String? = null,
        maxTokens: Int = 500,
        temperature: Double = 0.7,
        callback: (String) -> Unit
    ) {
        val modelInfo = resolveModel(model)

        val messagesArray = JSONArray()
        for (msg in messages) {
            messagesArray.put(
                JSONObject()
                    .put("role", msg.role)
                    .put("content", msg.content)
            )
        }

        val payload = JSONObject()
            .put("model", modelInfo.id)
            .put("messages", messagesArray)
            .put("max_tokens", maxTokens)
            .put("temperature", temperature)
            .put("stream", true)

        // Use Java's HttpClient with streaming
        val url = java.net.URL("${modelInfo.endpoint}/chat/completions")
        val connection = url.openConnection() as java.net.HttpURLConnection
        connection.requestMethod = "POST"
        connection.setRequestProperty("Content-Type", "application/json")
        if (apiKey != null) {
            connection.setRequestProperty("Authorization", "Bearer $apiKey")
        }
        connection.doOutput = true
        connection.connectTimeout = timeout.toInt()
        connection.readTimeout = timeout.toInt()

        connection.outputStream.use { os ->
            os.write(payload.toString().toByteArray())
        }

        BufferedReader(InputStreamReader(connection.inputStream)).use { reader ->
            var line: String?
            while (reader.readLine().also { line = it } != null) {
                val currentLine = line ?: continue

                if (currentLine.startsWith("data: ")) {
                    val data = currentLine.substring(6).trim()

                    if (data == "[DONE]") {
                        break
                    }

                    try {
                        val chunk = JSONObject(data)
                        val choices = chunk.optJSONArray("choices")
                        if (choices != null && choices.length() > 0) {
                            val delta = choices.getJSONObject(0).optJSONObject("delta")
                            val content = delta?.optString("content", "")
                            if (content != null && content.isNotEmpty()) {
                                callback(content)
                            }
                        }
                    } catch (e: Exception) {
                        if (debug) {
                            println("[DEBUG] Parse error: ${e.message}")
                        }
                    }
                }
            }
        }
    }

    fun tts(
        text: String,
        voice: String = "alloy",
        model: String = "tts-1",
        responseFormat: String = "mp3"
    ): ByteArray {
        if (ttsEndpointList.isEmpty()) {
            throw IllegalStateException("No TTS endpoints available")
        }

        val endpoint = ttsEndpointList[0]

        val payload = JSONObject()
            .put("model", model)
            .put("voice", voice)
            .put("input", text)
            .put("response_format", responseFormat)

        val requestBuilder = HttpRequest.newBuilder()
            .uri(URI.create("$endpoint/audio/speech"))
            .header("Content-Type", "application/json")
            .POST(HttpRequest.BodyPublishers.ofString(payload.toString()))

        if (apiKey != null) {
            requestBuilder.header("Authorization", "Bearer $apiKey")
        }

        val response = client.send(requestBuilder.build(), HttpResponse.BodyHandlers.ofByteArray())
        return response.body()
    }

    private fun resolveModel(model: String?): ModelInfo {
        if (models.isEmpty()) {
            throw IllegalStateException("No models available")
        }

        if (model == null) {
            return models[0]
        }

        return models.find { it.id == model }
            ?: throw IllegalArgumentException("Model '$model' not found")
    }
}

// Demo when run as application
fun main() {
    println("=== UncloseAI Kotlin Client (with Streaming) ===\n")

    val client = UncloseAI(debug = true)

    val models = client.listModels()
    if (models.isEmpty()) {
        println("ERROR: No models discovered. Set environment variables:")
        println("  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.")
        return
    }

    println("\nDiscovered ${models.size} model(s):")
    for (model in models) {
        println("  - ${model.id} (max_tokens: ${model.maxTokens})")
    }
    println()

    // Non-streaming chat
    println("=== Non-Streaming Chat ===")
    try {
        val response = client.chat(
            messages = listOf(
                ChatMessage("system", "You are a helpful AI assistant."),
                ChatMessage("user", "Explain quantum computing in one sentence.")
            ),
            maxTokens = 100
        )
        val content = response.getJSONArray("choices")
            .getJSONObject(0)
            .getJSONObject("message")
            .getString("content")
        println("Response: $content\n")
    } catch (e: Exception) {
        println("Error: ${e.message}\n")
    }

    // Streaming chat
    println("=== Streaming Chat ===")
    val modelId = if (models.size > 1) models[1].id else null
    println("Model: ${modelId ?: models[0].id}")
    print("Response: ")
    try {
        client.chatStream(
            messages = listOf(
                ChatMessage("system", "You are a coding assistant."),
                ChatMessage("user", "Write a Kotlin function to check if a number is prime")
            ),
            model = modelId,
            maxTokens = 200
        ) { content ->
            print(content)
        }
        println("\n")
    } catch (e: Exception) {
        println("\nError: ${e.message}\n")
    }

    // TTS
    if (client.listModels().isNotEmpty()) {
        println("=== TTS Speech Generation ===")
        try {
            val audioData = client.tts("Hello from UncloseAI Kotlin client! This demonstrates streaming support.")
            File("speech.mp3").writeBytes(audioData)
            println("[OK] Speech file created: speech.mp3 (${audioData.size} bytes)\n")
        } catch (e: Exception) {
            println("[ERROR] TTS Error: ${e.message}\n")
        }
    }

    println("=== Examples Complete ===")
}
