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

# UncloseAI - PowerShell client for OpenAI-compatible APIs with streaming support

class UncloseAI {
    [System.Collections.ArrayList]$Models
    [System.Collections.ArrayList]$TtsEndpoints
    [string]$ApiKey
    [int]$Timeout
    [bool]$Debug

    UncloseAI([hashtable]$Config) {
        $this.Models = [System.Collections.ArrayList]::new()
        $this.TtsEndpoints = [System.Collections.ArrayList]::new()
        $this.ApiKey = $Config.ApiKey
        $this.Timeout = if ($Config.Timeout) { $Config.Timeout } else { 30 }
        $this.Debug = if ($Config.Debug) { $Config.Debug } else { $false }

        $modelEndpoints = if ($Config.Endpoints) { $Config.Endpoints } else { $this.DiscoverEndpointsFromEnv("MODEL_ENDPOINT") }
        $ttsEps = if ($Config.TtsEndpoints) { $Config.TtsEndpoints } else { $this.DiscoverEndpointsFromEnv("TTS_ENDPOINT") }

        if ($this.Debug) {
            Write-Host "[DEBUG] Initialized with $($modelEndpoints.Count) endpoint(s)"
        }

        $this.DiscoverModels($modelEndpoints)

        # Ensure TtsEndpoints is ArrayList even if $ttsEps is a single string
        if ($ttsEps -is [System.Collections.ArrayList]) {
            $this.TtsEndpoints = $ttsEps
        } else {
            foreach ($ep in $ttsEps) {
                $this.TtsEndpoints.Add($ep) | Out-Null
            }
        }
    }

    [System.Collections.ArrayList] ListModels() {
        return $this.Models
    }

    [object] Chat([array]$Messages, [hashtable]$Options) {
        $modelInfo = $this.ResolveModel($Options.Model)

        $payload = @{
            model = $modelInfo.id
            messages = $Messages
            max_tokens = if ($Options.MaxTokens) { $Options.MaxTokens } else { 100 }
            temperature = if ($Options.Temperature) { $Options.Temperature } else { 0.7 }
        }

        $response = $this.HttpRequest("$($modelInfo.endpoint)/chat/completions", "POST", $payload)
        return $response
    }

    [void] ChatStream([array]$Messages, [scriptblock]$Callback, [hashtable]$Options) {
        $modelInfo = $this.ResolveModel($Options.Model)

        $payload = @{
            model = $modelInfo.id
            messages = $Messages
            max_tokens = if ($Options.MaxTokens) { $Options.MaxTokens } else { 500 }
            temperature = if ($Options.Temperature) { $Options.Temperature } else { 0.7 }
            stream = $true
        } | ConvertTo-Json -Depth 10

        $url = "$($modelInfo.endpoint)/chat/completions"

        $headers = @{
            "Content-Type" = "application/json"
        }
        if ($this.ApiKey) {
            $headers["Authorization"] = "Bearer $($this.ApiKey)"
        }

        $buffer = ""
        $httpClient = $null
        $stream = $null
        $reader = $null

        try {
            $httpClient = [System.Net.Http.HttpClient]::new()
            $httpClient.Timeout = [System.TimeSpan]::FromSeconds($this.Timeout)

            $content = [System.Net.Http.StringContent]::new($payload, [System.Text.Encoding]::UTF8, "application/json")
            $request = [System.Net.Http.HttpRequestMessage]::new([System.Net.Http.HttpMethod]::Post, $url)
            $request.Content = $content

            foreach ($key in $headers.Keys) {
                $request.Headers.TryAddWithoutValidation($key, $headers[$key]) | Out-Null
            }

            $response = $httpClient.SendAsync($request, [System.Net.Http.HttpCompletionOption]::ResponseHeadersRead).Result
            $stream = $response.Content.ReadAsStreamAsync().Result
            $reader = [System.IO.StreamReader]::new($stream)

            while (-not $reader.EndOfStream) {
                $line = $reader.ReadLine()

                if ($line -match "^data: (.*)$") {
                    $data = $matches[1]
                    if ($data -eq "[DONE]") { break }

                    try {
                        $chunk = $data | ConvertFrom-Json
                        & $Callback $chunk
                    }
                    catch {
                        if ($this.Debug) {
                            Write-Host "[DEBUG] Parse error: $_"
                        }
                    }
                }
            }
        }
        finally {
            if ($reader) { $reader.Dispose() }
            if ($stream) { $stream.Dispose() }
            if ($httpClient) { $httpClient.Dispose() }
        }
    }

    [byte[]] Tts([string]$Text, [hashtable]$Options) {
        if ($this.TtsEndpoints.Count -eq 0) {
            throw "No TTS endpoints available"
        }

        $payload = @{
            model = if ($Options.Model) { $Options.Model } else { "tts-1" }
            voice = if ($Options.Voice) { $Options.Voice } else { "alloy" }
            input = $Text
        }

        $response = Invoke-RestMethod -Uri "$($this.TtsEndpoints[0])/audio/speech" -Method Post -Body ($payload | ConvertTo-Json) -ContentType "application/json" -TimeoutSec $this.Timeout
        return $response
    }

    hidden [System.Collections.ArrayList] DiscoverEndpointsFromEnv([string]$Prefix) {
        $endpoints = [System.Collections.ArrayList]::new()
        for ($i = 1; $i -lt 10000; $i++) {
            $endpoint = [Environment]::GetEnvironmentVariable("${Prefix}_$i")
            if (-not $endpoint) { break }
            $endpoints.Add($endpoint) | Out-Null
        }
        return $endpoints
    }

    hidden [void] DiscoverModels([array]$Endpoints) {
        foreach ($endpoint in $Endpoints) {
            if ($this.Debug) {
                Write-Host "[DEBUG] Discovering from: $endpoint"
            }

            try {
                $response = Invoke-RestMethod -Uri "$endpoint/models" -TimeoutSec 10
                foreach ($model in $response.data) {
                    $this.Models.Add(@{
                        id = $model.id
                        endpoint = $endpoint
                        max_tokens = if ($model.max_model_len) { $model.max_model_len } else { 8192 }
                    }) | Out-Null

                    if ($this.Debug) {
                        Write-Host "[DEBUG] Discovered: $($model.id)"
                    }
                }
            }
            catch {
                if ($this.Debug) {
                    Write-Host "[DEBUG] Error: $_"
                }
            }
        }
    }

    hidden [hashtable] ResolveModel([string]$Model) {
        if ($this.Models.Count -eq 0) {
            throw "No models available"
        }

        if (-not $Model) {
            return $this.Models[0]
        }

        foreach ($m in $this.Models) {
            if ($m.id -eq $Model) {
                return $m
            }
        }

        throw "Model '$Model' not found"
    }

    hidden [object] HttpRequest([string]$Url, [string]$Method, [hashtable]$Payload) {
        $headers = @{
            "Content-Type" = "application/json"
        }
        if ($this.ApiKey) {
            $headers["Authorization"] = "Bearer $($this.ApiKey)"
        }

        $response = $null
        if ($Method -eq "GET") {
            $response = Invoke-RestMethod -Uri $Url -Method Get -Headers $headers -TimeoutSec $this.Timeout
        }
        elseif ($Method -eq "POST") {
            $body = $Payload | ConvertTo-Json -Depth 10
            $response = Invoke-RestMethod -Uri $Url -Method Post -Body $body -Headers $headers -TimeoutSec $this.Timeout
        }

        return $response
    }
}

# Demo when run as script
if ($MyInvocation.InvocationName -ne '.') {
    Write-Host "=== UncloseAI PowerShell Client (with Streaming) ===`n"

    $client = [UncloseAI]::new(@{ Debug = $true })

    if ($client.ListModels().Count -eq 0) {
        Write-Host "ERROR: No models discovered. Set environment variables:"
        Write-Host "  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc."
        exit 1
    }

    $models = $client.ListModels()
    Write-Host "`nDiscovered $($models.Count) model(s):"
    foreach ($m in $models) {
        Write-Host "  - $($m.id) (max_tokens: $($m.max_tokens))"
    }
    Write-Host ""

    # Non-streaming chat
    Write-Host "=== Non-Streaming Chat ==="
    $response = $client.Chat(@(
        @{ role = "system"; content = "You are a helpful AI assistant." }
        @{ role = "user"; content = "Explain quantum computing in one sentence." }
    ), @{})
    Write-Host "Response: $($response.choices[0].message.content)`n"

    # Streaming chat
    Write-Host "=== Streaming Chat ==="
    $modelId = if ($models.Count -gt 1) { $models[1].id } else { $null }
    Write-Host "Model: $(if ($modelId) { $modelId } else { $models[0].id })"
    Write-Host "Response: " -NoNewline

    $client.ChatStream(@(
        @{ role = "system"; content = "You are a coding assistant." }
        @{ role = "user"; content = "Write a PowerShell function to check if a number is prime" }
    ), {
        param($chunk)
        $content = $chunk.choices[0].delta.content
        if ($content) {
            Write-Host $content -NoNewline
        }
    }, @{ Model = $modelId; MaxTokens = 200 })

    Write-Host "`n"

    # TTS
    if ($client.TtsEndpoints.Count -gt 0) {
        Write-Host "=== TTS Speech Generation ==="
        try {
            $audioData = $client.Tts("Hello from UncloseAI PowerShell client! This demonstrates streaming support.", @{})
            [System.IO.File]::WriteAllBytes("speech.mp3", $audioData)
            $fileSize = (Get-Item "speech.mp3").Length
            Write-Host "[OK] Speech file created: speech.mp3 ($fileSize bytes)`n"
        }
        catch {
            Write-Host "[ERROR] TTS Error: $_`n"
        }
    }

    Write-Host "=== Examples Complete ==="
}
