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

import httpclient, json, strformat, os, strutils, asyncdispatch, streams

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

type
  ModelInfo* = object
    id*: string
    endpoint*: string
    maxTokens*: int

  ChatMessage* = object
    role*: string
    content*: string

  UncloseAI* = ref object
    models: seq[ModelInfo]
    ttsEndpoints: seq[string]
    apiKey: string
    timeout: int
    debug: bool

# Forward declarations
proc discoverEnvEndpoints(self: UncloseAI, prefix: string): seq[string]
proc discoverModels(self: UncloseAI, endpoints: seq[string])

proc newUncloseAI*(
  modelEndpoints: seq[string] = @[],
  ttsEndpoints: seq[string] = @[],
  apiKey: string = "",
  timeout: int = 30000,
  debug: bool = false
): UncloseAI =
  result = UncloseAI(
    models: @[],
    ttsEndpoints: @[],
    apiKey: apiKey,
    timeout: timeout,
    debug: debug
  )

  # Discover endpoints from environment
  let modelEnds = if modelEndpoints.len > 0: modelEndpoints else: result.discoverEnvEndpoints("MODEL_ENDPOINT")
  let ttsEnds = if ttsEndpoints.len > 0: ttsEndpoints else: result.discoverEnvEndpoints("TTS_ENDPOINT")

  if result.debug:
    echo fmt"[DEBUG] Initialized with {modelEnds.len} endpoint(s)"

  result.discoverModels(modelEnds)
  result.ttsEndpoints = ttsEnds

proc discoverEnvEndpoints(self: UncloseAI, prefix: string): seq[string] =
  result = @[]
  for i in 1..<10000:
    let endpoint = getEnv(prefix & "_" & $i)
    if endpoint == "":
      break
    result.add(endpoint)

proc discoverModels(self: UncloseAI, endpoints: seq[string]) =
  for endpoint in endpoints:
    if self.debug:
      echo "[DEBUG] Discovering from: ", endpoint

    try:
      let client = newHttpClient(timeout = 10000)
      let response = client.getContent(endpoint & "/models")
      let jsonData = parseJson(response)

      for model in jsonData["data"]:
        let modelId = model["id"].getStr()
        # Skip permission entries
        if modelId.startsWith("modelperm-") or modelId.startsWith("chatcmpl-"):
          continue

        let maxTokens = if model.hasKey("max_model_len"):
          model["max_model_len"].getInt()
        else:
          8192

        self.models.add(ModelInfo(
          id: modelId,
          endpoint: endpoint,
          maxTokens: maxTokens
        ))

        if self.debug:
          echo "[DEBUG] Discovered: ", modelId
    except:
      if self.debug:
        echo "[DEBUG] Error: ", getCurrentExceptionMsg()

proc listModels*(self: UncloseAI): seq[ModelInfo] =
  return self.models

proc resolveModel(self: UncloseAI, model: string): ModelInfo =
  if self.models.len == 0:
    raise newException(ValueError, "No models available")

  if model == "":
    return self.models[0]

  for m in self.models:
    if m.id == model:
      return m

  raise newException(ValueError, fmt"Model '{model}' not found")

proc chat*(
  self: UncloseAI,
  messages: seq[ChatMessage],
  model: string = "",
  maxTokens: int = 100,
  temperature: float = 0.7
): JsonNode =
  let modelInfo = self.resolveModel(model)

  var messagesJson = newJArray()
  for msg in messages:
    messagesJson.add(%* {"role": msg.role, "content": msg.content})

  let payload = %* {
    "model": modelInfo.id,
    "messages": messagesJson,
    "max_tokens": maxTokens,
    "temperature": temperature,
    "stream": false
  }

  let client = newHttpClient(timeout = self.timeout)
  client.headers = newHttpHeaders({"Content-Type": "application/json"})

  if self.apiKey != "":
    client.headers["Authorization"] = "Bearer " & self.apiKey

  let response = client.request(
    modelInfo.endpoint & "/chat/completions",
    httpMethod = HttpPost,
    body = $payload
  )

  return parseJson(response.body)

proc chatStream*(
  self: UncloseAI,
  messages: seq[ChatMessage],
  model: string = "",
  maxTokens: int = 500,
  temperature: float = 0.7,
  callback: proc(content: string)
) =
  let modelInfo = self.resolveModel(model)

  var messagesJson = newJArray()
  for msg in messages:
    messagesJson.add(%* {"role": msg.role, "content": msg.content})

  let payload = %* {
    "model": modelInfo.id,
    "messages": messagesJson,
    "max_tokens": maxTokens,
    "temperature": temperature,
    "stream": true
  }

  let client = newHttpClient(timeout = self.timeout)
  client.headers = newHttpHeaders({
    "Content-Type": "application/json",
    "Accept": "text/event-stream"
  })

  if self.apiKey != "":
    client.headers["Authorization"] = "Bearer " & self.apiKey

  try:
    let response = client.request(
      modelInfo.endpoint & "/chat/completions",
      httpMethod = HttpPost,
      body = $payload
    )

    # Parse streaming response
    var buffer = ""
    for line in response.bodyStream.lines:
      let trimmed = line.strip()

      if trimmed.startsWith("data: "):
        let data = trimmed[6..^1].strip()

        if data == "[DONE]":
          break

        try:
          let chunk = parseJson(data)
          if chunk.hasKey("choices") and chunk["choices"].len > 0:
            let delta = chunk["choices"][0]["delta"]
            if delta.hasKey("content"):
              let content = delta["content"].getStr()
              if content.len > 0:
                callback(content)
        except:
          if self.debug:
            echo "[DEBUG] Parse error: ", getCurrentExceptionMsg()
  except:
    if self.debug:
      echo "[DEBUG] Stream error: ", getCurrentExceptionMsg()

proc tts*(
  self: UncloseAI,
  text: string,
  voice: string = "alloy",
  model: string = "tts-1",
  responseFormat: string = "mp3"
): string =
  if self.ttsEndpoints.len == 0:
    raise newException(ValueError, "No TTS endpoints available")

  let endpoint = self.ttsEndpoints[0]

  let payload = %* {
    "model": model,
    "voice": voice,
    "input": text,
    "response_format": responseFormat
  }

  let client = newHttpClient(timeout = self.timeout)
  client.headers = newHttpHeaders({"Content-Type": "application/json"})

  if self.apiKey != "":
    client.headers["Authorization"] = "Bearer " & self.apiKey

  let response = client.request(
    endpoint & "/audio/speech",
    httpMethod = HttpPost,
    body = $payload
  )

  return response.body

# Demo when run as main module
when isMainModule:
  echo "=== UncloseAI Nim Client (with Streaming) ===\n"

  let client = newUncloseAI(debug = true)

  let models = client.listModels()
  if models.len == 0:
    echo "ERROR: No models discovered. Set environment variables:"
    echo "  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc."
    quit(1)

  echo fmt"\nDiscovered {models.len} model(s):"
  for model in models:
    echo fmt"  - {model.id} (max_tokens: {model.maxTokens})"
  echo ""

  # Non-streaming chat
  echo "=== Non-Streaming Chat ==="
  try:
    let response = client.chat(
      @[
        ChatMessage(role: "system", content: "You are a helpful AI assistant."),
        ChatMessage(role: "user", content: "Explain quantum computing in one sentence.")
      ],
      maxTokens = 100
    )
    let content = response["choices"][0]["message"]["content"].getStr()
    echo "Response: ", content, "\n"
  except:
    echo "Error: ", getCurrentExceptionMsg(), "\n"

  # Streaming chat
  echo "=== Streaming Chat ==="
  let modelId = if models.len > 1: models[1].id else: ""
  echo "Model: ", if modelId != "": modelId else: models[0].id
  stdout.write("Response: ")
  stdout.flushFile()

  try:
    client.chatStream(
      @[
        ChatMessage(role: "system", content: "You are a coding assistant."),
        ChatMessage(role: "user", content: "Write a Nim function to check if a number is prime")
      ],
      model = modelId,
      maxTokens = 200,
      callback = proc(content: string) =
        stdout.write(content)
        stdout.flushFile()
    )
    echo "\n"
  except:
    echo "\nError: ", getCurrentExceptionMsg(), "\n"

  # TTS
  if client.ttsEndpoints.len > 0:
    echo "=== TTS Speech Generation ==="
    try:
      let audioData = client.tts("Hello from UncloseAI Nim client! This demonstrates streaming support.")
      writeFile("speech.mp3", audioData)
      echo fmt"[OK] Speech file created: speech.mp3 ({audioData.len} bytes)\n"
    except:
      echo "[ERROR] TTS Error: ", getCurrentExceptionMsg(), "\n"

  echo "=== Examples Complete ==="
