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


;;; uncloseai - Common Lisp client library for OpenAI-compatible APIs
;;; Supports streaming and non-streaming chat, model discovery, and TTS
;;; Compatible with vLLM, Ollama, and OpenAI-compatible endpoints
;;;
;;; Uses Dexador HTTP client and Jonathan JSON library

(load "~/quicklisp/setup.lisp")
(ql:quickload '(:dexador :jonathan) :silent t)

;;; Data structures

(defclass uncloseai ()
  ((models :initform nil :accessor models
           :documentation "List of discovered models with metadata")
   (tts-endpoints :initform nil :accessor tts-endpoints
                  :documentation "List of TTS endpoint URLs")
   (api-key :initform nil :initarg :api-key :accessor api-key
            :documentation "Optional API key for authentication")
   (request-timeout :initform 30 :initarg :request-timeout :accessor request-timeout
            :documentation "Request timeout in seconds")))

(defun make-model-info (id endpoint max-tokens)
  "Create a model info plist"
  (list :id id :endpoint endpoint :max-tokens max-tokens))

;;; Environment discovery

(defun discover-env-endpoints (prefix)
  "Discover endpoints from environment variables like PREFIX_1, PREFIX_2, ..."
  (loop for i from 1 to 9999
        for var-name = (format nil "~A_~D" prefix i)
        for endpoint = (uiop:getenv var-name)
        while endpoint
        collect endpoint))

;;; Model discovery

(defun discover-models-from-endpoint (client endpoint)
  "Discover available models from an endpoint"
  (handler-case
      (let* ((url (concatenate 'string endpoint "/models"))
             (headers (when (api-key client)
                       (list (cons "Authorization"
                                   (format nil "Bearer ~A" (api-key client))))))
             (response (dex:get url :headers headers))
             (parsed (jonathan:parse response))
             (data (getf parsed :|data|))
             ;; Coerce to list to handle both vector (vLLM) and list (Ollama) responses
             (models-list (coerce data 'list)))
        (dolist (model models-list)
          (let ((model-id (getf model :|id|))
                (max-tokens (or (getf model :|max_model_len|) 8192)))
            (push (make-model-info model-id endpoint max-tokens)
                  (models client)))))
    (error (e)
      (format t "Warning: Failed to discover models from ~A: ~A~%" endpoint e))))

(defun initialize-client (client model-endpoints tts-endpoints)
  "Initialize client with endpoint discovery and model detection"
  ;; Discover endpoints from environment if not provided
  (let ((model-eps (or model-endpoints (discover-env-endpoints "MODEL_ENDPOINT")))
        (tts-eps (or tts-endpoints (discover-env-endpoints "TTS_ENDPOINT"))))

    ;; Discover models from each endpoint
    (dolist (endpoint model-eps)
      (discover-models-from-endpoint client endpoint))

    ;; Reverse models list (they were pushed in reverse order)
    (setf (models client) (nreverse (models client)))

    ;; Store TTS endpoints
    (setf (tts-endpoints client) tts-eps))

  client)

(defun make-uncloseai (&key model-endpoints tts-endpoints api-key (request-timeout 30))
  "Create a new uncloseai client with auto-discovery from environment variables"
  (let ((client (make-instance 'uncloseai
                              :api-key api-key
                              :request-timeout request-timeout)))
    (initialize-client client model-endpoints tts-endpoints)))

;;; Helper functions

(defun get-model-info (client model-id)
  "Get model info by ID or return first available model"
  (when (null (models client))
    (error "No models available. Check endpoint configuration."))

  (if (null model-id)
      (first (models client))
      (or (find model-id (models client) :key (lambda (m) (getf m :id)) :test #'string=)
          (error "Model '~A' not found in discovered models" model-id))))

(defun make-headers (client &optional (content-type t))
  "Create HTTP headers with optional authorization"
  (let ((headers nil))
    (when content-type
      (push (cons "Content-Type" "application/json") headers))
    (when (api-key client)
      (push (cons "Authorization"
                  (format nil "Bearer ~A" (api-key client)))
            headers))
    headers))

;;; Non-streaming chat completion

(defun chat (client messages &key model (max-tokens 100) (temperature 0.7))
  "Non-streaming chat completion

   Args:
     messages - List of message plists with :role and :content
     model - Model ID (defaults to first available model)
     max-tokens - Maximum tokens in response
     temperature - Sampling temperature

   Returns:
     Response plist with :choices containing the completion"
  (let* ((model-info (get-model-info client model))
         (endpoint (getf model-info :endpoint))
         (model-id (getf model-info :id))
         (url (concatenate 'string endpoint "/chat/completions"))
         (payload (jonathan:to-json
                   (list :|model| model-id
                         :|messages| (coerce messages 'vector)
                         :|max_tokens| max-tokens
                         :|temperature| temperature
                         :|stream| :false)))
         (headers (make-headers client))
         (response (dex:post url
                            :headers headers
                            :content payload)))
    (jonathan:parse response)))

;;; Streaming chat completion

(defun process-sse-line (line)
  "Process a single Server-Sent Event line, return parsed content or nil"
  (when (and line (> (length line) 6) (string= (subseq line 0 6) "data: "))
    (let ((data (subseq line 6)))
      (unless (string= data "[DONE]")
        (handler-case
            (let* ((parsed (jonathan:parse data))
                   ;; Coerce choices to list to handle both vector and list responses
                   (choices (coerce (getf parsed :|choices|) 'list))
                   (delta (when choices
                           (getf (first choices) :|delta|)))
                   (content (when delta (getf delta :|content|))))
              content)
          (error (e) nil))))))

(defun chat-stream (client messages callback &key model (max-tokens 500) (temperature 0.7))
  "Streaming chat completion using Server-Sent Events

   Args:
     messages - List of message plists with :role and :content
     callback - Function to call for each content chunk (receives string)
     model - Model ID (defaults to first available model)
     max-tokens - Maximum tokens in response
     temperature - Sampling temperature"
  (let* ((model-info (get-model-info client model))
         (endpoint (getf model-info :endpoint))
         (model-id (getf model-info :id))
         (url (concatenate 'string endpoint "/chat/completions"))
         (payload (jonathan:to-json
                   (list :|model| model-id
                         :|messages| (coerce messages 'vector)
                         :|max_tokens| max-tokens
                         :|temperature| temperature
                         :|stream| t)))
         (headers (make-headers client))
         (stream (dex:request url
                              :method :post
                              :headers headers
                              :content payload
                              :want-stream t)))
    (unwind-protect
         (loop for line = (read-line stream nil nil)
               while line
               do (let ((content (process-sse-line line)))
                    (when content
                      (funcall callback content))))
      (close stream))))

;;; Text-to-Speech

(defun tts (client text &key (voice "alloy") (model "tts-1") (response-format "mp3"))
  "Generate speech from text

   Args:
     text - Input text to convert to speech
     voice - Voice name (alloy, echo, fable, onyx, nova, shimmer)
     model - TTS model (tts-1 or tts-1-hd)
     response-format - Audio format (mp3, opus, aac, flac)

   Returns:
     Audio data as byte array"
  (when (null (tts-endpoints client))
    (error "No TTS endpoints available"))

  (let* ((endpoint (first (tts-endpoints client)))
         (url (concatenate 'string endpoint "/audio/speech"))
         (payload (jonathan:to-json
                   (list :|model| model
                         :|voice| voice
                         :|input| text
                         :|response_format| response-format)))
         (headers (make-headers client))
         (response (dex:post url
                            :headers headers
                            :content payload
                            :force-binary t)))
    response))

;;; Demo usage

(defun demo-nonstreaming (client)
  "Demonstrate non-streaming chat"
  (format t "=== Non-Streaming Chat ===~%")
  (let* ((response (chat client
                        (list (list :|role| "system"
                                   :|content| "You are a helpful AI assistant.")
                              (list :|role| "user"
                                   :|content| "Explain quantum computing in one sentence."))
                        :max-tokens 100
                        :temperature 0.7))
         (model-used (getf response :|model|))
         ;; Coerce choices to list to handle both vector and list responses
         (choices (coerce (getf response :|choices|) 'list))
         (choice (first choices))
         (message (getf choice :|message|))
         (content (getf message :|content|)))
    (format t "Model: ~A~%" model-used)
    (format t "Response: ~A~%~%" content)))

(defun demo-streaming (client)
  "Demonstrate streaming chat"
  (format t "=== Streaming Chat ===~%")
  (let ((model-id (if (> (length (models client)) 1)
                      (getf (second (models client)) :id)
                      nil)))
    (format t "Model: ~A~%" (or model-id (getf (first (models client)) :id)))
    (format t "Response: ")
    (force-output)

    (chat-stream client
                (list (list :|role| "system"
                           :|content| "You are a coding assistant.")
                      (list :|role| "user"
                           :|content| "Write a Common Lisp function to check if a number is prime"))
                (lambda (content)
                  (format t "~A" content)
                  (force-output))
                :model model-id
                :max-tokens 200
                :temperature 0.7)

    (format t "~%~%")))

(defun demo-tts (client)
  "Demonstrate text-to-speech"
  (when (tts-endpoints client)
    (format t "=== TTS Speech Generation ===~%")
    (let ((audio-data (tts client
                          "Hello from uncloseai Common Lisp client! This demonstrates text to speech with streaming support."
                          :voice "alloy")))
      (with-open-file (out "speech.mp3"
                          :direction :output
                          :if-exists :supersede
                          :if-does-not-exist :create
                          :element-type '(unsigned-byte 8))
        (write-sequence audio-data out))
      (format t "[OK] Speech file created: speech.mp3 (~D bytes)~%~%" (length audio-data)))))

(defun main ()
  "Main demo function"
  (format t "=== uncloseai Common Lisp Client (with Streaming) ===~%~%")

  ;; Initialize client (auto-discovers from environment)
  (let ((client (make-uncloseai)))

    (when (null (models client))
      (format t "ERROR: No models discovered. Set environment variables:~%")
      (format t "  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.~%")
      (uiop:quit 1))

    (format t "Discovered ~D model(s)~%" (length (models client)))
    (dolist (model (models client))
      (format t "  - ~A (max_tokens: ~D)~%"
              (getf model :id)
              (getf model :max-tokens)))
    (format t "~%")

    ;; Run demos
    (handler-case
        (progn
          (demo-nonstreaming client)
          (demo-streaming client)
          (demo-tts client)
          (format t "=== Examples Complete ===~%"))
      (error (e)
        (format t "~%Error: ~A~%" e)))))

;; Main execution
(main)
