;;; llm.scm - Provider-agnostic LLM chat API client with tool calling support
;;;
;;; BSD-3-Clause License
;;; Copyright (c) 2025, Rolando Abarca

(module llm

(llm/chat
 llm/send
 llm/register-tool!
 llm/get-registered-tools
 llm/get-last-response
 llm/get-cost
 llm/get-tokens
 llm/use-provider
 ;; Simple convenience functions
 llm/generate-text
 llm/generate-image
 llm/generate-image-save
 llm/transcribe-audio)

(import
  scheme
  chicken.base
  chicken.string
  chicken.condition
  chicken.file
  chicken.pathname
  chicken.port
  chicken.io
  srfi-1  ;; lists
  srfi-13 ;; strings
  base64  ;; base64 encoding
  medea   ;; JSON
  llm-provider
  llm-common
  llm-openai
  logger)

(logger/install LLM)

;;; ================================================================
;;; Provider Setup
;;; ================================================================

;; Set OpenAI as the default provider
(current-provider openai-provider)

;; Re-export current-provider as llm/use-provider
(define llm/use-provider current-provider)

;;; ================================================================
;;; File Content Building
;;; ================================================================

;; Build message content for API
;; Returns either a string (text-only) or a vector of content parts (text + file)
;; FILE-PATH: optional local file path
(define (build-message-content text #!key file-path)
  (cond
   ;; No file, return simple text
   ((not file-path)
    text)

   ;; Local file path provided
   (file-path
    (let* ((mime-type (detect-mime-type file-path))
           (is-vision (vision-mime-type? mime-type)))
      (if is-vision
          ;; Image or PDF: embed as base64
          (let* ((base64-data (read-file-base64 file-path))
                 (data-url (conc "data:" mime-type ";base64," base64-data)))
            (if (pdf-mime-type? mime-type)
                ;; PDF: use file format
                (let ((filename (pathname-strip-directory file-path)))
                  `#(((type . "text")
                      (text . ,text))
                     ((type . "file")
                      (file . ((filename . ,filename)
                               (file_data . ,data-url))))))
                ;; Image: use image_url format
                `#(((type . "text")
                    (text . ,text))
                   ((type . "image_url")
                    (image_url . ((url . ,data-url)))))))
          ;; Text file: inline the content
          (let* ((file-content (read-text-file file-path))
                 (filename (pathname-strip-directory file-path)))
            (conc text "\n\n--- File: " filename " ---\n"
                  file-content)))))

   (else text)))

;;; ================================================================
;;; Tool Registry System
;;; ================================================================

;; Global tool registry: ((tool-name . (schema . implementation)) ...)
;; - schema: OpenAI function schema (alist)
;; - implementation: procedure that takes params alist, returns result alist
(define *llm-tools* '())

;; Register a tool in the global registry
;; NAME: symbol - tool identifier (kebab-case like 'get-current-time)
;; SCHEMA: alist - function call schema (use snake_case in function.name)
;; IMPLEMENTATION: procedure (params-alist) -> result-alist
;;
;; Example:
;; (llm/register-tool!
;;   'get-current-time  ;; Register with kebab-case
;;   '((type . "function")
;;     (function . ((name . "get_current_time")  ;; Schema uses snake_case
;;                  (description . "Get the current time")
;;                  (parameters . ((type . "object")
;;                                 (properties . ())
;;                                 (required . #()))))))
;;   (lambda (params)
;;     `((success . #t)
;;       (time . ,(current-seconds)))))
(define (llm/register-tool! name schema implementation)
  (let ((snake-name (kebab->snake name)))
    (set! *llm-tools*
      (alist-update snake-name
                    (cons schema implementation)
                    *llm-tools*))
    (d "[TOOL] Registered: " name " as " snake-name)))

;; Get all registered tool schemas as a vector (for API)
;; TOOL-NAMES: optional list of symbols to filter tools (kebab-case)
;; Returns: vector of tool schemas
(define (llm/get-registered-tools #!optional tool-names)
  (let* ((tools (if tool-names
                    ;; Filter to requested tools (convert to snake_case for lookup)
                    (filter-map (lambda (name)
                                  (let ((entry (alist-ref (kebab->snake name) *llm-tools*)))
                                    (and entry (car entry))))
                                tool-names)
                    ;; Get all tools
                    (map (lambda (entry) (car (cdr entry)))
                         *llm-tools*))))
    (list->vector tools)))

;; Internal: Execute a tool by name (API sends snake_case names)
;; Returns result alist with success/error
(define (execute-tool tool-name arguments)
  (d "[TOOL] Executing: " tool-name)
  (d "[TOOL] Arguments: " (json->string arguments))
  (condition-case
      (let* ((tool-sym (string->symbol tool-name))  ;; API sends "get_current_time"
             (entry (alist-ref tool-sym *llm-tools*)))  ;; Direct lookup with snake_case
        (if entry
            (let* ((implementation (cdr entry))
                   (result (implementation arguments)))
              (d "[TOOL] Result: " (json->string result))
              result)
            (begin
              (e "Unknown tool: " tool-name)
              (e "Registered tools: " (map car *llm-tools*))
              `((success . #f)
                (error . ,(conc "Unknown tool: " tool-name))))))
    [exn ()
         (e "Tool execution exception: " ((condition-property-accessor 'exn 'message) exn))
         `((success . #f)
           (error . ,(->string exn)))]))

;;; ================================================================
;;; Cost Calculation
;;; ================================================================

;; Calculate cost from token counts and model pricing
(define (calculate-cost input-tokens output-tokens provider model-name)
  (let* ((get-pricing (llm-provider-get-model-pricing provider))
         (pricing (get-pricing model-name))
         (input-price (alist-ref 'input-price-per-1m pricing))
         (output-price (alist-ref 'output-price-per-1m pricing)))
    (+ (* (/ input-tokens 1000000.0) input-price)
       (* (/ output-tokens 1000000.0) output-price))))

;;; ================================================================
;;; Response Processing
;;; ================================================================

;; Process API response recursively, handling tool calls
;; Uses provider abstraction for parsing and formatting
;; Returns: (success: #t/f, history: messages, finish-reason: ..., input-tokens: N, output-tokens: N)
;; acc-input/acc-output: accumulated token counts from previous recursive calls
(define (process-response response-data depth messages hooks provider
                          #!optional (acc-input 0) (acc-output 0))
  (d "Processing response at depth " depth)
  (cond
   ((> depth 10)
    `((success . #f)
      (error . "Max recursion depth reached (too many tool calls)")
      (history . ,messages)
      (input-tokens . ,acc-input)
      (output-tokens . ,acc-output)))
   ((not response-data)
    `((success . #f)
      (error . "No response from LLM")
      (history . ,messages)
      (input-tokens . ,acc-input)
      (output-tokens . ,acc-output)))
   (else
    ;; Use provider's parse-response to extract data
    (let* ((parse-response (llm-provider-parse-response provider))
           (parsed (parse-response response-data))
           (success (alist-ref 'success parsed)))
      (if (not success)
          ;; Parse failed
          `((success . #f)
            (error . ,(alist-ref 'error parsed))
            (history . ,messages)
            (input-tokens . ,(+ acc-input (or (alist-ref 'input-tokens parsed) 0)))
            (output-tokens . ,(+ acc-output (or (alist-ref 'output-tokens parsed) 0))))
          ;; Parse succeeded
          (let* ((message (alist-ref 'message parsed))
                 (content (alist-ref 'content parsed))
                 (tool-calls (alist-ref 'tool-calls parsed))
                 (finish-reason (alist-ref 'finish-reason parsed))
                 (input-tokens (or (alist-ref 'input-tokens parsed) 0))
                 (output-tokens (or (alist-ref 'output-tokens parsed) 0))
                 ;; Accumulate tokens
                 (total-input (+ acc-input input-tokens))
                 (total-output (+ acc-output output-tokens))
                 (history (append messages (list message)))
                 ;; Debug logging
                 (_ (d "Finish reason: " finish-reason))
                 (_ (d "Content: " (if (and content (string? content))
                                       (string-take content (min 100 (string-length content)))
                                       "null")))
                 ;; Call on-response-received hook
                 (_ (when (alist-ref 'on-response-received hooks)
                      ((alist-ref 'on-response-received hooks) message))))

            ;; Check if LLM wants to call tools
            (if (and tool-calls (> (vector-length tool-calls) 0))
                ;; Execute tool calls and recurse
                (let* ((extract-tool-calls (llm-provider-extract-tool-calls provider))
                       (format-tool-result (llm-provider-format-tool-result provider))
                       (normalized-calls (extract-tool-calls message))
                       (_ (d "Tool calls: " (length normalized-calls)))
                       (_ (d "Tool calls details: " (map (lambda (tc) (alist-ref 'name tc)) normalized-calls)))
                       (tool-results
                        (map (lambda (tool-call)
                               (let* ((tool-id (alist-ref 'id tool-call))
                                      (tool-name (alist-ref 'name tool-call))
                                      (arguments-str (alist-ref 'arguments tool-call))
                                      (arguments (read-json arguments-str))
                                      (result (execute-tool tool-name arguments)))
                                 ;; Call on-tool-executed hook
                                 (when (alist-ref 'on-tool-executed hooks)
                                   ((alist-ref 'on-tool-executed hooks) tool-name arguments result))
                                 `((tool-call-id . ,tool-id)
                                   (result . ,result))))
                             normalized-calls))
                       ;; Build tool result messages using provider's format
                       (tool-messages
                        (map (lambda (tr)
                               (format-tool-result (alist-ref 'tool-call-id tr)
                                                   (alist-ref 'result tr)))
                             tool-results))
                       ;; Continue conversation with tool results
                       (messages-with-tools (append history tool-messages))
                       ;; Build and send next payload
                       (build-payload (llm-provider-build-payload provider))
                       (call-api (llm-provider-call-api provider))
                       (next-payload (build-payload messages-with-tools
                                                    (alist-ref 'tools hooks)
                                                    (alist-ref 'model hooks)
                                                    (alist-ref 'temperature hooks)
                                                    (alist-ref 'max-tokens hooks)))
                       (next-response (call-api "chat/completions" next-payload)))
                  ;; Recursive call with accumulated tokens
                  (process-response next-response (+ 1 depth) messages-with-tools hooks provider
                                    total-input total-output))

                ;; No tool calls, return final response with token counts
                `((success . #t)
                  (history . ,history)
                  (finish-reason . ,finish-reason)
                  (content . ,content)
                  (input-tokens . ,total-input)
                  (output-tokens . ,total-output)))))))))

;;; ================================================================
;;; Public API
;;; ================================================================

;; Create a new conversation
;; Returns conversation state alist
;;
;; Usage:
;;   (llm/chat)
;;   (llm/chat system: "You are a helpful assistant")
;;   (llm/chat system: "..." tools: '(tool1 tool2))
;;   (llm/chat history: previous-history tools: '(tool1))
;;   (llm/chat on-response-received: (lambda (msg) ...))
;;   (llm/chat on-tool-executed: (lambda (name args result) ...))
;;   (llm/chat provider: anthropic-provider)
(define (llm/chat #!key
                  system
                  tools
                  history
                  model
                  temperature
                  max-tokens
                  on-response-received
                  on-tool-executed
                  provider)
  (let* ((conv-provider (or provider (current-provider)))
         (system-message (if system
                             `((role . "system")
                               (content . ,system))
                             #f))
         (initial-history (if system-message
                              (if (and history (pair? history))
                                  ;; Only prepend system if history doesn't start with one
                                  (if (equal? (alist-ref 'role (car history)) "system")
                                      history
                                      (cons system-message history))
                                  (list system-message))
                              (or history '())))
         (tool-schemas (if tools
                           (llm/get-registered-tools tools)
                           #())))
    ;; Explicitly construct alist with cons to ensure proper dotted pairs
    ;; Initialize cost tracking fields to 0
    (list (cons 'history initial-history)
          (cons 'tools tool-schemas)
          (cons 'tool-names (or tools '()))
          (cons 'model model)
          (cons 'temperature temperature)
          (cons 'max-tokens max-tokens)
          (cons 'on-response-received on-response-received)
          (cons 'on-tool-executed on-tool-executed)
          (cons 'provider conv-provider)
          (cons 'total-input-tokens 0)
          (cons 'total-output-tokens 0)
          (cons 'total-cost 0.0))))

;; Send a message and get response
;; CONVERSATION: conversation state from llm/chat
;; MESSAGE: string to send
;; FILE: optional local file path to attach (image, PDF, or text file)
;; Returns: (values conversation success)
;;   - conversation: updated conversation state
;;   - success: #t if successful, #f if error occurred
;;
;; Usage:
;;   ;; Text only
;;   (let-values ([(conv ok?) (llm/send conv "What is 2+2?")])
;;     (if ok?
;;         (llm/get-last-response conv)  ;; => "4"
;;         (error "LLM call failed")))
;;
;;   ;; With local file
;;   (let-values ([(conv ok?) (llm/send conv "What's in this image?" file: "photo.jpg")])
;;     (if ok? (llm/get-last-response conv) #f))
(define (llm/send conversation message #!key file)
  (let* ((provider (or (alist-ref 'provider conversation) (current-provider)))
         (history (alist-ref 'history conversation))
         (message-content (build-message-content message file-path: file))
         (user-message `((role . "user")
                         (content . ,message-content)))
         (messages (append history (list user-message)))
         (tools (alist-ref 'tools conversation))
         (model (or (alist-ref 'model conversation) *openai-default-model*))
         (temperature (alist-ref 'temperature conversation))
         (max-tokens (alist-ref 'max-tokens conversation))
         ;; Debug: Log tools being sent
         (_ (d "Sending message with " (if tools (vector-length tools) 0) " tools"))
         (_ (when (and tools (> (vector-length tools) 0))
              (d "Available tools: "
                 (map (lambda (t)
                        (alist-ref 'name (alist-ref 'function t)))
                      (vector->list tools)))))
         (hooks `((on-response-received . ,(alist-ref 'on-response-received conversation))
                  (on-tool-executed . ,(alist-ref 'on-tool-executed conversation))
                  (tools . ,tools)
                  (model . ,model)
                  (temperature . ,temperature)
                  (max-tokens . ,max-tokens)))
         ;; Use provider's build-payload and call-api
         (build-payload (llm-provider-build-payload provider))
         (call-api (llm-provider-call-api provider))
         (payload (build-payload messages tools model temperature max-tokens))
         (response-data (call-api "chat/completions" payload))
         (result (process-response response-data 0 messages hooks provider))
         (success (alist-ref 'success result))
         ;; Extract token usage from this send
         (send-input-tokens (or (alist-ref 'input-tokens result) 0))
         (send-output-tokens (or (alist-ref 'output-tokens result) 0))
         ;; Get previous accumulated totals
         (prev-input-tokens (or (alist-ref 'total-input-tokens conversation) 0))
         (prev-output-tokens (or (alist-ref 'total-output-tokens conversation) 0))
         (prev-cost (or (alist-ref 'total-cost conversation) 0.0))
         ;; Calculate new totals
         (new-total-input (+ prev-input-tokens send-input-tokens))
         (new-total-output (+ prev-output-tokens send-output-tokens))
         ;; Calculate cost for this send
         (send-cost (calculate-cost send-input-tokens send-output-tokens provider model))
         (new-total-cost (+ prev-cost send-cost))
         ;; Log cost info
         (_ (i "Tokens: " send-input-tokens " in / " send-output-tokens " out | "
               "Cost: $" (number->string send-cost) " (total: $" (number->string new-total-cost) ")"))
         ;; Update conversation with history and cost tracking
         (updated-conv (alist-update 'total-cost new-total-cost
                         (alist-update 'total-output-tokens new-total-output
                           (alist-update 'total-input-tokens new-total-input
                             (alist-update 'history
                                           (alist-ref 'history result)
                                           conversation))))))
    ;; Return multiple values: conversation and success flag
    (values updated-conv success)))

;; Helper: Get last assistant message from conversation
(define (llm/get-last-response conversation)
  (let* ((history (alist-ref 'history conversation))
         (reversed (reverse history))
         (last-assistant (find (lambda (msg)
                                 (equal? (alist-ref 'role msg) "assistant"))
                               reversed)))
    (if last-assistant
        (alist-ref 'content last-assistant)
        #f)))

;; Get total cost of conversation in USD
(define (llm/get-cost conversation)
  (or (alist-ref 'total-cost conversation) 0.0))

;; Get total token counts as (input . output) pair
(define (llm/get-tokens conversation)
  (cons (or (alist-ref 'total-input-tokens conversation) 0)
        (or (alist-ref 'total-output-tokens conversation) 0)))

;;; ================================================================
;;; Simple Convenience Functions
;;; ================================================================

;; Simple text generation - one-shot prompt to response
;; PROMPT: string - the prompt to send
;; SYSTEM: optional system message
;; MODEL: optional model name
;; TEMPERATURE: optional temperature (0-2)
;; MAX-TOKENS: optional max completion tokens
;; HISTORY: optional list of previous messages
;; PROVIDER: optional provider (defaults to current-provider)
;; Returns: string (the response text) or #f on error
(define (llm/generate-text prompt #!key system model temperature max-tokens history provider)
  (let-values (((conv ok?)
                (llm/send (llm/chat system: system
                                    model: model
                                    temperature: temperature
                                    max-tokens: max-tokens
                                    history: history
                                    provider: provider)
                          prompt)))
    (if ok?
        (llm/get-last-response conv)
        #f)))

;; Generate image using current provider
;; PROMPT: string - the image description
;; MODEL: optional model (e.g., "gpt-image-1", "dall-e-3")
;; SIZE: optional size (e.g., "1024x1024")
;; QUALITY: optional quality (e.g., "high", "medium", "low")
;; OUTPUT-FORMAT: optional format (e.g., "png", "jpeg")
;; N: optional number of images to generate
;; PROVIDER: optional provider (defaults to current-provider)
;; Returns: list of base64-encoded image strings, or #f on error
(define (llm/generate-image prompt #!key model size quality output-format n provider)
  (let* ((prov (or provider (current-provider)))
         (gen-fn (llm-provider-generate-image prov)))
    (if (not gen-fn)
        (error "Provider does not support image generation"
               (llm-provider-name prov))
        (let ((params `(,@(if model `((model . ,model)) '())
                        ,@(if size `((size . ,size)) '())
                        ,@(if quality `((quality . ,quality)) '())
                        ,@(if output-format `((output_format . ,output-format)) '())
                        ,@(if n `((n . ,n)) '()))))
          (gen-fn prompt params)))))

;; Generate image and save to file
;; PROMPT: string - the image description
;; DESTINATION: pair - (type . path) where type is 'file
;; MODEL, SIZE, QUALITY, PROVIDER: same as llm/generate-image
;; Returns: the save path on success, #f on failure
(define (llm/generate-image-save prompt destination #!key model size quality provider)
  (let ((dest-type (car destination))
        (path (cdr destination)))
    (unless (string? path)
      (error "Path should be a string"))
    (case dest-type
      ((file)
       (and-let* ((images (llm/generate-image prompt
                                               model: model
                                               size: size
                                               quality: quality
                                               provider: provider))
                  ((pair? images)))
         (condition-case
             (let ((binary-data (base64-decode (car images))))
               (call-with-output-file path
                 (lambda (port)
                   (display binary-data port))
                 #:binary)
               path)
           [exn ()
                (e "Failed to save image to file: " path)
                (e "Exception message: " ((condition-property-accessor 'exn 'message) exn))
                #f])))
      (else
       (error "Unknown destination type" dest-type)))))

;; Transcribe audio using current provider
;; FILE-PATH: string - path to the audio file
;; MODEL: optional model (e.g., "gpt-4o-transcribe", "whisper-1")
;; RESPONSE-FORMAT: optional format ("text", "json", "verbose_json")
;; LANGUAGE: optional language code (e.g., "en", "es")
;; PROVIDER: optional provider (defaults to current-provider)
;; Returns: string (for "text" format) or alist (for "json" formats), #f on error
(define (llm/transcribe-audio file-path #!key model response-format language provider)
  (let* ((prov (or provider (current-provider)))
         (transcribe-fn (llm-provider-transcribe-audio prov)))
    (if (not transcribe-fn)
        (error "Provider does not support audio transcription"
               (llm-provider-name prov))
        (let ((params `(,@(if model `((model . ,model)) '())
                        ,@(if response-format `((response_format . ,response-format)) '())
                        ,@(if language `((language . ,language)) '()))))
          (transcribe-fn file-path params)))))

) ;; end module