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

(ns uncloseai
  "UncloseAI Clojure Library - OpenAI-compatible API client with streaming support
   Compatible with vLLM, Ollama, and OpenAI-compatible endpoints"
  (:require [clj-http.client :as client]
            [cheshire.core :as json]
            [clojure.string :as str]))

;; Client record for managing models and endpoints
(defrecord UncloseAIClient [models tts-endpoints])

(defn filter-modelperm
  "Filter out modelperm entries from model list"
  [models]
  (remove #(str/starts-with? (:id %) "modelperm-") models))

(defn discover-models-from-endpoint
  "Discover models from a single endpoint"
  [endpoint]
  (try
    (let [response (client/get (str endpoint "/models") {:as :json})
          body (:body response)
          model-list (:data body)]
      (->> model-list
           (map (fn [model]
                  {:id (:id model)
                   :endpoint endpoint
                   :max-tokens (or (:max_model_len model) 8192)}))
           (filter-modelperm)))
    (catch Exception e
      (println "Warning: Failed to discover models from" endpoint ":" (.getMessage e))
      [])))

(defn discover-tts-endpoints
  "Discover TTS endpoints from environment variables"
  []
  (loop [i 1
         endpoints []]
    (if-let [endpoint (System/getenv (str "TTS_ENDPOINT_" i))]
      (recur (inc i) (conj endpoints endpoint))
      endpoints)))

(defn init-client
  "Initialize UncloseAI client with model discovery from environment variables"
  ([]
   (init-client nil nil))
  ([model-endpoints tts-endpoints]
   (let [model-eps (or model-endpoints
                       (loop [i 1 eps []]
                         (if-let [ep (System/getenv (str "MODEL_ENDPOINT_" i))]
                           (recur (inc i) (conj eps ep))
                           eps)))
         tts-eps (or tts-endpoints (discover-tts-endpoints))
         discovered-models (mapcat discover-models-from-endpoint model-eps)]
     (->UncloseAIClient discovered-models tts-eps))))

(defn list-models
  "List all discovered models"
  [client]
  (:models client))

(defn get-model
  "Get model by ID or return first model if ID is nil"
  [client model-id]
  (if model-id
    (first (filter #(= (:id %) model-id) (:models client)))
    (first (:models client))))

(defn chat
  "Non-streaming chat completion

   Args:
     client: UncloseAI client instance
     messages: Vector of message maps with :role and :content
     options: Map with optional :model-id, :max-tokens, :temperature"
  ([client messages]
   (chat client messages {}))
  ([client messages {:keys [model-id max-tokens temperature]
                     :or {max-tokens 100 temperature 0.7}}]
   (let [model (get-model client model-id)]
     (if-not model
       (throw (ex-info "Model not found" {:model-id model-id}))
       (let [url (str (:endpoint model) "/chat/completions")
             payload {:model (:id model)
                      :messages messages
                      :max_tokens max-tokens
                      :temperature temperature
                      :stream false}
             response (client/post url
                                   {:content-type :json
                                    :body (json/generate-string payload)
                                    :as :json})
             body (:body response)]
         {:model (:id model)
          :content (get-in body [:choices 0 :message :content])
          :response body})))))

(defn parse-sse-line
  "Parse a single SSE line and extract content"
  [line]
  (when (str/starts-with? line "data: ")
    (let [data (subs line 6)]
      (when-not (= data "[DONE]")
        (try
          (let [parsed (json/parse-string data true)]
            (get-in parsed [:choices 0 :delta :content]))
          (catch Exception e
            nil))))))

(defn chat-stream
  "Streaming chat completion using Server-Sent Events

   Returns a lazy sequence of content chunks

   Args:
     client: UncloseAI client instance
     messages: Vector of message maps with :role and :content
     options: Map with optional :model-id, :max-tokens, :temperature"
  ([client messages]
   (chat-stream client messages {}))
  ([client messages {:keys [model-id max-tokens temperature]
                     :or {max-tokens 500 temperature 0.7}}]
   (let [model (get-model client model-id)]
     (if-not model
       (throw (ex-info "Model not found" {:model-id model-id}))
       (let [url (str (:endpoint model) "/chat/completions")
             payload {:model (:id model)
                      :messages messages
                      :max_tokens max-tokens
                      :temperature temperature
                      :stream true}
             response (client/post url
                                   {:content-type :json
                                    :body (json/generate-string payload)
                                    :as :stream})
             stream (:body response)]
         (->> (line-seq (clojure.java.io/reader stream))
              (keep parse-sse-line)
              (remove nil?)))))))

(defn tts
  "Text-to-speech generation

   Args:
     client: UncloseAI client instance
     text: Input text to convert to speech
     options: Map with optional :voice, :model, :output-file"
  ([client text]
   (tts client text {}))
  ([client text {:keys [voice model output-file]
                 :or {voice "alloy" model "tts-1" output-file "/tmp/speech.mp3"}}]
   (if (empty? (:tts-endpoints client))
     (throw (ex-info "No TTS endpoints available" {}))
     (let [endpoint (first (:tts-endpoints client))
           url (str endpoint "/audio/speech")
           payload {:model model
                    :voice voice
                    :input text}
           response (client/post url
                                 {:content-type :json
                                  :body (json/generate-string payload)
                                  :as :byte-array})
           audio-data (:body response)]
       (with-open [out (clojure.java.io/output-stream output-file)]
         (.write out audio-data))
       {:file output-file
        :size (count audio-data)}))))

;; Demo usage when run as script
(defn -main [& args]
  (println "=== UncloseAI Clojure Client (with Streaming) ===")
  (println)

  ;; Initialize client with auto-discovery
  (let [client (init-client)]

    (when (empty? (:models client))
      (println "ERROR: No models discovered. Set environment variables:")
      (println "  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.")
      (System/exit 1))

    (println (str "Discovered " (count (:models client)) " model(s)"))
    (doseq [model (:models client)]
      (println (str "  - " (:id model) " (max_tokens: " (:max-tokens model) ")")))
    (println)

    ;; Non-streaming chat example
    (println "=== Non-Streaming Chat ===")
    (try
      (let [result (chat client
                         [{:role "system" :content "You are a helpful AI assistant."}
                          {:role "user" :content "Explain quantum computing in one sentence."}])]
        (println "Model:" (:model result))
        (println "Response:" (:content result)))
      (catch Exception e
        (println "Error:" (.getMessage e))))
    (println)

    ;; Streaming chat example
    (println "=== Streaming Chat ===")
    (let [model-id (if (>= (count (:models client)) 2)
                     (:id (nth (:models client) 1))
                     nil)]
      (println "Model:" (or model-id (:id (first (:models client)))))
      (print "Response: ")
      (flush)
      (try
        (doseq [chunk (chat-stream client
                                   [{:role "system" :content "You are a coding assistant."}
                                    {:role "user" :content "Write a hello world function in Clojure."}]
                                   {:model-id model-id :max-tokens 200})]
          (print chunk)
          (flush))
        (println)
        (catch Exception e
          (println "\nError:" (.getMessage e)))))
    (println)

    ;; TTS example
    (when-not (empty? (:tts-endpoints client))
      (println "=== TTS Speech Generation ===")
      (try
        (let [result (tts client "Hello from UncloseAI Clojure client! This demonstrates text to speech with streaming support.")]
          (println (str "[OK] Speech file created: " (:file result) " (" (:size result) " bytes)")))
        (catch Exception e
          (println "Error:" (.getMessage e))))
      (println))

    (println "=== Examples Complete ===")))
