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

' UncloseAI - VB.NET client for OpenAI-compatible APIs with streaming support

Imports System
Imports System.Net.Http
Imports System.Text
Imports System.Threading.Tasks
Imports System.Text.Json
Imports System.IO
Imports System.Collections.Generic
Imports System.Linq

Public Class ChatMessage
    Public Property Role As String
    Public Property Content As String
End Class

Public Class ModelInfo
    Public Property Id As String
    Public Property Endpoint As String
    Public Property MaxTokens As Integer
End Class

Public Class UncloseAI
    Private ReadOnly models As List(Of ModelInfo)
    Private ReadOnly ttsEndpoints As List(Of String)
    Private ReadOnly apiKey As String
    Private ReadOnly timeout As Integer
    Private ReadOnly debug As Boolean
    Private ReadOnly httpClient As HttpClient

    Public Sub New(Optional endpoints As List(Of String) = Nothing,
                   Optional ttsEps As List(Of String) = Nothing,
                   Optional key As String = "",
                   Optional tm As Integer = 30,
                   Optional dbg As Boolean = False)
        Me.apiKey = key
        Me.timeout = tm
        Me.debug = dbg
        Me.httpClient = New HttpClient()
        Me.httpClient.Timeout = TimeSpan.FromSeconds(tm)

        ' Discover endpoints from environment if not provided
        Dim modelEnds = If(endpoints Is Nothing OrElse endpoints.Count = 0,
                           DiscoverEnvEndpoints("MODEL_ENDPOINT"),
                           endpoints)

        Dim ttsEnds = If(ttsEps Is Nothing OrElse ttsEps.Count = 0,
                         DiscoverEnvEndpoints("TTS_ENDPOINT"),
                         ttsEps)

        If debug Then
            Console.WriteLine($"[DEBUG] Initialized with {modelEnds.Count} endpoint(s)")
        End If

        Me.models = DiscoverModels(modelEnds).GetAwaiter().GetResult()
        Me.ttsEndpoints = ttsEnds
    End Sub

    Private Shared Function DiscoverEnvEndpoints(prefix As String) As List(Of String)
        Dim endpoints As New List(Of String)()
        For i As Integer = 1 To 9999
            Dim endpoint = Environment.GetEnvironmentVariable($"{prefix}_{i}")
            If String.IsNullOrEmpty(endpoint) Then Exit For
            endpoints.Add(endpoint)
        Next
        Return endpoints
    End Function

    Private Async Function DiscoverModels(endpoints As List(Of String)) As Task(Of List(Of ModelInfo))
        Dim modelList As New List(Of ModelInfo)()

        For Each endpoint In endpoints
            If debug Then
                Console.WriteLine($"[DEBUG] Discovering from: {endpoint}")
            End If

            Try
                Dim response = Await httpClient.GetAsync($"{endpoint}/models")
                Dim json = Await response.Content.ReadAsStringAsync()

                ' Parse JSON to find models
                Dim doc = JsonDocument.Parse(json)
                Dim data = doc.RootElement.GetProperty("data")

                For Each model In data.EnumerateArray()
                    Dim modelId = model.GetProperty("id").GetString()

                    ' Skip permission entries
                    If Not modelId.StartsWith("modelperm-") AndAlso Not modelId.StartsWith("chatcmpl-") Then
                        Dim maxTokens = 8192
                        Dim maxTokensProp As JsonElement
                        If model.TryGetProperty("max_model_len", maxTokensProp) Then
                            maxTokens = maxTokensProp.GetInt32()
                        End If

                        modelList.Add(New ModelInfo With {
                            .Id = modelId,
                            .Endpoint = endpoint,
                            .MaxTokens = maxTokens
                        })

                        If debug Then
                            Console.WriteLine($"[DEBUG] Discovered: {modelId}")
                        End If
                    End If
                Next
            Catch ex As Exception
                If debug Then
                    Console.WriteLine($"[DEBUG] Error: {ex.Message}")
                End If
            End Try
        Next

        Return modelList
    End Function

    Public Function ListModels() As List(Of ModelInfo)
        Return models
    End Function

    Private Function ResolveModel(modelId As String) As ModelInfo
        If models.Count = 0 Then
            Throw New Exception("No models available")
        End If

        If String.IsNullOrEmpty(modelId) Then
            Return models(0)
        End If

        For Each model In models
            If model.Id = modelId Then
                Return model
            End If
        Next

        Throw New Exception($"Model '{modelId}' not found")
    End Function

    Public Async Function Chat(messages As List(Of ChatMessage),
                                Optional modelId As String = "",
                                Optional maxTokens As Integer = 100,
                                Optional temperature As Double = 0.7) As Task(Of String)
        Dim model = ResolveModel(modelId)

        Dim request = New With {
            .model = model.Id,
            .messages = messages.Select(Function(m) New With {.role = m.Role, .content = m.Content}).ToArray(),
            .max_tokens = maxTokens,
            .temperature = temperature,
            .stream = False
        }

        Dim jsonRequest = JsonSerializer.Serialize(request)
        Dim content = New StringContent(jsonRequest, Encoding.UTF8, "application/json")

        If Not String.IsNullOrEmpty(apiKey) Then
            httpClient.DefaultRequestHeaders.Authorization = New Headers.AuthenticationHeaderValue("Bearer", apiKey)
        End If

        Dim response = Await httpClient.PostAsync($"{model.Endpoint}/chat/completions", content)
        Dim result = Await response.Content.ReadAsStringAsync()

        Return result
    End Function

    Public Async Function ChatStream(messages As List(Of ChatMessage),
                                      callback As Action(Of String),
                                      Optional modelId As String = "",
                                      Optional maxTokens As Integer = 500,
                                      Optional temperature As Double = 0.7) As Task
        Dim model = ResolveModel(modelId)

        Dim request = New With {
            .model = model.Id,
            .messages = messages.Select(Function(m) New With {.role = m.Role, .content = m.Content}).ToArray(),
            .max_tokens = maxTokens,
            .temperature = temperature,
            .stream = True
        }

        Dim jsonRequest = JsonSerializer.Serialize(request)
        Dim content = New StringContent(jsonRequest, Encoding.UTF8, "application/json")

        If Not String.IsNullOrEmpty(apiKey) Then
            httpClient.DefaultRequestHeaders.Authorization = New Headers.AuthenticationHeaderValue("Bearer", apiKey)
        End If

        Try
            Dim httpRequest = New HttpRequestMessage(HttpMethod.Post, $"{model.Endpoint}/chat/completions")
            httpRequest.Content = content
            Using response = Await httpClient.SendAsync(httpRequest, HttpCompletionOption.ResponseHeadersRead)
                Using stream = Await response.Content.ReadAsStreamAsync()
                    Using reader = New StreamReader(stream)
                        While Not reader.EndOfStream
                            Dim line = Await reader.ReadLineAsync()

                            If Not String.IsNullOrEmpty(line) AndAlso line.StartsWith("data: ") Then
                                Dim data = line.Substring(6).Trim()
                                If data = "[DONE]" Then Exit While

                                Try
                                    Dim doc = JsonDocument.Parse(data)
                                    Dim choices = doc.RootElement.GetProperty("choices")

                                    If choices.GetArrayLength() > 0 Then
                                        Dim choice = choices(0)
                                        If choice.TryGetProperty("delta", Nothing) Then
                                            Dim delta = choice.GetProperty("delta")
                                            If delta.TryGetProperty("content", Nothing) Then
                                                Dim contentText = delta.GetProperty("content").GetString()
                                                If Not String.IsNullOrEmpty(contentText) Then
                                                    callback(contentText)
                                                End If
                                            End If
                                        End If
                                    End If
                                Catch
                                    ' Ignore parse errors
                                End Try
                            End If
                        End While
                    End Using
                End Using
            End Using
        Catch ex As Exception
            If debug Then
                Console.WriteLine($"[DEBUG] Stream error: {ex.Message}")
            End If
        End Try
    End Function

    Public Async Function Tts(text As String,
                               Optional voice As String = "alloy",
                               Optional model As String = "tts-1",
                               Optional format As String = "mp3") As Task(Of Byte())
        If ttsEndpoints.Count = 0 Then
            Throw New Exception("No TTS endpoints available")
        End If

        Dim endpoint = ttsEndpoints(0)

        Dim request = New With {
            .model = model,
            .voice = voice,
            .input = text,
            .response_format = format
        }

        Dim jsonRequest = JsonSerializer.Serialize(request)
        Dim content = New StringContent(jsonRequest, Encoding.UTF8, "application/json")

        If Not String.IsNullOrEmpty(apiKey) Then
            httpClient.DefaultRequestHeaders.Authorization = New Headers.AuthenticationHeaderValue("Bearer", apiKey)
        End If

        Dim response = Await httpClient.PostAsync($"{endpoint}/audio/speech", content)
        Dim audioBytes = Await response.Content.ReadAsByteArrayAsync()

        Return audioBytes
    End Function
End Class

' Demo when run as application
Module Program
    Sub Main(args As String())
        MainAsync(args).GetAwaiter().GetResult()
    End Sub

    Private Async Function MainAsync(args As String()) As Task
        Console.WriteLine("=== UncloseAI VB.NET Client (with Streaming) ===")
        Console.WriteLine()

        Dim client = New UncloseAI(dbg:=True)

        If client.ListModels().Count = 0 Then
            Console.WriteLine("ERROR: No models discovered. Set environment variables:")
            Console.WriteLine("  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.")
            Return
        End If

        Console.WriteLine()
        Console.WriteLine($"Discovered {client.ListModels().Count} model(s):")
        For Each model In client.ListModels()
            Console.WriteLine($"  - {model.Id} (max_tokens: {model.MaxTokens})")
        Next
        Console.WriteLine()

        ' Non-streaming chat
        Console.WriteLine("=== Non-Streaming Chat ===")
        Try
            Dim messages = New List(Of ChatMessage) From {
                New ChatMessage With {.Role = "system", .Content = "You are a helpful AI assistant."},
                New ChatMessage With {.Role = "user", .Content = "Explain quantum computing in one sentence."}
            }

            Dim response = Await client.Chat(messages)
            Dim doc = JsonDocument.Parse(response)
            Dim content = doc.RootElement.GetProperty("choices")(0).GetProperty("message").GetProperty("content").GetString()
            Console.WriteLine($"Response: {content}")
            Console.WriteLine()
        Catch ex As Exception
            Console.WriteLine($"Error: {ex.Message}")
            Console.WriteLine()
        End Try

        ' Streaming chat
        Console.WriteLine("=== Streaming Chat ===")
        Dim modelId = If(client.ListModels().Count > 1, client.ListModels()(1).Id, "")
        Dim modelName = If(String.IsNullOrEmpty(modelId), client.ListModels()(0).Id, modelId)
        Console.WriteLine($"Model: {modelName}")
        Console.Write("Response: ")

        Dim streamMessages = New List(Of ChatMessage) From {
            New ChatMessage With {.Role = "system", .Content = "You are a coding assistant."},
            New ChatMessage With {.Role = "user", .Content = "Write a VB.NET function to check if a number is prime"}
        }

        Await client.ChatStream(streamMessages, Sub(content) Console.Write(content), modelId, 200)
        Console.WriteLine()
        Console.WriteLine()

        ' TTS
        Console.WriteLine("=== TTS Speech Generation ===")
        Try
            Dim audioData = Await client.Tts("Hello from UncloseAI VB.NET client! This demonstrates streaming support.")
            File.WriteAllBytes("speech.mp3", audioData)
            Console.WriteLine($"✓ Speech file created: speech.mp3 ({audioData.Length} bytes)")
            Console.WriteLine()
        Catch ex As Exception
            Console.WriteLine($"✗ TTS Error: {ex.Message}")
            Console.WriteLine()
        End Try

        Console.WriteLine("=== Examples Complete ===")
    End Function
End Module
