// 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 Foundation

extension [ModelContent] {
  // TODO: Rename and refactor this.
  /// Throws a ``GenerateContentError/internalError(underlying:)`` wrapping the error of the first
  /// part that failed to convert to model content (for example, an image that could not be
  /// converted to JPEG) so that a request is not sent with that part silently missing.
  func throwIfError() throws {
    for content in self {
      for part in content.internalParts {
        if case let .error(errorPart) = part.data {
          throw GenerateContentError.internalError(underlying: errorPart.error)
        }
      }
    }
  }
}

struct InternalPart: Equatable, Sendable {
  enum OneOfData: Equatable, Sendable {
    case text(String)
    case inlineData(InlineData)
    case fileData(FileData)
    case functionCall(FunctionCall)
    case functionResponse(FunctionResponse)
    case executableCode(ExecutableCode)
    case codeExecutionResult(CodeExecutionResult)
    /// A part that failed to convert to model content on the client (for example, an image that
    /// could not be converted to JPEG); it is never sent to or received from the backend.
    case error(ErrorPart)

    struct UnsupportedDataError: Error {
      let decodingError: DecodingError

      var localizedDescription: String {
        decodingError.localizedDescription
      }
    }
  }

  let data: OneOfData?

  let isThought: Bool?

  let thoughtSignature: String?

  let speechMetadata: SpeechMetadata?

  init(_ data: OneOfData,
       speechMetadata: SpeechMetadata? = nil,
       isThought: Bool?,
       thoughtSignature: String?) {
    self.data = data
    self.speechMetadata = speechMetadata
    self.isThought = isThought
    self.thoughtSignature = thoughtSignature
  }
}

/// A type describing data in media formats interpretable by an AI model. Each generative AI
/// request or response contains an `Array` of ``ModelContent``s, and each ``ModelContent`` value
/// may comprise multiple heterogeneous ``Part``s.
public struct ModelContent: Equatable, Sendable {
  /// The role of the entity creating the ``ModelContent``. For user-generated client requests,
  /// for example, the role is `user`.
  public let role: String?

  /// The data parts comprising this ``ModelContent`` value.
  public var parts: [any Part] {
    return internalParts.compactMap { part -> (any Part)? in
      switch part.data {
      case let .text(text):
        return TextPart(
          text,
          speechMetadata: part.speechMetadata,
          isThought: part.isThought,
          thoughtSignature: part.thoughtSignature
        )
      case let .inlineData(inlineData):
        return InlineDataPart(
          inlineData, isThought: part.isThought, thoughtSignature: part.thoughtSignature
        )
      case let .fileData(fileData):
        return FileDataPart(
          fileData, isThought: part.isThought, thoughtSignature: part.thoughtSignature
        )
      case let .functionCall(functionCall):
        return FunctionCallPart(
          functionCall, isThought: part.isThought, thoughtSignature: part.thoughtSignature
        )
      case let .functionResponse(functionResponse):
        return FunctionResponsePart(
          functionResponse, isThought: part.isThought, thoughtSignature: part.thoughtSignature
        )
      case let .executableCode(executableCode):
        return ExecutableCodePart(
          executableCode, isThought: part.isThought, thoughtSignature: part.thoughtSignature
        )
      case let .codeExecutionResult(codeExecutionResult):
        return CodeExecutionResultPart(
          codeExecutionResult: codeExecutionResult,
          isThought: part.isThought,
          thoughtSignature: part.thoughtSignature
        )
      case let .error(errorPart):
        return errorPart
      case .none:
        // Filter out parts that contain missing or unrecognized data
        return nil
      }
    }
  }

  // TODO: Refactor this
  let internalParts: [InternalPart]

  /// Creates a new value from a list of ``Part``s.
  public init(role: String? = "user", parts: [any Part]) {
    self.role = role
    var convertedParts = [InternalPart]()
    for part in parts {
      switch part {
      case let textPart as TextPart:
        convertedParts.append(InternalPart(
          .text(textPart.text),
          speechMetadata: textPart.speechMetadata,
          isThought: textPart._isThought,
          thoughtSignature: textPart.thoughtSignature
        ))
      case let inlineDataPart as InlineDataPart:
        convertedParts.append(InternalPart(
          .inlineData(inlineDataPart.inlineData),
          isThought: inlineDataPart._isThought,
          thoughtSignature: inlineDataPart.thoughtSignature
        ))
      case let fileDataPart as FileDataPart:
        convertedParts.append(InternalPart(
          .fileData(fileDataPart.fileData),
          isThought: fileDataPart._isThought,
          thoughtSignature: fileDataPart.thoughtSignature
        ))
      case let functionCallPart as FunctionCallPart:
        convertedParts.append(InternalPart(
          .functionCall(functionCallPart.functionCall),
          isThought: functionCallPart._isThought,
          thoughtSignature: functionCallPart.thoughtSignature
        ))
      case let functionResponsePart as FunctionResponsePart:
        convertedParts.append(InternalPart(
          .functionResponse(functionResponsePart.functionResponse),
          isThought: functionResponsePart._isThought,
          thoughtSignature: functionResponsePart.thoughtSignature
        ))
      case let executableCodePart as ExecutableCodePart:
        convertedParts.append(InternalPart(
          .executableCode(executableCodePart.executableCode),
          isThought: executableCodePart._isThought,
          thoughtSignature: executableCodePart.thoughtSignature
        ))
      case let codeExecutionResultPart as CodeExecutionResultPart:
        convertedParts.append(InternalPart(
          .codeExecutionResult(codeExecutionResultPart.codeExecutionResult),
          isThought: codeExecutionResultPart._isThought,
          thoughtSignature: codeExecutionResultPart.thoughtSignature
        ))
      case let errorPart as ErrorPart:
        // Keep the error so that `throwIfError()` fails the request instead of the part being
        // silently dropped.
        convertedParts.append(InternalPart(
          .error(errorPart),
          isThought: nil,
          thoughtSignature: nil
        ))
      default:
        AILog.error(
          code: .modelContentUnsupportedPartType,
          "Skipping a part with unsupported type: \(type(of: part))"
        )
      }
    }
    internalParts = convertedParts
  }

  /// Creates a new value from any data interpretable as a ``Part``.
  /// See ``PartsRepresentable`` for types that can be interpreted as `Part`s.
  public init(role: String? = "user", parts: any PartsRepresentable...) {
    let content = parts.flatMap { $0.partsValue }
    self.init(role: role, parts: content)
  }

  init(role: String?, parts: [InternalPart]) {
    self.role = role
    internalParts = parts
  }
}

// MARK: Codable Conformances

extension ModelContent: Codable {
  enum CodingKeys: String, CodingKey {
    case role
    case internalParts = "parts"
  }

  public init(from decoder: any Decoder) throws {
    let container = try decoder.container(keyedBy: CodingKeys.self)
    role = try container.decodeIfPresent(String.self, forKey: .role)
    internalParts = try container.decodeIfPresent([InternalPart].self, forKey: .internalParts) ?? []
  }
}

extension InternalPart: Codable {
  enum CodingKeys: String, CodingKey {
    case isThought = "thought"
    case thoughtSignature
    case speechMetadata
  }

  public func encode(to encoder: Encoder) throws {
    // `data` is `nil` for parts with missing or unrecognized data received from the backend (e.g.,
    // when sent back as chat history); encoding `nil` as a single value before requesting a keyed
    // container below would trap in `JSONEncoder`.
    if let data {
      try data.encode(to: encoder)
    }
    var container = encoder.container(keyedBy: CodingKeys.self)
    try container.encodeIfPresent(isThought, forKey: .isThought)
    try container.encodeIfPresent(thoughtSignature, forKey: .thoughtSignature)
    try container.encodeIfPresent(speechMetadata, forKey: .speechMetadata)
  }

  public init(from decoder: Decoder) throws {
    do {
      data = try OneOfData(from: decoder)
    } catch let error as OneOfData.UnsupportedDataError {
      AILog.error(code: .decodedUnsupportedPartData, error.localizedDescription)
      data = nil
    } catch { // Re-throw any other error types
      throw error
    }
    let container = try decoder.container(keyedBy: CodingKeys.self)
    isThought = try container.decodeIfPresent(Bool.self, forKey: .isThought)
    thoughtSignature = try container.decodeIfPresent(String.self, forKey: .thoughtSignature)
    speechMetadata = try container.decodeIfPresent(SpeechMetadata.self, forKey: .speechMetadata)
  }
}

extension InternalPart.OneOfData: Codable {
  enum CodingKeys: String, CodingKey {
    case text
    case inlineData
    case fileData
    case functionCall
    case functionResponse
    case executableCode
    case codeExecutionResult
  }

  public func encode(to encoder: Encoder) throws {
    var container = encoder.container(keyedBy: CodingKeys.self)
    switch self {
    case let .text(text):
      try container.encode(text, forKey: .text)
    case let .inlineData(inlineData):
      try container.encode(inlineData, forKey: .inlineData)
    case let .fileData(fileData):
      try container.encode(fileData, forKey: .fileData)
    case let .functionCall(functionCall):
      try container.encode(functionCall, forKey: .functionCall)
    case let .functionResponse(functionResponse):
      try container.encode(functionResponse, forKey: .functionResponse)
    case let .executableCode(executableCode):
      try container.encode(executableCode, forKey: .executableCode)
    case let .codeExecutionResult(codeExecutionResult):
      try container.encode(codeExecutionResult, forKey: .codeExecutionResult)
    case let .error(errorPart):
      // Requests reject these in `throwIfError()` before encoding; anything else that encodes one
      // (e.g., a `LiveSession`) gets an `EncodingError` instead of a crash.
      try errorPart.encode(to: encoder)
    }
  }

  public init(from decoder: Decoder) throws {
    let values = try decoder.container(keyedBy: CodingKeys.self)
    if values.contains(.text) {
      self = try .text(values.decode(String.self, forKey: .text))
    } else if values.contains(.inlineData) {
      self = try .inlineData(values.decode(InlineData.self, forKey: .inlineData))
    } else if values.contains(.fileData) {
      self = try .fileData(values.decode(FileData.self, forKey: .fileData))
    } else if values.contains(.functionCall) {
      self = try .functionCall(values.decode(FunctionCall.self, forKey: .functionCall))
    } else if values.contains(.functionResponse) {
      self = try .functionResponse(values.decode(FunctionResponse.self, forKey: .functionResponse))
    } else if values.contains(.executableCode) {
      self = try .executableCode(values.decode(ExecutableCode.self, forKey: .executableCode))
    } else if values.contains(.codeExecutionResult) {
      self = try .codeExecutionResult(
        values.decode(CodeExecutionResult.self, forKey: .codeExecutionResult)
      )
    } else {
      let unexpectedKeys = values.allKeys.map { $0.stringValue }
      throw UnsupportedDataError(decodingError: DecodingError.dataCorrupted(
        DecodingError.Context(
          codingPath: values.codingPath,
          debugDescription: "Unexpected Part type(s): \(unexpectedKeys)"
        )
      ))
    }
  }
}