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

// UncloseAI - Scala client for OpenAI-compatible APIs with streaming support

import sttp.client3._
import io.circe._
import io.circe.generic.auto._
import io.circe.parser._
import io.circe.syntax._
import java.nio.file.{Files, Paths}
import scala.util.{Try, Success, Failure}

case class ChatMessage(role: String, content: String)
case class ModelInfo(id: String, endpoint: String, max_tokens: Int = 8192)

// Public types
case class ResponseMessage(content: String)
case class Choice(message: ResponseMessage)
case class ChatResponse(choices: List[Choice])

// Internal types
private case class ChatRequest(model: String, messages: List[ChatMessage], max_tokens: Int, temperature: Double, stream: Boolean = false)
private case class DeltaMessage(content: Option[String])
private case class StreamChoice(delta: DeltaMessage)
private case class ChatStreamChunk(choices: List[StreamChoice])
private case class TTSRequest(model: String, voice: String, input: String)
private case class ModelData(id: String, max_model_len: Option[Int])
private case class ModelsResponse(data: List[ModelData])

class UncloseAI(
  modelEndpoints: Option[List[String]] = None,
  ttsEndpointsIn: Option[List[String]] = None,
  val apiKey: Option[String] = None,
  val timeout: Int = 30,
  val debug: Boolean = false
) {
  private val backend = HttpURLConnectionBackend()

  val models: List[ModelInfo] = discoverModels()
  val ttsEndpoints: List[String] = ttsEndpointsIn.getOrElse(discoverEnvEndpoints("TTS_ENDPOINT"))

  if (debug) {
    println(s"[DEBUG] Initialized with ${models.length} model(s)")
  }

  private def discoverEnvEndpoints(prefix: String): List[String] = {
    var endpoints = List[String]()
    for (i <- 1 to 9999) {
      sys.env.get(s"${prefix}_$i") match {
        case Some(endpoint) => endpoints = endpoints :+ endpoint
        case None => return endpoints
      }
    }
    endpoints
  }

  private def discoverModels(): List[ModelInfo] = {
    val endpoints = modelEndpoints.getOrElse(discoverEnvEndpoints("MODEL_ENDPOINT"))
    var modelsList = List[ModelInfo]()

    endpoints.foreach { endpoint =>
      if (debug) {
        println(s"[DEBUG] Discovering from: $endpoint")
      }

      Try {
        val response = basicRequest
          .get(uri"$endpoint/models")
          .readTimeout(scala.concurrent.duration.Duration(10, "seconds"))
          .send(backend)

        response.body match {
          case Right(body) =>
            decode[ModelsResponse](body) match {
              case Right(modelsResp) =>
                modelsResp.data.foreach { model =>
                  // Skip permission entries
                  if (!model.id.startsWith("modelperm-") && !model.id.startsWith("chatcmpl-")) {
                    modelsList = modelsList :+ ModelInfo(
                      model.id,
                      endpoint,
                      model.max_model_len.getOrElse(8192)
                    )
                    if (debug) {
                      println(s"[DEBUG] Discovered: ${model.id}")
                    }
                  }
                }
              case Left(_) =>
                if (debug) {
                  println(s"[DEBUG] Error parsing models from $endpoint")
                }
            }
          case Left(_) =>
            if (debug) {
              println(s"[DEBUG] Error fetching from $endpoint")
            }
        }
      }.recover {
        case e =>
          if (debug) {
            println(s"[DEBUG] Error: ${e.getMessage}")
          }
      }
    }

    modelsList
  }

  def listModels(): List[ModelInfo] = models

  private def resolveModel(modelId: Option[String]): ModelInfo = {
    if (models.isEmpty) {
      throw new Exception("No models available")
    }

    modelId match {
      case None => models.head
      case Some(id) => models.find(_.id == id).getOrElse {
        throw new Exception(s"Model '$id' not found")
      }
    }
  }

  def chat(
    messages: List[ChatMessage],
    model: Option[String] = None,
    maxTokens: Int = 100,
    temperature: Double = 0.7
  ): Either[String, ChatResponse] = {
    val modelInfo = resolveModel(model)

    val request = ChatRequest(
      model = modelInfo.id,
      messages = messages,
      max_tokens = maxTokens,
      temperature = temperature,
      stream = false
    )

    Try {
      val req = basicRequest
        .post(uri"${modelInfo.endpoint}/chat/completions")
        .contentType("application/json")
        .body(request.asJson.noSpaces)
        .readTimeout(scala.concurrent.duration.Duration(timeout, "seconds"))

      val finalReq = apiKey match {
        case Some(key) => req.header("Authorization", s"Bearer $key")
        case None => req
      }

      val response = finalReq.send(backend)

      response.body match {
        case Right(body) =>
          decode[ChatResponse](body) match {
            case Right(chatResponse) => Right(chatResponse)
            case Left(error) => Left(s"Parse error: ${error.getMessage}")
          }
        case Left(error) => Left(s"Request error: $error")
      }
    }.recover {
      case e: Exception => Left(s"Error: ${e.getMessage}")
    }.get
  }

  def chatStream(
    messages: List[ChatMessage],
    callback: String => Unit,
    model: Option[String] = None,
    maxTokens: Int = 500,
    temperature: Double = 0.7
  ): Unit = {
    val modelInfo = resolveModel(model)

    val request = ChatRequest(
      model = modelInfo.id,
      messages = messages,
      max_tokens = maxTokens,
      temperature = temperature,
      stream = true
    )

    Try {
      val req = basicRequest
        .post(uri"${modelInfo.endpoint}/chat/completions")
        .contentType("application/json")
        .body(request.asJson.noSpaces)
        .readTimeout(scala.concurrent.duration.Duration(timeout, "seconds"))

      val finalReq = apiKey match {
        case Some(key) => req.header("Authorization", s"Bearer $key")
        case None => req
      }

      val response = finalReq.send(backend)

      response.body match {
        case Right(body) =>
          // Parse SSE stream line by line
          body.split("\n").foreach { line =>
            if (line.startsWith("data: ")) {
              val data = line.substring(6).trim
              if (data == "[DONE]") {
                return
              }

              Try {
                decode[ChatStreamChunk](data) match {
                  case Right(chunk) =>
                    chunk.choices.headOption.foreach { choice =>
                      choice.delta.content.foreach { content =>
                        if (content.nonEmpty) {
                          callback(content)
                        }
                      }
                    }
                  case Left(_) => // Ignore parse errors
                }
              }.recover {
                case _ => // Ignore errors
              }
            }
          }
        case Left(error) =>
          if (debug) {
            println(s"[DEBUG] Stream error: $error")
          }
      }
    }.recover {
      case e: Exception =>
        if (debug) {
          println(s"[DEBUG] Error: ${e.getMessage}")
        }
    }
  }

  def tts(
    text: String,
    voice: String = "alloy",
    model: String = "tts-1",
    responseFormat: String = "mp3"
  ): Either[String, Array[Byte]] = {
    if (ttsEndpoints.isEmpty) {
      return Left("No TTS endpoints available")
    }

    val endpoint = ttsEndpoints.head

    val request = TTSRequest(
      model = model,
      voice = voice,
      input = text
    )

    Try {
      val req = basicRequest
        .post(uri"$endpoint/audio/speech")
        .contentType("application/json")
        .body(request.asJson.noSpaces)
        .readTimeout(scala.concurrent.duration.Duration(timeout, "seconds"))
        .response(asByteArray)

      val finalReq = apiKey match {
        case Some(key) => req.header("Authorization", s"Bearer $key")
        case None => req
      }

      val response = finalReq.send(backend)

      response.body match {
        case Right(audioData) => Right(audioData)
        case Left(error) => Left(s"Error: $error")
      }
    }.recover {
      case e: Exception => Left(s"Error: ${e.getMessage}")
    }.get
  }

  def close(): Unit = {
    backend.close()
  }
}

// Demo when run as application
object UncloseAI extends App {
  println("=== UncloseAI Scala Client (with Streaming) ===\n")

  val client = new UncloseAI(debug = true)

  try {
    if (client.listModels().isEmpty) {
      println("ERROR: No models discovered. Set environment variables:")
      println("  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.")
      System.exit(1)
    }

    println(s"\nDiscovered ${client.listModels().length} model(s):")
    client.listModels().foreach { model =>
      println(s"  - ${model.id} (max_tokens: ${model.max_tokens})")
    }
    println()

    // Non-streaming chat
    println("=== Non-Streaming Chat ===")
    client.chat(
      List(
        ChatMessage("system", "You are a helpful AI assistant."),
        ChatMessage("user", "Explain quantum computing in one sentence.")
      )
    ) match {
      case Right(response) =>
        println(s"Response: ${response.choices.head.message.content}\n")
      case Left(error) =>
        println(s"Error: $error\n")
    }

    // Streaming chat
    println("=== Streaming Chat ===")
    val modelId = if (client.listModels().length > 1) Some(client.listModels()(1).id) else None
    val modelName = modelId.getOrElse(client.listModels().head.id)
    println(s"Model: $modelName")
    print("Response: ")

    client.chatStream(
      List(
        ChatMessage("system", "You are a coding assistant."),
        ChatMessage("user", "Write a Scala function to check if a number is prime")
      ),
      content => print(content),
      modelId,
      200
    )
    println("\n")

    // TTS
    if (client.ttsEndpoints.nonEmpty) {
      println("=== TTS Speech Generation ===")
      client.tts("Hello from UncloseAI Scala client! This demonstrates streaming support.") match {
        case Right(audioData) =>
          Files.write(Paths.get("speech.mp3"), audioData)
          println(s"[OK] Speech file created: speech.mp3 (${audioData.length} bytes)\n")
        case Left(error) =>
          println(s"[ERROR] TTS Error: $error\n")
      }
    }

    println("=== Examples Complete ===")
  } finally {
    client.close()
  }
}
