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

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

library(httr)
library(jsonlite)
library(R6)

UncloseAI <- R6Class("UncloseAI",
  public = list(
    models = NULL,
    tts_endpoints = NULL,
    api_key = NULL,
    timeout = NULL,
    debug = NULL,

    initialize = function(endpoints = NULL, tts_endpoints = NULL, api_key = NULL, timeout = 30, debug = FALSE) {
      self$models <- list()
      self$tts_endpoints <- character(0)
      self$api_key <- api_key
      self$timeout <- timeout
      self$debug <- debug

      if (is.null(endpoints)) {
        endpoints <- private$discover_endpoints_from_env("MODEL_ENDPOINT")
      }
      if (is.null(tts_endpoints)) {
        tts_endpoints <- private$discover_endpoints_from_env("TTS_ENDPOINT")
      }

      if (self$debug) {
        cat("[DEBUG] Initialized with", length(endpoints), "endpoint(s)\n")
      }

      private$discover_models(endpoints)
      self$tts_endpoints <- tts_endpoints
    },

    list_models = function() {
      return(self$models)
    },

    chat = function(messages, model = NULL, max_tokens = 100, temperature = 0.7) {
      model_info <- private$resolve_model(model)

      payload <- list(
        model = model_info$id,
        messages = messages,
        max_tokens = max_tokens,
        temperature = temperature
      )

      response <- private$http_request(paste0(model_info$endpoint, "/chat/completions"), "POST", payload)
      return(response)
    },

    chat_stream = function(messages, callback, model = NULL, max_tokens = 500, temperature = 0.7) {
      model_info <- private$resolve_model(model)

      payload <- list(
        model = model_info$id,
        messages = messages,
        max_tokens = max_tokens,
        temperature = temperature,
        stream = TRUE
      )

      url <- paste0(model_info$endpoint, "/chat/completions")

      buffer <- ""
      handle_chunk <- function(chunk) {
        buffer <<- paste0(buffer, rawToChar(chunk))

        lines <- strsplit(buffer, "\n")[[1]]
        if (length(lines) > 0) {
          buffer <<- lines[length(lines)]
          lines <- lines[-length(lines)]

          for (line in lines) {
            if (grepl("^data: ", line)) {
              data <- sub("^data: ", "", line)
              if (trimws(data) == "[DONE]") break

              tryCatch({
                parsed <- fromJSON(data, simplifyVector = FALSE)
                callback(parsed)
              }, error = function(e) {
                if (self$debug) {
                  cat("[DEBUG] Parse error:", conditionMessage(e), "\n")
                }
              })
            }
          }
        }
        TRUE
      }

      headers <- add_headers("Content-Type" = "application/json")
      if (!is.null(self$api_key)) {
        headers <- add_headers("Content-Type" = "application/json",
                             "Authorization" = paste("Bearer", self$api_key))
      }

      POST(url,
           body = toJSON(payload, auto_unbox = TRUE),
           headers,
           timeout(self$timeout),
           write_stream(handle_chunk))
    },

    tts = function(text, voice = "alloy", model = "tts-1") {
      if (length(self$tts_endpoints) == 0) {
        stop("No TTS endpoints available")
      }

      payload <- list(
        model = model,
        voice = voice,
        input = text
      )

      response <- POST(
        paste0(self$tts_endpoints[1], "/audio/speech"),
        body = toJSON(payload, auto_unbox = TRUE),
        add_headers("Content-Type" = "application/json"),
        timeout(self$timeout)
      )

      if (http_error(response)) {
        stop(http_status(response)$message)
      }

      return(content(response, as = "raw"))
    }
  ),

  private = list(
    discover_endpoints_from_env = function(prefix) {
      endpoints <- character(0)
      for (i in 1:9999) {
        endpoint <- Sys.getenv(paste0(prefix, "_", i), unset = NA)
        if (is.na(endpoint)) break
        endpoints <- c(endpoints, endpoint)
      }
      return(endpoints)
    },

    discover_models = function(endpoints) {
      for (endpoint in endpoints) {
        if (self$debug) {
          cat("[DEBUG] Discovering from:", endpoint, "\n")
        }

        tryCatch({
          response <- GET(paste0(endpoint, "/models"), timeout(10))
          result <- content(response, as = "parsed", type = "application/json")

          for (model in result$data) {
            self$models[[length(self$models) + 1]] <- list(
              id = model$id,
              endpoint = endpoint,
              max_tokens = if (!is.null(model$max_model_len)) model$max_model_len else 8192
            )

            if (self$debug) {
              cat("[DEBUG] Discovered:", model$id, "\n")
            }
          }
        }, error = function(e) {
          if (self$debug) {
            cat("[DEBUG] Error:", conditionMessage(e), "\n")
          }
        })
      }
    },

    resolve_model = function(model) {
      if (length(self$models) == 0) {
        stop("No models available")
      }

      if (is.null(model)) {
        return(self$models[[1]])
      }

      for (m in self$models) {
        if (m$id == model) {
          return(m)
        }
      }

      stop(paste("Model '", model, "' not found", sep = ""))
    },

    http_request = function(url, method = "GET", payload = NULL) {
      headers <- list()
      if (!is.null(self$api_key)) {
        headers <- add_headers("Authorization" = paste("Bearer", self$api_key))
      }

      if (method == "GET") {
        response <- GET(url, headers, timeout(self$timeout))
      } else if (method == "POST") {
        headers <- add_headers("Content-Type" = "application/json")
        if (!is.null(self$api_key)) {
          headers <- add_headers("Content-Type" = "application/json",
                               "Authorization" = paste("Bearer", self$api_key))
        }
        response <- POST(url, body = toJSON(payload, auto_unbox = TRUE), headers, timeout(self$timeout))
      }

      if (http_error(response)) {
        stop(http_status(response)$message)
      }

      return(content(response, as = "parsed", type = "application/json"))
    }
  )
)

# Demo when run as script
if (!interactive()) {
  cat("=== UncloseAI R Client (with Streaming) ===\n\n")

  client <- UncloseAI$new(debug = TRUE)

  if (length(client$list_models()) == 0) {
    cat("ERROR: No models discovered. Set environment variables:\n")
    cat("  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.\n")
    quit(status = 1)
  }

  models <- client$list_models()
  cat("\nDiscovered", length(models), "model(s):\n")
  for (m in models) {
    cat("  -", m$id, "(max_tokens:", m$max_tokens, ")\n")
  }
  cat("\n")

  # Non-streaming chat
  cat("=== Non-Streaming Chat ===\n")
  response <- client$chat(list(
    list(role = "system", content = "You are a helpful AI assistant."),
    list(role = "user", content = "Explain quantum computing in one sentence.")
  ))
  cat("Response:", response$choices[[1]]$message$content, "\n\n")

  # Streaming chat
  cat("=== Streaming Chat ===\n")
  model_id <- if (length(models) > 1) models[[2]]$id else NULL
  cat("Model:", if (!is.null(model_id)) model_id else models[[1]]$id, "\n")
  cat("Response: ")

  client$chat_stream(list(
    list(role = "system", content = "You are a coding assistant."),
    list(role = "user", content = "Write an R function to check if a number is prime")
  ), function(chunk) {
    content <- chunk$choices[[1]]$delta$content
    if (!is.null(content)) {
      cat(content)
    }
  }, model = model_id, max_tokens = 200)

  cat("\n\n")

  # TTS
  if (length(client$tts_endpoints) > 0) {
    cat("=== TTS Speech Generation ===\n")
    tryCatch({
      audio_data <- client$tts("Hello from UncloseAI R client! This demonstrates streaming support.")
      writeBin(audio_data, "speech.mp3")
      cat("[OK] Speech file created: speech.mp3 (", length(audio_data), " bytes)\n\n", sep = "")
    }, error = function(e) {
      cat("[ERROR] TTS Error:", conditionMessage(e), "\n\n")
    })
  }

  cat("=== Examples Complete ===\n")
}
