/*
 * Copyright 2026 Signal Messenger, LLC
 * SPDX-License-Identifier: AGPL-3.0-only
 */

package org.signal.network.api

import okio.ByteString.Companion.toByteString
import org.signal.core.models.AccountEntropyPool
import org.signal.core.models.MasterKey
import org.signal.core.models.ServiceId.ACI
import org.signal.core.models.ServiceId.PNI
import org.signal.core.models.backup.MediaRootBackupKey
import org.signal.core.models.backup.MessageBackupKey
import org.signal.core.util.Base64
import org.signal.core.util.logging.Log
import org.signal.core.util.urlEncode
import org.signal.libsignal.net.AuthDevicesService
import org.signal.libsignal.net.DeviceIdNotFoundException
import org.signal.libsignal.net.RequestResult
import org.signal.libsignal.protocol.IdentityKeyPair
import org.signal.libsignal.protocol.ecc.ECPublicKey
import org.signal.libsignal.zkgroup.profiles.ProfileKey
import org.signal.network.NetworkResult
import org.signal.network.websocket.WebSocketRequestMessage
import org.signal.network.websocket.get
import org.signal.network.websocket.put
import org.whispersystems.signalservice.api.fromWebSocketRequest
import org.whispersystems.signalservice.api.link.LinkedDeviceVerificationCodeResponse
import org.whispersystems.signalservice.api.link.SetLinkedDeviceTransferArchiveRequest
import org.whispersystems.signalservice.api.link.TransferArchiveError
import org.whispersystems.signalservice.api.link.TransferArchiveResponse
import org.whispersystems.signalservice.api.link.WaitForLinkedDeviceResponse
import org.whispersystems.signalservice.api.provisioning.ProvisioningMessage
import org.whispersystems.signalservice.api.websocket.SignalWebSocket
import org.whispersystems.signalservice.internal.crypto.PrimaryProvisioningCipher
import org.whispersystems.signalservice.internal.push.ProvisionMessage
import org.whispersystems.signalservice.internal.push.ProvisioningVersion
import kotlin.time.Duration
import kotlin.time.Duration.Companion.seconds
import org.signal.libsignal.net.LinkedDevice as LibSignalLinkedDevice

/**
 * Class to interact with device-linking endpoints.
 */
class LinkDeviceApi(
  private val authWebSocket: SignalWebSocket.AuthenticatedWebSocket
) {

  companion object {
    private val TAG = Log.tag(LinkDeviceApi::class)
  }

  /**
   * Fetches a list of linked devices.
   */
  suspend fun getDevices(): RequestResult<List<LibSignalLinkedDevice>, Nothing> {
    return authWebSocket.runCatchingWithChatConnection { connection ->
      AuthDevicesService(connection).getDevices()
    }
  }

  /**
   * Remove and unlink a linked device.
   */
  suspend fun removeDevice(deviceId: Int): RequestResult<Unit, Nothing> {
    return authWebSocket.runCatchingWithChatConnection { connection ->
      AuthDevicesService(connection).removeDevice(deviceId)
    }
  }

  /**
   * Fetches a new verification code that lets you link a new device.
   *
   * GET /v1/devices/provisioning/code
   *
   * - 200: Success.
   * - 411: Account is already at the device limit.
   * - 429: Rate-limited.
   */
  fun getDeviceVerificationCode(): NetworkResult<LinkedDeviceVerificationCodeResponse> {
    val request = WebSocketRequestMessage.get("/v1/devices/provisioning/code")
    return NetworkResult
      .fromWebSocketRequest(authWebSocket, request, LinkedDeviceVerificationCodeResponse::class)
  }

  /**
   * Links a new device to the account.
   *
   * PUT /v1/provisioning/[deviceIdentifier]
   *
   * - 200: Success.
   * - 403: Account not found or incorrect verification code.
   * - 409: The new device is missing a required capability.
   * - 411: Account is already at the device limit.
   * - 422: Bad request.
   * - 429: Rate-limited.
   *
   * [e164], [pni], and [pniIdentityKeyPair] are absent for an account with no phone number, and are only ever sent
   * as a complete set.
   */
  fun linkDevice(
    e164: String?,
    aci: ACI,
    pni: PNI?,
    deviceIdentifier: String,
    deviceKey: ECPublicKey,
    aciIdentityKeyPair: IdentityKeyPair,
    pniIdentityKeyPair: IdentityKeyPair?,
    profileKey: ProfileKey,
    accountEntropyPool: AccountEntropyPool,
    masterKey: MasterKey,
    mediaRootBackupKey: MediaRootBackupKey,
    code: String,
    ephemeralMessageBackupKey: MessageBackupKey?
  ): NetworkResult<Unit> {
    val sendPhoneNumberData = e164 != null && pni != null && pniIdentityKeyPair != null
    if (!sendPhoneNumberData && (e164 != null || pni != null || pniIdentityKeyPair != null)) {
      Log.w(TAG, "[linkDevice] Incomplete phone number data! hasNumber: ${e164 != null}, hasPni: ${pni != null}, hasPniIdentityKey: ${pniIdentityKeyPair != null}. Linking the new device without a phone number.")
    }

    val sentE164 = e164.takeIf { sendPhoneNumberData }
    val sentPni = pni.takeIf { sendPhoneNumberData }
    val sentPniIdentityKeyPair = pniIdentityKeyPair.takeIf { sendPhoneNumberData }

    val cipher = PrimaryProvisioningCipher(deviceKey)
    val message = ProvisionMessage(
      aciIdentityKeyPublic = aciIdentityKeyPair.publicKey.serialize().toByteString(),
      aciIdentityKeyPrivate = aciIdentityKeyPair.privateKey.serialize().toByteString(),
      pniIdentityKeyPublic = sentPniIdentityKeyPair?.publicKey?.serialize()?.toByteString(),
      pniIdentityKeyPrivate = sentPniIdentityKeyPair?.privateKey?.serialize()?.toByteString(),
      aci = aci.toString(),
      pni = sentPni?.toStringWithoutPrefix(),
      number = sentE164,
      provisioningCode = code,
      userAgent = null,
      profileKey = profileKey.serialize().toByteString(),
      provisioningVersion = ProvisioningVersion.CURRENT.value,
      ephemeralBackupKey = ephemeralMessageBackupKey?.value?.toByteString(),
      accountEntropyPool = accountEntropyPool.value,
      mediaRootBackupKey = mediaRootBackupKey.value.toByteString(),
      aciBinary = aci.toByteString(),
      pniBinary = sentPni?.toByteStringWithoutPrefix()
    )
    val ciphertext: ByteArray = cipher.encrypt(message)
    val body = ProvisioningMessage(Base64.encodeWithPadding(ciphertext))

    val request = WebSocketRequestMessage.put("/v1/provisioning/${deviceIdentifier.urlEncode()}", body)
    return NetworkResult.fromWebSocketRequest(authWebSocket, request)
  }

  /**
   * A "long-polling" endpoint that will return once the device has successfully been linked.
   *
   * @param timeout The max amount of time to wait. Capped at 30 seconds.
   *
   * GET /v1/devices/wait_for_linked_device/[token]?timeout=[timeout]
   *
   * - 200: Success, a new device was linked associated with the provided token.
   * - 204: No device was linked before the max waiting time elapsed.
   * - 400: Invalid token/timeout.
   * - 429: Rate-limited.
   */
  fun waitForLinkedDevice(token: String, timeout: Duration = 30.seconds): NetworkResult<WaitForLinkedDeviceResponse> {
    val request = WebSocketRequestMessage.get("/v1/devices/wait_for_linked_device/${token.urlEncode()}?timeout=${timeout.inWholeSeconds}")
    return NetworkResult
      .fromWebSocketRequest(
        signalWebSocket = authWebSocket,
        request = request,
        timeout = timeout,
        webSocketResponseConverter = NetworkResult.LongPollingWebSocketConverter(WaitForLinkedDeviceResponse::class)
      )
  }

  /**
   * After a device has been linked and an archive has been uploaded, you can call this endpoint to share the archive with the linked device.
   *
   * PUT /v1/devices/transfer_archive
   *
   * - 204: Success.
   * - 422: Bad inputs.
   * - 429: Rate-limited.
   */
  fun setTransferArchive(destinationDeviceId: Int, destinationDeviceRegistrationId: Int, cdn: Int, cdnKey: String): NetworkResult<Unit> {
    val body = SetLinkedDeviceTransferArchiveRequest(
      destinationDeviceId = destinationDeviceId,
      destinationDeviceRegistrationId = destinationDeviceRegistrationId,
      transferArchive = SetLinkedDeviceTransferArchiveRequest.TransferArchive.CdnInfo(
        cdn = cdn,
        key = cdnKey
      )
    )
    val request = WebSocketRequestMessage.put("/v1/devices/transfer_archive", body)
    return NetworkResult.fromWebSocketRequest(authWebSocket, request)
  }

  /**
   * If creating an archive has failed after linking a device, notify the linked
   * device of the failure and if you are going to try relinking or skip syncing
   *
   * PUT /v1/devices/transfer_archive
   *
   * - 204: Success.
   * - 422: Bad inputs.
   * - 429: Rate-limited.
   */
  fun setTransferArchiveError(destinationDeviceId: Int, destinationDeviceRegistrationId: Int, error: TransferArchiveError): NetworkResult<Unit> {
    val body = SetLinkedDeviceTransferArchiveRequest(
      destinationDeviceId = destinationDeviceId,
      destinationDeviceRegistrationId = destinationDeviceRegistrationId,
      transferArchive = SetLinkedDeviceTransferArchiveRequest.TransferArchive.Error(error)
    )
    val request = WebSocketRequestMessage.put("/v1/devices/transfer_archive", body)
    return NetworkResult.fromWebSocketRequest(authWebSocket, request)
  }

  /**
   * Sets the name for a linked device.
   *
   * @param encryptedDeviceName Must be between 1 and 225 bytes long.
   *
   * A [DeviceIdNotFoundException] means there is no device with the given [deviceId].
   */
  suspend fun setDeviceName(encryptedDeviceName: ByteArray, deviceId: Int): RequestResult<Unit, DeviceIdNotFoundException> {
    return authWebSocket.runCatchingWithChatConnection { connection ->
      AuthDevicesService(connection).setDeviceName(deviceId, encryptedDeviceName)
    }
  }

  /**
   * A "long-polling" endpoint that will return once the primary device has successfully sent sync data.
   *
   * @param timeout The max amount of time to wait. Capped at 30 seconds.
   *
   * GET /v1/devices/transfer_archive?timeout=[timeout]
   *
   * - 200: Success, the primary device was sent backup sync data.
   * - 204: The primary didn't provide data before the max waiting time elapsed.
   * - 400: Invalid timeout.
   * - 429: Rate-limited.
   */
  fun waitForPrimaryDevice(timeout: Duration = 30.seconds): NetworkResult<TransferArchiveResponse> {
    val request = WebSocketRequestMessage.get("/v1/devices/transfer_archive?timeout=${timeout.inWholeSeconds}")
    return NetworkResult
      .fromWebSocketRequest(
        signalWebSocket = authWebSocket,
        request = request,
        timeout = timeout,
        webSocketResponseConverter = NetworkResult.LongPollingWebSocketConverter(TransferArchiveResponse::class)
      )
  }
}
