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

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

open System
open System.Net.Http
open System.Text
open System.Text.Json
open System.IO
open System.Collections.Generic
open System.Threading.Tasks

type ModelInfo = {
    Id: string
    Endpoint: string
    MaxTokens: int
}

type Message = {
    Role: string
    Content: string
}

/// UncloseAI Client class
type UncloseAIClient(timeout: int) =
    let httpClient = new HttpClient(Timeout = TimeSpan.FromSeconds(float timeout))
    let models = List<ModelInfo>()
    let ttsEndpoints = List<string>()

    /// Initialize client with auto-discovery
    member this.Init() =
        task {
            printfn "Initializing UncloseAI client..."

            // Discover chat/code models
            let mutable i = 1
            let mutable continueLoop = true
            while continueLoop && i <= 9999 do
                let endpoint = Environment.GetEnvironmentVariable($"MODEL_ENDPOINT_{i}")
                if String.IsNullOrEmpty(endpoint) then
                    continueLoop <- false
                else
                    printfn $"Endpoint {i}: {endpoint}"
                    do! this.DiscoverModelsFromEndpoint(endpoint)
                    i <- i + 1

            // Discover TTS endpoints
            i <- 1
            continueLoop <- true
            while continueLoop && i <= 9999 do
                let endpoint = Environment.GetEnvironmentVariable($"TTS_ENDPOINT_{i}")
                if String.IsNullOrEmpty(endpoint) then
                    continueLoop <- false
                else
                    ttsEndpoints.Add(endpoint)
                    i <- i + 1

            printfn $"Discovered {models.Count} models, {ttsEndpoints.Count} TTS endpoints\n"
        }

    member private this.DiscoverModelsFromEndpoint(endpoint: string) =
        task {
            try
                let! response = httpClient.GetAsync($"{endpoint}/models")
                let! body = response.Content.ReadAsStringAsync()

                use doc = JsonDocument.Parse(body)
                let root = doc.RootElement

                match root.TryGetProperty("data") with
                | (true, data) ->
                    for model in data.EnumerateArray() do
                        let modelId = model.GetProperty("id").GetString()

                        // Skip modelperm-* entries
                        if not (modelId.StartsWith("modelperm-")) then
                            let maxTokens =
                                match model.TryGetProperty("max_model_len") with
                                | (true, prop) -> prop.GetInt32()
                                | (false, _) -> 8192

                            models.Add({ Id = modelId; Endpoint = endpoint; MaxTokens = maxTokens })
                | (false, _) -> ()
            with ex ->
                () // Silently skip failed endpoints
        }

    member this.Models = models :> IReadOnlyList<ModelInfo>
    member this.TtsEndpoints = ttsEndpoints :> IReadOnlyList<string>

    /// Non-streaming chat completion
    member this.Chat(messages: Message[], ?modelIdx: int, ?maxTokens: int, ?temperature: float) =
        task {
            let idx = defaultArg modelIdx 0
            let tokens = defaultArg maxTokens 100
            let temp = defaultArg temperature 0.7

            if idx >= models.Count then
                return Error "Invalid model index"
            else
                let model = models.[idx]
                let url = $"{model.Endpoint}/chat/completions"

                let payload = {|
                    model = model.Id
                    messages = messages |> Array.map (fun m -> {| role = m.Role; content = m.Content |})
                    stream = false
                    max_tokens = tokens
                    temperature = temp
                |}

                let json = JsonSerializer.Serialize(payload)
                let content = new StringContent(json, Encoding.UTF8, "application/json")

                try
                    let! response = httpClient.PostAsync(url, content)
                    let! body = response.Content.ReadAsStringAsync()

                    use doc = JsonDocument.Parse(body)
                    let root = doc.RootElement
                    let responseContent = root.GetProperty("choices").[0].GetProperty("message").GetProperty("content").GetString()
                    return Ok responseContent
                with ex ->
                    return Error ex.Message
        }

    /// Streaming chat completion - yields content chunks
    member this.ChatStream(messages: Message[], ?modelIdx: int, ?maxTokens: int, ?temperature: float) =
        seq {
            let idx = defaultArg modelIdx 0
            let tokens = defaultArg maxTokens 500
            let temp = defaultArg temperature 0.7

            if idx >= models.Count then
                yield Error "Invalid model index"
            else
                let model = models.[idx]
                let url = $"{model.Endpoint}/chat/completions"

                let payload = {|
                    model = model.Id
                    messages = messages |> Array.map (fun m -> {| role = m.Role; content = m.Content |})
                    stream = true
                    max_tokens = tokens
                    temperature = temp
                |}

                let json = JsonSerializer.Serialize(payload)
                let content = new StringContent(json, Encoding.UTF8, "application/json")

                try
                    use request = new HttpRequestMessage(HttpMethod.Post, url, Content = content)
                    let response = httpClient.Send(request, HttpCompletionOption.ResponseHeadersRead)
                    use stream = response.Content.ReadAsStream()
                    use reader = new StreamReader(stream)

                    let mutable line = reader.ReadLine()
                    while not (isNull line) do
                        if line.StartsWith("data: ") then
                            let data = line.Substring(6)
                            if data = "[DONE]" then
                                line <- null
                            else
                                try
                                    use doc = JsonDocument.Parse(data)
                                    let root = doc.RootElement
                                    match root.TryGetProperty("choices") with
                                    | (true, choices) ->
                                        let choice = choices.[0]
                                        match choice.TryGetProperty("delta") with
                                        | (true, delta) ->
                                            match delta.TryGetProperty("content") with
                                            | (true, contentProp) ->
                                                let content = contentProp.GetString()
                                                if not (String.IsNullOrEmpty(content)) then
                                                    yield Ok content
                                            | (false, _) -> ()
                                        | (false, _) -> ()
                                    | (false, _) -> ()
                                with _ ->
                                    () // Skip malformed JSON

                                line <- reader.ReadLine()
                        else
                            line <- reader.ReadLine()
                with ex ->
                    yield Error ex.Message
        }

    /// Text-to-speech generation
    member this.Tts(text: string, ?voice: string, ?outputFile: string) =
        task {
            let voiceStr = defaultArg voice "alloy"
            let fileStr = defaultArg outputFile "/tmp/speech.mp3"

            if ttsEndpoints.Count = 0 then
                return Error "No TTS endpoints available"
            else
                let endpoint = ttsEndpoints.[0]
                let url = $"{endpoint}/audio/speech"

                let payload = {|
                    model = "tts-1"
                    voice = voiceStr
                    input = text
                |}

                let json = JsonSerializer.Serialize(payload)
                let content = new StringContent(json, Encoding.UTF8, "application/json")

                try
                    let! response = httpClient.PostAsync(url, content)
                    let! audioData = response.Content.ReadAsByteArrayAsync()

                    File.WriteAllBytes(fileStr, audioData)
                    return Ok fileStr
                with ex ->
                    return Error ex.Message
        }

    interface IDisposable with
        member this.Dispose() =
            httpClient.Dispose()

/// Demo program showing library usage
[<EntryPoint>]
let main argv =
    printfn "=== UncloseAI F# Client (with Streaming) ===\n"

    use client = new UncloseAIClient(30)

    task {
        // Initialize client
        do! client.Init()

        if client.Models.Count = 0 then
            printfn "ERROR: No models discovered"
            Environment.Exit(1)

        // Non-streaming chat example
        printfn "=== Non-Streaming Chat ==="
        printfn $"Model: {client.Models.[0].Id}"

        let! result = client.Chat([| { Role = "user"; Content = "Explain quantum computing in one sentence" } |])
        match result with
        | Ok response -> printfn $"Response: {response}\n"
        | Error err -> printfn $"Error: {err}\n"

        // Streaming chat example
        let modelIdx = if client.Models.Count >= 2 then 1 else 0
        printfn "=== Streaming Chat ==="
        printfn $"Model: {client.Models.[modelIdx].Id}"
        printf "Response: "

        let messages = [| { Role = "user"; Content = "Write a hello world program in F#" } |]
        for chunk in client.ChatStream(messages, modelIdx) do
            match chunk with
            | Ok content -> printf $"{content}"
            | Error _ -> ()
        printfn "\n"

        // TTS example
        if client.TtsEndpoints.Count > 0 then
            printfn "=== TTS Speech Generation ==="
            printfn "Model: tts-1"

            let! ttsResult = client.Tts("Hello from UncloseAI F# client!", "alloy", "/tmp/speech.mp3")
            match ttsResult with
            | Ok file -> printfn $"Audio saved to {file}"
            | Error err -> printfn $"TTS failed: {err}"

        printfn "\n=== Examples Complete ==="
    } |> Async.AwaitTask |> Async.RunSynchronously

    0
