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

defmodule UncloseAI do
  @moduledoc """
  UncloseAI Elixir Library
  OpenAI-compatible API client with streaming support
  Compatible with vLLM, Ollama, and OpenAI-compatible endpoints
  """

  defmodule ModelInfo do
    @moduledoc "Struct to hold model information"
    defstruct [:id, :endpoint, :max_tokens]
  end

  defmodule Client do
    @moduledoc "Client struct to hold discovered models and endpoints"
    defstruct models: [], tts_endpoints: [], timeout: 30_000

    @doc """
    Initialize a new UncloseAI client with auto-discovery
    """
    def new(opts \\ []) do
      timeout = Keyword.get(opts, :timeout, 30_000)
      IO.puts("Initializing UncloseAI client...")

      # Discover chat/code models
      models =
        Stream.iterate(1, &(&1 + 1))
        |> Stream.take_while(fn i ->
          System.get_env("MODEL_ENDPOINT_#{i}") != nil
        end)
        |> Stream.map(fn i ->
          endpoint = System.get_env("MODEL_ENDPOINT_#{i}")
          IO.puts("Endpoint #{i}: #{endpoint}")
          discover_models_from_endpoint(endpoint)
        end)
        |> Enum.to_list()
        |> List.flatten()

      # Discover TTS endpoints
      tts_endpoints =
        Stream.iterate(1, &(&1 + 1))
        |> Stream.take_while(fn i ->
          System.get_env("TTS_ENDPOINT_#{i}") != nil
        end)
        |> Stream.map(fn i ->
          System.get_env("TTS_ENDPOINT_#{i}")
        end)
        |> Enum.to_list()

      IO.puts("Discovered #{length(models)} models, #{length(tts_endpoints)} TTS endpoints\n")

      %Client{models: models, tts_endpoints: tts_endpoints, timeout: timeout}
    end

    defp discover_models_from_endpoint(endpoint) do
      case HTTPoison.get("#{endpoint}/models", [], recv_timeout: 10_000) do
        {:ok, %HTTPoison.Response{body: body}} ->
          data = Jason.decode!(body)

          case data["data"] do
            nil ->
              []

            models ->
              Enum.filter(models, fn model ->
                # Skip modelperm-* entries
                not String.starts_with?(model["id"], "modelperm-")
              end)
              |> Enum.map(fn model ->
                model_id = model["id"]
                max_tokens = model["max_model_len"] || 8192
                %ModelInfo{id: model_id, endpoint: endpoint, max_tokens: max_tokens}
              end)
          end

        {:error, _reason} ->
          # Silently skip failed endpoints
          []
      end
    end

    @doc """
    Non-streaming chat completion
    """
    def chat(client, messages, opts \\ []) do
      model_idx = Keyword.get(opts, :model_idx, 0)
      max_tokens = Keyword.get(opts, :max_tokens, 100)
      temperature = Keyword.get(opts, :temperature, 0.7)

      model = Enum.at(client.models, model_idx)

      if model == nil do
        {:error, "Invalid model index"}
      else
        url = "#{model.endpoint}/chat/completions"

        request = %{
          model: model.id,
          messages: messages,
          stream: false,
          max_tokens: max_tokens,
          temperature: temperature
        }

        case HTTPoison.post(
               url,
               Jason.encode!(request),
               [{"Content-Type", "application/json"}],
               recv_timeout: client.timeout
             ) do
          {:ok, %HTTPoison.Response{body: body}} ->
            data = Jason.decode!(body)
            {:ok, get_in(data, ["choices", Access.at(0), "message", "content"])}

          {:error, reason} ->
            {:error, reason}
        end
      end
    end

    @doc """
    Streaming chat completion - returns a Stream that yields content chunks
    """
    def chat_stream(client, messages, opts \\ []) do
      model_idx = Keyword.get(opts, :model_idx, 0)
      max_tokens = Keyword.get(opts, :max_tokens, 500)
      temperature = Keyword.get(opts, :temperature, 0.7)

      model = Enum.at(client.models, model_idx)

      if model == nil do
        raise "Invalid model index"
      end

      url = "#{model.endpoint}/chat/completions"

      request = %{
        model: model.id,
        messages: messages,
        stream: true,
        max_tokens: max_tokens,
        temperature: temperature
      }

      # Use HTTPoison stream with async response handling
      Stream.resource(
        fn ->
          {:ok, response} =
            HTTPoison.post(
              url,
              Jason.encode!(request),
              [{"Content-Type", "application/json"}],
              stream_to: self(),
              async: :once,
              recv_timeout: client.timeout
            )

          {response, ""}
        end,
        fn {response, buffer} ->
          receive do
            %HTTPoison.AsyncStatus{} ->
              HTTPoison.stream_next(response)
              {[], {response, buffer}}

            %HTTPoison.AsyncHeaders{} ->
              HTTPoison.stream_next(response)
              {[], {response, buffer}}

            %HTTPoison.AsyncChunk{chunk: chunk} ->
              # Append chunk to buffer and process lines
              new_buffer = buffer <> chunk
              {lines, remaining} = extract_lines(new_buffer)

              contents =
                lines
                |> Enum.filter(&String.starts_with?(&1, "data: "))
                |> Enum.map(&String.replace_prefix(&1, "data: ", ""))
                |> Enum.reject(&(&1 == "[DONE]"))
                |> Enum.map(&extract_content/1)
                |> Enum.reject(&is_nil/1)

              HTTPoison.stream_next(response)
              {contents, {response, remaining}}

            %HTTPoison.AsyncEnd{} ->
              {:halt, {response, buffer}}
          after
            client.timeout ->
              {:halt, {response, buffer}}
          end
        end,
        fn {response, _buffer} ->
          :hackney.close(response.id)
        end
      )
    end

    defp extract_lines(buffer) do
      lines = String.split(buffer, "\n")

      case List.last(lines) do
        "" -> {Enum.drop(lines, -1), ""}
        partial -> {Enum.drop(lines, -1), partial}
      end
    end

    defp extract_content(data) do
      case Jason.decode(data) do
        {:ok, json} ->
          get_in(json, ["choices", Access.at(0), "delta", "content"])

        {:error, _} ->
          nil
      end
    end

    @doc """
    Text-to-speech generation
    """
    def tts(client, text, opts \\ []) do
      voice = Keyword.get(opts, :voice, "alloy")
      output_file = Keyword.get(opts, :output_file, "/tmp/speech.mp3")

      if Enum.empty?(client.tts_endpoints) do
        {:error, "No TTS endpoints available"}
      else
        endpoint = List.first(client.tts_endpoints)
        url = "#{endpoint}/audio/speech"

        request = %{
          model: "tts-1",
          voice: voice,
          input: text
        }

        case HTTPoison.post(
               url,
               Jason.encode!(request),
               [{"Content-Type", "application/json"}],
               recv_timeout: client.timeout
             ) do
          {:ok, %HTTPoison.Response{body: body}} ->
            File.write!(output_file, body)
            {:ok, output_file}

          {:error, reason} ->
            {:error, reason}
        end
      end
    end
  end
end
