import Foundation

enum MP3EncodingError: LocalizedError {
  case invalidRecording
  case encodingFailed

  var errorDescription: String? {
    switch self {
    case .invalidRecording: "The recording could not be read."
    case .encodingFailed: "The recording could not be compressed."
    }
  }
}

enum MP3Encoder {
  static func encode(wav url: URL) throws -> URL {
    let wav = try Data(contentsOf: url, options: .mappedIfSafe)
    let pcm = try pcmData(fromWAV: wav)
    var encodedBytes: UnsafeMutablePointer<UInt8>?
    var encodedSize = 0
    let result = pcm.withUnsafeBytes { bytes in
      transui_mp3_encode_mono_16k_32kbps(
        bytes.bindMemory(to: UInt8.self).baseAddress,
        bytes.count,
        &encodedBytes,
        &encodedSize
      )
    }
    guard result == TRANSUI_MP3_OK, let encodedBytes, encodedSize > 0 else {
      transui_mp3_free(encodedBytes)
      throw MP3EncodingError.encodingFailed
    }
    defer { transui_mp3_free(encodedBytes) }

    let output = FileManager.default.temporaryDirectory
      .appendingPathComponent("transui-\(UUID().uuidString).mp3")
    try Data(bytes: encodedBytes, count: encodedSize).write(to: output, options: .atomic)
    return output
  }

  static func pcmData(fromWAV wav: Data) throws -> Data {
    guard wav.count >= 12,
      String(decoding: wav[0..<4], as: UTF8.self) == "RIFF",
      String(decoding: wav[8..<12], as: UTF8.self) == "WAVE"
    else { throw MP3EncodingError.invalidRecording }

    var formatIsSupported = false
    var audio: Data?
    var offset = 12
    while offset <= wav.count - 8 {
      let chunkName = String(decoding: wav[offset..<(offset + 4)], as: UTF8.self)
      let chunkSize = Int(readUInt32(wav, at: offset + 4))
      let contentOffset = offset + 8
      guard chunkSize >= 0, contentOffset <= wav.count,
        chunkSize <= wav.count - contentOffset
      else { throw MP3EncodingError.invalidRecording }

      if chunkName == "fmt ", chunkSize >= 16 {
        formatIsSupported = readUInt16(wav, at: contentOffset) == 1
          && readUInt16(wav, at: contentOffset + 2) == 1
          && readUInt32(wav, at: contentOffset + 4) == 16_000
          && readUInt16(wav, at: contentOffset + 14) == 16
      } else if chunkName == "data" {
        audio = wav.subdata(in: contentOffset..<(contentOffset + chunkSize))
      }

      let paddedSize = chunkSize + (chunkSize & 1)
      guard paddedSize <= Int.max - contentOffset else {
        throw MP3EncodingError.invalidRecording
      }
      offset = contentOffset + paddedSize
    }

    guard formatIsSupported, let audio, !audio.isEmpty, audio.count.isMultiple(of: 2) else {
      throw MP3EncodingError.invalidRecording
    }
    return audio
  }

  private static func readUInt16(_ data: Data, at offset: Int) -> UInt16 {
    UInt16(data[offset]) | UInt16(data[offset + 1]) << 8
  }

  private static func readUInt32(_ data: Data, at offset: Int) -> UInt32 {
    UInt32(data[offset])
      | UInt32(data[offset + 1]) << 8
      | UInt32(data[offset + 2]) << 16
      | UInt32(data[offset + 3]) << 24
  }
}