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


# UncloseAI AWK Library - OpenAI-compatible API client with streaming support
# Compatible with vLLM, Ollama, and OpenAI-compatible endpoints
#
# Library Functions:
#   uncloseai_init()                    - Initialize client with model discovery
#   uncloseai_list_models()             - List discovered models
#   uncloseai_chat(messages, model)     - Non-streaming chat completion
#   uncloseai_chat_stream(messages, model) - Streaming chat completion
#   uncloseai_tts(text, voice, model)   - Text-to-speech generation

# Global client state
# client_models[i,"id|endpoint|max_tokens"] - Discovered models
# client_tts_endpoints[i] - TTS endpoints
# client_model_count - Number of discovered models
# client_tts_count - Number of TTS endpoints

function uncloseai_init(    i, endpoint, cmd, response, model_id, max_tokens) {
    # Initialize client state
    client_model_count = 0
    client_tts_count = 0

    # Discover chat/code models from MODEL_ENDPOINT_N
    for (i = 1; i <= 9999; i++) {
        endpoint = ENVIRON["MODEL_ENDPOINT_" i]
        if (endpoint == "") break

        cmd = "curl -s " endpoint "/models"

        response = ""
        while ((cmd | getline line) > 0) {
            response = response line
        }
        close(cmd)

        # Parse models from JSON response
        # Look for "id":"model-name" patterns
        while (match(response, /"id":"([^"]+)"/, arr)) {
            model_id = arr[1]

            # Skip modelperm entries (they are permission tokens, not models)
            if (index(model_id, "modelperm-") == 1) {
                sub(/"id":"[^"]+"/, "", response)
                continue
            }

            # Extract max_model_len if present (vLLM style)
            if (match(response, /"max_model_len":([0-9]+)/, max_arr)) {
                max_tokens = max_arr[1]
            } else {
                max_tokens = 8192  # Default for Ollama
            }

            client_model_count++
            client_models[client_model_count,"id"] = model_id
            client_models[client_model_count,"endpoint"] = endpoint
            client_models[client_model_count,"max_tokens"] = max_tokens

            # Remove this model from response to find next one
            sub(/"id":"[^"]+"/, "", response)
        }
    }

    # Discover TTS endpoints from TTS_ENDPOINT_N
    for (i = 1; i <= 9999; i++) {
        endpoint = ENVIRON["TTS_ENDPOINT_" i]
        if (endpoint == "") break

        client_tts_count++
        client_tts_endpoints[client_tts_count] = endpoint
    }

    return client_model_count
}

function uncloseai_list_models(    i) {
    # Print all discovered models
    for (i = 1; i <= client_model_count; i++) {
        printf "  - %s (max_tokens: %d)\n", \
            client_models[i,"id"], \
            client_models[i,"max_tokens"]
    }
}

function uncloseai_get_model_idx(model_id,    i) {
    # Get model index by ID or return 1 for first model
    if (model_id == "" && client_model_count > 0) {
        return 1  # Return first model index
    }

    # Search for specific model
    for (i = 1; i <= client_model_count; i++) {
        if (client_models[i,"id"] == model_id) {
            return i
        }
    }

    return 0  # Model not found
}

function uncloseai_chat(messages_json, model_id, max_tokens, temperature,    model_idx, cmd, response, content) {
    # Non-streaming chat completion
    # messages_json: JSON array string like '[{"role":"user","content":"..."}]'
    # Returns: content string

    model_idx = uncloseai_get_model_idx(model_id)
    if (model_idx == 0) {
        return "ERROR: Model not found"
    }

    if (max_tokens == "") max_tokens = 100
    if (temperature == "") temperature = 0.7

    cmd = "curl -s " client_models[model_idx,"endpoint"] "/chat/completions " \
          "-H 'Content-Type: application/json' " \
          "-d '{\"model\":\"" client_models[model_idx,"id"] "\"," \
          "\"messages\":" messages_json "," \
          "\"max_tokens\":" max_tokens "," \
          "\"temperature\":" temperature "," \
          "\"stream\":false}'"

    # Execute curl and capture response
    response = ""
    while ((cmd | getline line) > 0) {
        response = response line
    }
    close(cmd)

    # Extract content field from JSON
    if (match(response, /"content":"([^"\\]*(\\.[^"\\]*)*)"/, arr)) {
        content = arr[1]
        # Unescape common JSON escape sequences
        gsub(/\\n/, "\n", content)
        gsub(/\\"/, "\"", content)
        gsub(/\\\\/, "\\", content)
        return content
    }

    return "ERROR: No response content"
}

function uncloseai_chat_stream(messages_json, model_id, max_tokens, temperature,    model_idx, cmd) {
    # Streaming chat completion with SSE parsing
    # Prints content chunks as they arrive
    # messages_json: JSON array string like '[{"role":"user","content":"..."}]'

    model_idx = uncloseai_get_model_idx(model_id)
    if (model_idx == 0) {
        print "ERROR: Model not found"
        return 0
    }

    if (max_tokens == "") max_tokens = 500
    if (temperature == "") temperature = 0.7

    # Use curl with --no-buffer for line-by-line streaming
    cmd = "curl -s --no-buffer " client_models[model_idx,"endpoint"] "/chat/completions " \
          "-H 'Content-Type: application/json' " \
          "-d '{\"model\":\"" client_models[model_idx,"id"] "\"," \
          "\"messages\":" messages_json "," \
          "\"max_tokens\":" max_tokens "," \
          "\"temperature\":" temperature "," \
          "\"stream\":true}'"

    # Process SSE stream line by line
    while ((cmd | getline line) > 0) {
        # SSE format: "data: {...}"
        if (match(line, /^data: (.+)$/, arr)) {
            data = arr[1]

            # Check for stream termination
            if (data == "[DONE]") {
                break
            }

            # Extract delta content from streaming chunk
            # Format: {"choices":[{"delta":{"content":"..."}}]}
            if (match(data, /"delta":\{[^}]*"content":"([^"\\]*(\\.[^"\\]*)*)"/, content_arr)) {
                content = content_arr[1]
                # Unescape JSON sequences
                gsub(/\\n/, "\n", content)
                gsub(/\\"/, "\"", content)
                gsub(/\\\\/, "\\", content)
                # Print chunk immediately (no newline for streaming effect)
                printf "%s", content
                fflush()  # Flush output for real-time display
            }
        }
    }
    close(cmd)

    return 1
}

function uncloseai_tts(text, voice, model_name, output_file,    endpoint, cmd, size_cmd, file_size) {
    # Text-to-speech generation
    # Returns: file size in bytes (0 on error)

    if (client_tts_count == 0) {
        print "ERROR: No TTS endpoints available"
        return 0
    }

    endpoint = client_tts_endpoints[1]

    if (voice == "") voice = "alloy"
    if (model_name == "") model_name = "tts-1"
    if (output_file == "") output_file = "/tmp/speech.mp3"

    cmd = "curl -s " endpoint "/audio/speech " \
          "-H 'Content-Type: application/json' " \
          "-d '{\"model\":\"" model_name "\"," \
          "\"voice\":\"" voice "\"," \
          "\"input\":\"" text "\"}' " \
          "-o " output_file

    system(cmd)

    # Check file size
    size_cmd = "stat -f%z " output_file " 2>/dev/null || stat -c%s " output_file " 2>/dev/null"
    size_cmd | getline file_size
    close(size_cmd)

    return file_size + 0  # Convert to number
}

# Demo usage when run as script
BEGIN {
    print "=== UncloseAI AWK Client (with Streaming) ===\n"

    # Initialize client with model discovery
    model_count = uncloseai_init()

    if (model_count == 0) {
        print "ERROR: No models discovered. Set environment variables:"
        print "  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc."
        exit 1
    }

    print "Discovered " model_count " model(s)"
    uncloseai_list_models()
    print ""

    # Non-streaming chat example
    print "=== Non-Streaming Chat ==="
    messages = "[{\"role\":\"system\",\"content\":\"You are a helpful AI assistant.\"}," \
               "{\"role\":\"user\",\"content\":\"Explain quantum computing in one sentence.\"}]"

    response = uncloseai_chat(messages, "", 100, 0.7)
    print "Model: " client_models[1,"id"]
    print "Response: " response
    print ""

    # Streaming chat example
    print "=== Streaming Chat ==="

    # Use second model if available, otherwise first
    model_id = ""
    if (client_model_count >= 2) {
        model_id = client_models[2,"id"]
    } else {
        model_id = client_models[1,"id"]
    }

    print "Model: " model_id
    print "Response: "

    messages = "[{\"role\":\"system\",\"content\":\"You are a coding assistant.\"}," \
               "{\"role\":\"user\",\"content\":\"Write a hello world function in AWK.\"}]"

    uncloseai_chat_stream(messages, model_id, 200, 0.7)
    print "\n"

    # TTS example
    if (client_tts_count > 0) {
        print "=== TTS Speech Generation ==="

        text = "Hello from UncloseAI AWK client! This demonstrates text to speech with streaming support."
        output_file = "/tmp/speech.mp3"

        file_size = uncloseai_tts(text, "alloy", "tts-1", output_file)

        if (file_size > 0) {
            printf "[OK] Speech file created: %s (%d bytes)\n\n", output_file, file_size
        } else {
            print "[ERROR] TTS generation failed\n"
        }
    }

    print "=== Examples Complete ==="
    exit
}
