defmodule Honeycomb.OpenAI do
  @moduledoc false

  alias Honeycomb.Serving

  defmodule Chunk do
    @moduledoc false
    defstruct [:choices, :created, :id, :model, :object]

    def new(model, text, id \\ id(), created \\ System.os_time(:second)) do
      struct(__MODULE__,
        choices: choices(text),
        created: created,
        id: id,
        model: model,
        object: "chat.completion.chunk"
      )
    end

    defp id() do
      encoded = Base.url_encode64(:crypto.strong_rand_bytes(21), padding: false)
      "honeycomb-#{String.slice(encoded, 0..21)}"
    end

    defp choices(generation) do
      [%{delta: %{content: generation}, finish_reason: nil, index: 0}]
    end
  end

  defmodule Response do
    @moduledoc false
    defstruct [:choices, :created, :id, :model, :object, :usage]

    def new(model, %{text: text, token_summary: usage}) do
      struct(__MODULE__,
        choices: choices(text),
        created: System.os_time(:second),
        id: id(),
        model: model,
        object: "chat.completion",
        usage: usage(usage)
      )
    end

    defp id() do
      encoded = Base.url_encode64(:crypto.strong_rand_bytes(21), padding: false)
      "honeycomb-#{String.slice(encoded, 0..21)}"
    end

    defp choices(generation) do
      %{
        finish_reason: "stop",
        index: 0,
        message: %{
          content: generation,
          role: "assistant"
        },
        logprobs: nil
      }
    end

    defp usage(%{input: inp, output: out}) do
      %{
        completion_tokens: out,
        prompt_tokens: inp,
        total_tokens: inp + out
      }
    end
  end

  @system_message_keys [
    content: [type: :string, required: true, doc: "The contents of the system message."],
    role: [
      type: {:in, ["system"]},
      required: true,
      doc: "The role of the messages author, in this case `system`."
    ],
    name: [
      type: :string,
      doc:
        "An optional name for the participant. Provides the model information to differentiate between participants of the same role."
    ]
  ]

  @user_message_keys [
    content: [type: :string, required: true, doc: "The contents of the user message."],
    role: [
      type: {:in, ["user"]},
      required: true,
      doc: "The role of the messages author, in this case `user`."
    ],
    name: [
      type: :string,
      doc:
        "An optional name for the participant. Provides the model information to differentiate between participants of the same role."
    ]
  ]

  @assistant_message_keys [
    content: [
      type: :string,
      required: true,
      doc:
        "The contents of the assistant message. Required unless `tool_calls` or `function_call` is specified."
    ],
    role: [
      type: {:in, ["assistant"]},
      required: true,
      doc: "The role of the messages author, in this case `assistant`."
    ],
    name: [
      type: :string,
      doc:
        "An optional name for the participant. Provides the model information to differentiate between participants of the same role."
    ],
    tool_calls: [
      doc: "The tool calls generated by the model, such as function calls.",
      type:
        {:list,
         [
           id: [type: :string, required: true, doc: "The ID of the tool call."],
           type: [
             type: :string,
             required: true,
             doc: "The type of the tool. Currently, only function is supported."
           ],
           function: [
             doc: "The function that the model called.",
             type: :map,
             keys: [
               name: [type: :string, required: true],
               arguments: [type: :string, required: true]
             ]
           ]
         ]}
    ],
    function_call: [
      deprecated:
        "Deprecated and replaced by tool_calls. The name and arguments of a function that should be called, as generated by the model."
    ]
  ]

  @tool_message_keys [
    content: [type: :string, required: true, doc: "The contents of the tool message."],
    role: [
      type: {:in, ["tool"]},
      required: true,
      doc: "The role of the messages author, in this case `tool`."
    ],
    tool_call_id: [
      type: :string,
      required: true,
      doc: "Tool call that this message is responding to."
    ]
  ]

  @function_message_schema [
    deprecated: "Deprecated in favor of 'tool_message'"
  ]

  @number_schema {:or, [:integer, :float]}

  # Implements logic for OpenAI compatible endpoints
  @body_params_schema [
    messages: [
      type:
        {:list,
         {:or,
          [
            map: @user_message_keys,
            map: @system_message_keys,
            map: @tool_message_keys,
            map: @assistant_message_keys,
            map: @function_message_schema
          ]}},
      required: true,
      doc: "A list of messages comprising the conversation so far. "
    ],
    model: [type: :string, doc: "ID of the model to use"],
    frequency_penalty: [
      type: {:or, [nil, @number_schema]},
      default: 0,
      doc:
        "Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim."
    ],
    logit_bias: [
      type: :map,
      doc:
        "Modify the likelihood of specified tokens appearing in the completion. Accepts a JSON object that maps tokens (specified by their token ID in the tokenizer) to an associated bias value from -100 to 100. Mathematically, the bias is added to the logits generated by the model prior to sampling. The exact effect will vary per model, but values between -1 and 1 should decrease or increase likelihood of selection; values like -100 or 100 should result in a ban or exclusive selection of the relevant token."
    ],
    logprobs: [
      type: {:or, [:boolean, nil]},
      default: false,
      doc:
        "Whether to return log probabilities of the output tokens or not. If true, returns the log probabilities of each output token returned in the content of message."
    ],
    top_logprobs: [
      type: {:or, [:integer, nil]},
      doc:
        "An integer between 0 and 20 specifying the number of most likely tokens to return at each token position, each with an associated log probability. logprobs must be set to true if this parameter is used."
    ],
    max_tokens: [
      type: {:or, [:integer, nil]},
      doc:
        "The maximum number of tokens that can be generated in the chat completion. The total length of input tokens and generated tokens is limited by the model's context length."
    ],
    n: [
      type: {:or, [:integer, nil]},
      default: 1,
      doc:
        "How many chat completion choices to generate for each input message. Note that you will be charged based on the number of generated tokens across all of the choices. "
    ],
    presence_penalty: [
      type: {:or, [@number_schema, nil]},
      default: 0,
      doc:
        "Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the text so far, increasing the model's likelihood to talk about new topics."
    ],
    response_format: [
      type: :map,
      keys: [
        type: [
          type: {:in, ["text", "json_object"]},
          default: "text",
          doc:
            "Setting to 'json_object' enables JSON mode, which guarantees the message the model generates is valid JSON.
            Important: when using JSON mode, you must also instruct the model to produce JSON yourself via a system or user message. Without this, the model may generate an unending stream of whitespace until the generation reaches the token limit, resulting in a long-running and seemingly 'stuck' request. Also note that the message content may be partially cut off if finish_reason='length', which indicates the generation exceeded max_tokens or the conversation exceeded the max context length."
        ]
      ],
      default: %{type: "text"},
      doc:
        "An object specifying the format that the model must output. Not compatible with all models."
    ],
    seed: [
      type: {:or, [:integer, nil]},
      doc:
        "This feature is in Beta. If specified, our system will make a best effort to sample deterministically, such that repeated requests with the same seed and parameters should return the same result. Determinism is not guaranteed, and you should refer to the system_fingerprint response parameter to monitor changes in the backend."
    ],
    service_tier: [
      type: {:in, [nil, "auto", "default"]},
      default: nil,
      doc:
        "Specifies the latency tier to use for processing the request. This parameter is relevant for customers subscribed to the scale tier service"
    ],
    stop: [
      type: {:or, [:string, {:list, :string}, nil]},
      doc: "Up to 4 sequences where the API will stop generating further tokens."
    ],
    stream: [
      type: {:or, [:boolean, nil]},
      default: false,
      doc: "If set, partial message deltas will be sent, like in ChatGPT. "
    ],
    stream_options: [
      type: {:or, [nil, map: [include_usage: [type: :boolean]]]},
      doc: "Options for streaming response. Only set this when you set `stream: true`."
    ],
    temperature: [
      type: {:or, [@number_schema, nil]},
      default: 1,
      doc:
        "What sampling temperature to use, between 0 and 2. Higher values like 0.8 will make the output more random, while lower values like 0.2 will make it more focused and deterministic. We generally recommend altering this or top_p but not both."
    ],
    top_p: [
      type: {:or, [@number_schema, nil]},
      default: 1,
      doc:
        "An alternative to sampling with `temperature`, called nucleus sampling, where the model considers the results of the tokens with top_p probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered. We generally recommend altering this or temperature but not both."
    ],
    tools: [
      doc:
        "A list of tools the model may call. Currently, only functions are supported as a tool. Use this to provide a list of functions the model may generate JSON inputs for. A max of 128 functions are supported.",
      type:
        {:list,
         {
           :map,
           [
             type: [
               type: :string,
               required: true,
               doc: "The type of the tool. Currently, only `function` is supported."
             ],
             function: [
               type: :map,
               required: true,
               keys: [
                 description: [
                   type: :string,
                   doc:
                     "A description of what the function does, used by the model to choose when and how to call the function."
                 ],
                 name: [
                   type: :string,
                   required: true,
                   doc:
                     "The name of the function to be called. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length of 64."
                 ],
                 parameters: [
                   type: :map,
                   doc: "The parameters the functions accepts, described as a JSON Schema object"
                 ]
               ]
             ]
           ]
         }}
    ],
    tool_choice: [
      doc:
        "Controls which (if any) tool is called by the model. `none` means the model will not call any tool and instead generates a message. `auto` means the model can pick between generating a message or calling one or more tools. `required` means the model must call one or more tools. Specifying a particular tool via `{'type': 'function', 'function': {'name': 'my_function'}}` forces the model to call that tool.",
      type:
        {:or,
         [
           {:in, ["auto", "none", "required"]},
           map: [
             type: [
               type: :string,
               required: true,
               doc: "The type of the tool. Currently, only `function` is supported."
             ],
             function: [
               type: :map,
               required: true,
               keys: [
                 name: [type: :string, required: true, doc: "The name of the function to call."]
               ]
             ]
           ]
         ]},
      default: "none"
    ],
    parallel_tool_calls: [
      type: :boolean,
      default: true,
      doc: "Whether to enable parallel function calling during tool use."
    ],
    user: [
      type: :string,
      doc:
        "A unique identifier representing your end-user, which can help to monitor and detect abuse."
    ],
    function_call: [deprecated: "Deprecated in favor of tool_choice."],
    functions: [deprecated: "Deprecated in favor of tools."]
  ]

  @body_params_validator NimbleOptions.new!(@body_params_schema)
  def chat_completion(opts \\ []) when is_list(opts) do
    assert_serving_started!()

    case NimbleOptions.validate(opts, @body_params_validator) do
      {:ok, opts} ->
        messages = Keyword.fetch!(opts, :messages)
        stream? = Keyword.fetch!(opts, :stream)

        model = Serving.model()
        stream = Serving.generate(messages)

        if stream? do
          do_stream_chunks(model, stream)
        else
          do_consume_stream(model, stream)
        end

      {:error, %NimbleOptions.ValidationError{message: msg}} ->
        {:error, msg}
    end
  end

  defp do_stream_chunks(model, stream) do
    %{id: id, created: created} = Chunk.new(model, "")

    Stream.map(stream, fn text ->
      model
      |> Chunk.new(text, id, created)
      |> Map.from_struct()
    end)
    |> Stream.drop(-1)
  end

  defp do_consume_stream(model, stream) do
    {text, summary} =
      Enum.reduce(stream, {"", nil}, fn res, {acc, _last} ->
        case res do
          text when is_binary(res) -> {acc <> text, nil}
          {:done, summary} -> {acc, summary}
        end
      end)

    summary
    |> Map.put(:text, text)
    |> then(&Response.new(model, &1))
    |> then(&{:ok, &1})
  end

  defp assert_serving_started! do
    with nil <- Process.whereis(Honeycomb.Serving) do
      raise "could not find started Honeycomb serving process"
    end
  end
end
