// Copyright 2023 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import FirebaseAppCheckInterop
import FirebaseAuthInterop
import Foundation

/// A type that represents a remote multimodal model (like Gemini), with the ability to generate
/// content based on various input types.
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
public final class GenerativeModel {
  /// The resource name of the model in the backend; has the format "models/model-name".
  let modelResourceName: String

  /// The backing service responsible for sending and receiving model requests to the backend.
  let generativeAIService: GenerativeAIService

  /// Configuration parameters used for the MultiModalModel.
  let generationConfig: GenerationConfig?

  /// The safety settings to be used for prompts.
  let safetySettings: [SafetySetting]?

  /// A list of tools the model may use to generate the next response.
  let tools: [Tool]?

  /// Tool configuration for any `Tool` specified in the request.
  let toolConfig: ToolConfig?

  /// Instructions that direct the model to behave a certain way.
  let systemInstruction: ModelContent?

  /// Configuration parameters for sending requests to the backend.
  let requestOptions: RequestOptions

  /// Initializes a new remote model with the given parameters.
  ///
  /// - Parameters:
  ///   - name: The name of the model to use, for example `"gemini-1.0-pro"`.
  ///   - projectID: The project ID from the Firebase console.
  ///   - apiKey: The API key for your project.
  ///   - generationConfig: The content generation parameters your model should use.
  ///   - safetySettings: A value describing what types of harmful content your model should allow.
  ///   - tools: A list of ``Tool`` objects that the model may use to generate the next response.
  ///   - toolConfig: Tool configuration for any `Tool` specified in the request.
  ///   - systemInstruction: Instructions that direct the model to behave a certain way; currently
  ///     only text content is supported.
  ///   - requestOptions: Configuration parameters for sending requests to the backend.
  ///   - urlSession: The `URLSession` to use for requests; defaults to `URLSession.shared`.
  init(name: String,
       projectID: String,
       apiKey: String,
       generationConfig: GenerationConfig? = nil,
       safetySettings: [SafetySetting]? = nil,
       tools: [Tool]?,
       toolConfig: ToolConfig? = nil,
       systemInstruction: ModelContent? = nil,
       requestOptions: RequestOptions,
       appCheck: AppCheckInterop?,
       auth: AuthInterop?,
       urlSession: URLSession = .shared) {
    modelResourceName = name
    generativeAIService = GenerativeAIService(
      projectID: projectID,
      apiKey: apiKey,
      appCheck: appCheck,
      auth: auth,
      urlSession: urlSession
    )
    self.generationConfig = generationConfig
    self.safetySettings = safetySettings
    self.tools = tools
    self.toolConfig = toolConfig
    self.systemInstruction = systemInstruction
    self.requestOptions = requestOptions

    if VertexLog.additionalLoggingEnabled() {
      VertexLog.debug(code: .verboseLoggingEnabled, "Verbose logging enabled.")
    } else {
      VertexLog.info(code: .verboseLoggingDisabled, """
      [FirebaseVertexAI] To enable additional logging, add \
      `\(VertexLog.enableArgumentKey)` as a launch argument in Xcode.
      """)
    }
    VertexLog.debug(code: .generativeModelInitialized, "Model \(name) initialized.")
  }

  /// Generates content from String and/or image inputs, given to the model as a prompt, that are
  /// representable as one or more ``Part``s.
  ///
  /// Since ``Part``s do not specify a role, this method is intended for generating content from
  /// [zero-shot](https://developers.google.com/machine-learning/glossary/generative#zero-shot-prompting)
  /// or "direct" prompts. For
  /// [few-shot](https://developers.google.com/machine-learning/glossary/generative#few-shot-prompting)
  /// prompts, see `generateContent(_ content: [ModelContent])`.
  ///
  /// - Parameters:
  ///   - parts: The input(s) given to the model as a prompt (see ``PartsRepresentable`` for
  ///   conforming types).
  /// - Returns: The content generated by the model.
  /// - Throws: A ``GenerateContentError`` if the request failed.
  public func generateContent(_ parts: any PartsRepresentable...)
    async throws -> GenerateContentResponse {
    return try await generateContent([ModelContent(parts: parts)])
  }

  /// Generates new content from input content given to the model as a prompt.
  ///
  /// - Parameter content: The input(s) given to the model as a prompt.
  /// - Returns: The generated content response from the model.
  /// - Throws: A ``GenerateContentError`` if the request failed.
  public func generateContent(_ content: [ModelContent]) async throws
    -> GenerateContentResponse {
    try content.throwIfError()
    let response: GenerateContentResponse
    let generateContentRequest = GenerateContentRequest(model: modelResourceName,
                                                        contents: content,
                                                        generationConfig: generationConfig,
                                                        safetySettings: safetySettings,
                                                        tools: tools,
                                                        toolConfig: toolConfig,
                                                        systemInstruction: systemInstruction,
                                                        isStreaming: false,
                                                        options: requestOptions)
    do {
      response = try await generativeAIService.loadRequest(request: generateContentRequest)
    } catch {
      throw GenerativeModel.generateContentError(from: error)
    }

    // Check the prompt feedback to see if the prompt was blocked.
    if response.promptFeedback?.blockReason != nil {
      throw GenerateContentError.promptBlocked(response: response)
    }

    // Check to see if an error should be thrown for stop reason.
    if let reason = response.candidates.first?.finishReason, reason != .stop {
      throw GenerateContentError.responseStoppedEarly(reason: reason, response: response)
    }

    return response
  }

  /// Generates content from String and/or image inputs, given to the model as a prompt, that are
  /// representable as one or more ``Part``s.
  ///
  /// Since ``Part``s do not specify a role, this method is intended for generating content from
  /// [zero-shot](https://developers.google.com/machine-learning/glossary/generative#zero-shot-prompting)
  /// or "direct" prompts. For
  /// [few-shot](https://developers.google.com/machine-learning/glossary/generative#few-shot-prompting)
  /// prompts, see `generateContentStream(_ content: @autoclosure () throws -> [ModelContent])`.
  ///
  /// - Parameters:
  ///   - parts: The input(s) given to the model as a prompt (see ``PartsRepresentable`` for
  ///   conforming types).
  /// - Returns: A stream wrapping content generated by the model or a ``GenerateContentError``
  ///     error if an error occurred.
  @available(macOS 12.0, *)
  public func generateContentStream(_ parts: any PartsRepresentable...) throws
    -> AsyncThrowingStream<GenerateContentResponse, Error> {
    return try generateContentStream([ModelContent(parts: parts)])
  }

  /// Generates new content from input content given to the model as a prompt.
  ///
  /// - Parameter content: The input(s) given to the model as a prompt.
  /// - Returns: A stream wrapping content generated by the model or a ``GenerateContentError``
  ///     error if an error occurred.
  @available(macOS 12.0, *)
  public func generateContentStream(_ content: [ModelContent]) throws
    -> AsyncThrowingStream<GenerateContentResponse, Error> {
    try content.throwIfError()
    let generateContentRequest = GenerateContentRequest(model: modelResourceName,
                                                        contents: content,
                                                        generationConfig: generationConfig,
                                                        safetySettings: safetySettings,
                                                        tools: tools,
                                                        toolConfig: toolConfig,
                                                        systemInstruction: systemInstruction,
                                                        isStreaming: true,
                                                        options: requestOptions)

    var responseIterator = generativeAIService.loadRequestStream(request: generateContentRequest)
      .makeAsyncIterator()
    return AsyncThrowingStream {
      let response: GenerateContentResponse?
      do {
        response = try await responseIterator.next()
      } catch {
        throw GenerativeModel.generateContentError(from: error)
      }

      // The responseIterator will return `nil` when it's done.
      guard let response = response else {
        // This is the end of the stream! Signal it by sending `nil`.
        return nil
      }

      // Check the prompt feedback to see if the prompt was blocked.
      if response.promptFeedback?.blockReason != nil {
        throw GenerateContentError.promptBlocked(response: response)
      }

      // If the stream ended early unexpectedly, throw an error.
      if let finishReason = response.candidates.first?.finishReason, finishReason != .stop {
        throw GenerateContentError.responseStoppedEarly(reason: finishReason, response: response)
      } else {
        // Response was valid content, pass it along and continue.
        return response
      }
    }
  }

  /// Creates a new chat conversation using this model with the provided history.
  public func startChat(history: [ModelContent] = []) -> Chat {
    return Chat(model: self, history: history)
  }

  /// Runs the model's tokenizer on String and/or image inputs that are representable as one or more
  /// ``Part``s.
  ///
  /// Since ``Part``s do not specify a role, this method is intended for tokenizing
  /// [zero-shot](https://developers.google.com/machine-learning/glossary/generative#zero-shot-prompting)
  /// or "direct" prompts. For
  /// [few-shot](https://developers.google.com/machine-learning/glossary/generative#few-shot-prompting)
  /// input, see `countTokens(_ content: @autoclosure () throws -> [ModelContent])`.
  ///
  /// - Parameters:
  ///   - parts: The input(s) given to the model as a prompt (see ``PartsRepresentable`` for
  ///   conforming types).
  /// - Returns: The results of running the model's tokenizer on the input; contains
  /// ``CountTokensResponse/totalTokens``.
  public func countTokens(_ parts: any PartsRepresentable...) async throws -> CountTokensResponse {
    return try await countTokens([ModelContent(parts: parts)])
  }

  /// Runs the model's tokenizer on the input content and returns the token count.
  ///
  /// - Parameter content: The input given to the model as a prompt.
  /// - Returns: The results of running the model's tokenizer on the input; contains
  /// ``CountTokensResponse/totalTokens``.
  public func countTokens(_ content: [ModelContent]) async throws -> CountTokensResponse {
    let countTokensRequest = CountTokensRequest(
      model: modelResourceName,
      contents: content,
      systemInstruction: systemInstruction,
      tools: tools,
      generationConfig: generationConfig,
      options: requestOptions
    )
    return try await generativeAIService.loadRequest(request: countTokensRequest)
  }

  /// Returns a `GenerateContentError` (for public consumption) from an internal error.
  ///
  /// If `error` is already a `GenerateContentError` the error is returned unchanged.
  private static func generateContentError(from error: Error) -> GenerateContentError {
    if let error = error as? GenerateContentError {
      return error
    }
    return GenerateContentError.internalError(underlying: error)
  }
}