/*
 * Copyright (C) 2024 The Android Open Source Project
 *
 * 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.
 */
package android.bluetooth

import android.Manifest
import android.annotation.SuppressLint
import android.bluetooth.test_utils.EnableBluetoothRule
import android.content.BroadcastReceiver
import android.content.Context
import android.content.Intent
import android.content.IntentFilter
import android.platform.test.annotations.RequiresFlagsEnabled
import android.platform.test.flag.junit.CheckFlagsRule
import android.platform.test.flag.junit.DeviceFlagsValueProvider
import android.util.Log
import androidx.test.core.app.ApplicationProvider
import androidx.test.ext.junit.runners.AndroidJUnit4
import androidx.test.platform.app.InstrumentationRegistry
import com.android.bluetooth.flags.Flags
import com.android.compatibility.common.util.AdoptShellPermissionsRule
import com.google.common.truth.Truth
import com.google.protobuf.ByteString
import java.io.IOException
import java.time.Duration
import java.util.UUID
import java.util.concurrent.TimeUnit
import kotlin.concurrent.thread
import kotlinx.coroutines.*
import kotlinx.coroutines.channels.*
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.SharingStarted
import kotlinx.coroutines.flow.callbackFlow
import kotlinx.coroutines.flow.filter
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.flow.shareIn
import kotlinx.coroutines.withTimeout
import org.junit.After
import org.junit.Before
import org.junit.Ignore
import org.junit.Rule
import org.junit.Test
import org.junit.runner.RunWith
import org.mockito.kotlin.argumentCaptor
import org.mockito.kotlin.eq
import org.mockito.kotlin.mock
import org.mockito.kotlin.timeout
import org.mockito.kotlin.verify
import pandora.BumbleConfigProto
import pandora.HostProto
import pandora.RfcommProto
import pandora.RfcommProto.ServerId

@SuppressLint("MissingPermission")
@RunWith(AndroidJUnit4::class)
@ExperimentalCoroutinesApi
class RfcommTest {
    private val mContext = ApplicationProvider.getApplicationContext<Context>()
    private val mManager = mContext.getSystemService(BluetoothManager::class.java)
    private val mAdapter = mManager!!.adapter

    @Rule(order = 0)
    @JvmField
    val mCheckFlagsRule: CheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule()

    // Gives shell permissions during the test.
    @Rule(order = 1)
    @JvmField
    val mPermissionsRule =
        AdoptShellPermissionsRule(
            InstrumentationRegistry.getInstrumentation().getUiAutomation(),
            Manifest.permission.BLUETOOTH_CONNECT,
            Manifest.permission.BLUETOOTH_PRIVILEGED,
            Manifest.permission.MODIFY_PHONE_STATE,
        )

    // Set up a Bumble Pandora device for the duration of the test.
    @Rule(order = 2) @JvmField val mBumble = PandoraDevice()

    @Rule(order = 3) @JvmField val enableBluetoothRule = EnableBluetoothRule(false, true)

    private lateinit var mRemoteDevice: BluetoothDevice
    private lateinit var mHost: Host
    private var mConnectionCounter = 1
    private var mProfileServiceListener = mock<BluetoothProfile.ServiceListener>()

    private val mFlow: Flow<Intent>
    private val mScope: CoroutineScope = CoroutineScope(Dispatchers.Default.limitedParallelism(2))

    @OptIn(ExperimentalStdlibApi::class)
    private val bdAddrFormat = HexFormat { bytes { byteSeparator = ":" } }
    @OptIn(ExperimentalStdlibApi::class)
    private val mLocalAddress: ByteString =
        ByteString.copyFrom("DA:4C:10:DE:17:00".hexToByteArray(bdAddrFormat))

    init {
        val intentFilter = IntentFilter(BluetoothDevice.ACTION_PAIRING_REQUEST)
        mFlow = intentFlow(mContext, intentFilter, mScope).shareIn(mScope, SharingStarted.Eagerly)
    }

    /*
        Setup:
        1. Initialize host and mRemoteDevice
        2. Override pairing config to enable insecure tests
        3. Disable A2DP, HFP, and HID profiles
        4. Disconnect devices, if they are connected
    */
    @Before
    fun setUp() {
        mRemoteDevice = mBumble.remoteDevice
        mHost = Host(mContext)

        // Set Bonding
        val pairingConfig =
            BumbleConfigProto.PairingConfig.newBuilder()
                .setBonding(false)
                .setMitm(false)
                .setSc(false)
                .setIdentityAddressType(HostProto.OwnAddressType.PUBLIC)
                .build()
        val overrideRequest =
            BumbleConfigProto.OverrideRequest.newBuilder().setPairingConfig(pairingConfig).build()
        mBumble.bumbleConfigBlocking().override(overrideRequest)

        val bluetoothA2dp = getProfileProxy(mContext, BluetoothProfile.A2DP) as BluetoothA2dp
        bluetoothA2dp.setConnectionPolicy(
            mRemoteDevice,
            BluetoothProfile.CONNECTION_POLICY_FORBIDDEN,
        )
        val bluetoothHfp = getProfileProxy(mContext, BluetoothProfile.HEADSET) as BluetoothHeadset
        bluetoothHfp.setConnectionPolicy(
            mRemoteDevice,
            BluetoothProfile.CONNECTION_POLICY_FORBIDDEN,
        )
        val bluetoothHidHost =
            getProfileProxy(mContext, BluetoothProfile.HID_HOST) as BluetoothHidHost
        bluetoothHidHost.setConnectionPolicy(
            mRemoteDevice,
            BluetoothProfile.CONNECTION_POLICY_FORBIDDEN,
        )
        if (mRemoteDevice.isConnected) {
            mHost.disconnectAndVerify(mRemoteDevice)
        }
    }

    /*
        TearDown:
        1. remove bond
        2. shutdown host
    */
    @After
    fun tearDown() {
        if (mAdapter.bondedDevices.contains(mRemoteDevice)) {
            mHost.removeBondAndVerify(mRemoteDevice)
        }
        mHost.close()
    }

    /*
       Test Steps:
       1. Create an insecure socket
       2. Connect to the socket
       3. Verify that devices are connected.
    */
    @Test
    fun clientConnectToOpenServerSocketInsecure() {
        startServer { serverId -> createConnectAcceptSocket(isSecure = false, serverId) }
    }

    /*
       Test Steps:
       1. Create an secure socket
       2. Connect to the socket
       3. Verify that devices are connected.
    */
    @Test
    fun clientConnectToOpenServerSocketSecure() {
        startServer { serverId -> createConnectAcceptSocket(isSecure = true, serverId) }
    }

    /*
        Test Steps:
        1. Create an insecure socket
        2. Connect to the socket
        3. Verify that devices are connected
        4. Write data to socket output stream
        5. Verify bumble received that data
    */
    @Test
    fun clientSendDataOverInsecureSocket() {
        startServer { serverId ->
            val (insecureSocket, connection) = createConnectAcceptSocket(isSecure = false, serverId)
            val data: ByteArray = "Test data for clientSendDataOverInsecureSocket".toByteArray()
            val socketOs = insecureSocket.outputStream

            socketOs.write(data)
            val rxResponse: RfcommProto.RxResponse =
                mBumble
                    .rfcommBlocking()
                    .withDeadlineAfter(GRPC_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)
                    .receive(RfcommProto.RxRequest.newBuilder().setConnection(connection).build())
            Truth.assertThat(rxResponse.data).isEqualTo(ByteString.copyFrom(data))
        }
    }

    /*
        Test Steps:
        1. Create a secure socket
        2. Connect to the socket
        3. Verify that devices are connected
        4. Write data to socket output stream
        5. Verify remote device received that data
    */
    @Test
    fun clientSendDataOverSecureSocket() {
        startServer { serverId ->
            val (secureSocket, connection) = createConnectAcceptSocket(isSecure = true, serverId)
            val data: ByteArray = "Test data for clientSendDataOverSecureSocket".toByteArray()
            val socketOs = secureSocket.outputStream

            socketOs.write(data)
            val rxResponse: RfcommProto.RxResponse =
                mBumble
                    .rfcommBlocking()
                    .withDeadlineAfter(GRPC_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)
                    .receive(RfcommProto.RxRequest.newBuilder().setConnection(connection).build())
            Truth.assertThat(rxResponse.data).isEqualTo(ByteString.copyFrom(data))
        }
    }

    /*
        Test Steps:
        1. Create an insecure socket
        2. Connect to the socket
        3. Verify that devices are connected
        4. Send data from remote device
        5. Read and verify data from socket input stream
    */
    @Test
    fun clientReceiveDataOverInsecureSocket() {
        startServer { serverId ->
            val (insecureSocket, connection) = createConnectAcceptSocket(isSecure = false, serverId)
            val buffer = ByteArray(64)
            val socketIs = insecureSocket.inputStream
            val data: ByteString =
                ByteString.copyFromUtf8("Test data for clientReceiveDataOverInsecureSocket")

            val txRequest =
                RfcommProto.TxRequest.newBuilder().setConnection(connection).setData(data).build()
            mBumble.rfcommBlocking().send(txRequest)
            val numBytesFromBumble = socketIs.read(buffer)
            Truth.assertThat(ByteString.copyFrom(buffer).substring(0, numBytesFromBumble))
                .isEqualTo(data)
        }
    }

    /*
        Test Steps:
        1. Create a secure socket
        2. Connect to the socket
        3. Verify that devices are connected
        4. Send data from remote device
        5. Read and verify data from socket input stream
    */
    @Test
    fun clientReceiveDataOverSecureSocket() {
        startServer { serverId ->
            val (secureSocket, connection) = createConnectAcceptSocket(isSecure = true, serverId)
            val buffer = ByteArray(64)
            val socketIs = secureSocket.inputStream
            val data: ByteString =
                ByteString.copyFromUtf8("Test data for clientReceiveDataOverSecureSocket")

            val txRequest =
                RfcommProto.TxRequest.newBuilder().setConnection(connection).setData(data).build()
            mBumble.rfcommBlocking().send(txRequest)
            val numBytesFromBumble = socketIs.read(buffer)
            Truth.assertThat(ByteString.copyFrom(buffer).substring(0, numBytesFromBumble))
                .isEqualTo(data)
        }
    }

    /*
        Test Steps:
        1. Create insecure socket 1
        2. Create insecure socket 2
        3. Remote device initiates connection to socket 1
        4. Remote device initiates connection to socket 2
        5. Accept socket 1 and verify connection
        6. Accept socket 2 and verify connection
    */
    @Test
    fun connectTwoInsecureClientsSimultaneously() {
        startServer("ServerPort1", TEST_UUID) { serverId1 ->
            startServer("ServerPort2", SERIAL_PORT_UUID) { serverId2 ->
                val socket1 = createSocket(mRemoteDevice, isSecure = false, TEST_UUID)
                val socket2 = createSocket(mRemoteDevice, isSecure = false, SERIAL_PORT_UUID)

                acceptSocket(serverId1)
                Truth.assertThat(socket1.isConnected).isTrue()

                acceptSocket(serverId2)
                Truth.assertThat(socket2.isConnected).isTrue()
            }
        }
    }

    /*
        Test Steps:
        1. Create insecure socket 1
        2. Remote device initiates connection to socket 1
        3. Accept socket 1 and verify connection
        4. Repeat for socket 2
    */
    @Test
    fun connectTwoInsecureClientsSequentially() {
        startServer("ServerPort1", TEST_UUID) { serverId1 ->
            startServer("ServerPort2", SERIAL_PORT_UUID) { serverId2 ->
                val socket1 = createSocket(mRemoteDevice, isSecure = false, TEST_UUID)
                acceptSocket(serverId1)
                Truth.assertThat(socket1.isConnected).isTrue()

                val socket2 = createSocket(mRemoteDevice, isSecure = false, SERIAL_PORT_UUID)
                acceptSocket(serverId2)
                Truth.assertThat(socket2.isConnected).isTrue()
            }
        }
    }

    /*
        Test Steps:
        1. Create secure socket 1
        2. Create secure socket 2
        3. Remote device initiates connection to socket 1
        4. Remote device initiates connection to socket 2
        5. Accept socket 1 and verify connection
        6. Accept socket 2 and verify connection
    */
    @Test
    fun connectTwoSecureClientsSimultaneously() {
        startServer("ServerPort1", TEST_UUID) { serverId1 ->
            startServer("ServerPort2", SERIAL_PORT_UUID) { serverId2 ->
                val socket2 = createSocket(mRemoteDevice, isSecure = true, SERIAL_PORT_UUID)
                val socket1 = createSocket(mRemoteDevice, isSecure = true, TEST_UUID)

                acceptSocket(serverId1)
                Truth.assertThat(socket1.isConnected).isTrue()

                acceptSocket(serverId2)
                Truth.assertThat(socket2.isConnected).isTrue()
            }
        }
    }

    /*
        Test Steps:
        1. Create insecure socket 1
        2. Remote device initiates connection to socket 1
        3. Accept socket 1 and verify connection
        4. Repeat for socket 2
    */
    @Test
    fun connectTwoSecureClientsSequentially() {
        startServer("ServerPort1", TEST_UUID) { serverId1 ->
            startServer("ServerPort2", SERIAL_PORT_UUID) { serverId2 ->
                val socket1 = createSocket(mRemoteDevice, isSecure = true, TEST_UUID)
                acceptSocket(serverId1)
                Truth.assertThat(socket1.isConnected).isTrue()

                val socket2 = createSocket(mRemoteDevice, isSecure = true, SERIAL_PORT_UUID)
                acceptSocket(serverId2)
                Truth.assertThat(socket2.isConnected).isTrue()
            }
        }
    }

    /*
        Test Steps:
        1. Create insecure socket 1
        2. Remote device initiates connection to socket 1
        3. Accept socket 1 and verify connection
        4. Repeat for secure socket 2
    */
    @Test
    @Ignore("b/380091558")
    fun connectTwoMixedClientsInsecureThenSecure() {
        startServer("ServerPort1", TEST_UUID) { serverId1 ->
            startServer("ServerPort2", SERIAL_PORT_UUID) { serverId2 ->
                val socket2 = createSocket(mRemoteDevice, isSecure = false, SERIAL_PORT_UUID)
                acceptSocket(serverId2)
                Truth.assertThat(socket2.isConnected).isTrue()
                Log.i(TAG, "Finished with socket number 2")
                val socket1 = createSocket(mRemoteDevice, isSecure = true, TEST_UUID)
                acceptSocket(serverId1)
                Truth.assertThat(socket1.isConnected).isTrue()
            }
        }
    }

    /*
        Test Steps:
        1. Create secure socket 2
        2. Remote device initiates connection to socket 2
        3. Accept socket 2 and verify connection
        4. Repeat for insecure socket 1
    */
    @Test
    fun connectTwoMixedClientsSecureThenInsecure() {
        startServer("ServerPort1", TEST_UUID) { serverId1 ->
            startServer("ServerPort2", SERIAL_PORT_UUID) { serverId2 ->
                val socket2 = createSocket(mRemoteDevice, isSecure = true, SERIAL_PORT_UUID)
                acceptSocket(serverId2)
                Truth.assertThat(socket2.isConnected).isTrue()

                val socket1 = createSocket(mRemoteDevice, isSecure = false, TEST_UUID)
                acceptSocket(serverId1)
                Truth.assertThat(socket1.isConnected).isTrue()
            }
        }
    }

    /*
      Test Steps:
      1. Create listening socket and connect
      2. Disconnect RFCOMM from remote device
    */
    @RequiresFlagsEnabled(Flags.FLAG_TRIGGER_SEC_PROC_ON_INC_ACCESS_REQ)
    @Test
    fun serverSecureConnectThenRemoteDisconnect() {
        // step 1
        val (serverSock, connection) = connectRemoteToListeningSocket()
        val disconnectRequest =
            RfcommProto.DisconnectionRequest.newBuilder().setConnection(connection).build()
        // step 2
        mBumble.rfcommBlocking().disconnect(disconnectRequest)
        Truth.assertThat(serverSock.channel).isEqualTo(-1) // ensure disconnected at RFCOMM Layer
    }

    /*
      Test Steps:
      1. Create listening socket and connect
      2. Disconnect RFCOMM from local device
    */
    @RequiresFlagsEnabled(Flags.FLAG_TRIGGER_SEC_PROC_ON_INC_ACCESS_REQ)
    @Test
    fun serverSecureConnectThenLocalDisconnect() {
        // step 1
        val (serverSock, _) = connectRemoteToListeningSocket()
        // step 2
        serverSock.close()
        Truth.assertThat(serverSock.channel).isEqualTo(-1) // ensure disconnected at RFCOMM Layer
    }

    private fun createConnectAcceptSocket(
        isSecure: Boolean,
        server: ServerId,
        uuid: String = TEST_UUID,
    ): Pair<BluetoothSocket, RfcommProto.RfcommConnection> {
        val socket = createSocket(mRemoteDevice, isSecure, uuid)
        val connection = acceptSocket(server)
        Truth.assertThat(socket.isConnected).isTrue()

        return Pair(socket, connection)
    }

    private fun createSocket(
        device: BluetoothDevice,
        isSecure: Boolean,
        uuid: String,
    ): BluetoothSocket {
        val socket =
            if (isSecure) {
                device.createRfcommSocketToServiceRecord(UUID.fromString(uuid))
            } else {
                device.createInsecureRfcommSocketToServiceRecord(UUID.fromString(uuid))
            }

        runBlocking(mScope.coroutineContext) {
            withTimeout(CONNECT_TIMEOUT.toMillis()) {
                // We need to reply to the pairing request in the case where the devices aren't
                // bonded yet
                if (isSecure && !mAdapter.bondedDevices.contains(device)) {
                    launch {
                        Log.i(TAG, "Waiting for ACTION_PAIRING_REQUEST")
                        mFlow
                            .filter { it.action == BluetoothDevice.ACTION_PAIRING_REQUEST }
                            .filter { it.getBluetoothDeviceExtra() == device }
                            .first()
                        device.setPairingConfirmation(true)
                    }
                }
                socket.connect()
            }
        }
        return socket
    }

    private fun acceptSocket(server: ServerId): RfcommProto.RfcommConnection {
        val connectionResponse =
            mBumble
                .rfcommBlocking()
                .withDeadlineAfter(GRPC_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)
                .acceptConnection(
                    RfcommProto.AcceptConnectionRequest.newBuilder().setServer(server).build()
                )
        Truth.assertThat(connectionResponse.connection.id).isEqualTo(mConnectionCounter)

        mConnectionCounter += 1
        return connectionResponse.connection
    }

    private fun startServer(
        name: String = TEST_SERVER_NAME,
        uuid: String = TEST_UUID,
        block: (ServerId) -> Unit,
    ) {
        val request =
            RfcommProto.StartServerRequest.newBuilder().setName(name).setUuid(uuid).build()
        val response = mBumble.rfcommBlocking().startServer(request)

        try {
            block(response.server)
        } finally {
            mBumble
                .rfcommBlocking()
                .withDeadlineAfter(GRPC_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS)
                .stopServer(
                    RfcommProto.StopServerRequest.newBuilder().setServer(response.server).build()
                )
        }
    }

    private fun connectRemoteToListeningSocket(
        name: String = TEST_SERVER_NAME,
        uuid: String = TEST_UUID,
    ): Pair<BluetoothServerSocket, RfcommProto.RfcommConnection> {
        var connection: RfcommProto.RfcommConnection? = null
        val connectRequest =
            RfcommProto.ConnectionRequest.newBuilder()
                .setAddress(mLocalAddress)
                .setUuid(uuid)
                .build()
        val t = thread {
            val connectResponse = mBumble.rfcommBlocking().connectToServer(connectRequest)
            connection = connectResponse.connection
        }
        val socket = mAdapter.listenUsingRfcommWithServiceRecord(name, UUID.fromString(uuid))

        try {
            socket.accept(3000) // 3 second timeout
        } catch (e: IOException) {
            Log.e(TAG, "Unexpected IOException: $e")
        }
        t.join()
        Truth.assertThat(connection).isNotNull()

        return Pair(socket, connection!!)
    }

    private fun getProfileProxy(context: Context, profile: Int): BluetoothProfile {
        mAdapter.getProfileProxy(context, mProfileServiceListener, profile)
        val proxyCaptor = argumentCaptor<BluetoothProfile>()
        verify(mProfileServiceListener, timeout(GRPC_TIMEOUT.toMillis()))
            .onServiceConnected(eq(profile), proxyCaptor.capture())
        return proxyCaptor.lastValue
    }

    fun Intent.getBluetoothDeviceExtra(): BluetoothDevice =
        this.getParcelableExtra(BluetoothDevice.EXTRA_DEVICE, BluetoothDevice::class.java)!!

    @kotlinx.coroutines.ExperimentalCoroutinesApi
    fun intentFlow(context: Context, intentFilter: IntentFilter, scope: CoroutineScope) =
        callbackFlow {
            val broadcastReceiver: BroadcastReceiver =
                object : BroadcastReceiver() {
                    override fun onReceive(context: Context, intent: Intent) {
                        Log.d(TAG, "intentFlow: onReceive: ${intent.action}")
                        scope.launch { trySendBlocking(intent) }
                    }
                }
            context.registerReceiver(broadcastReceiver, intentFilter, Context.RECEIVER_EXPORTED)

            awaitClose { context.unregisterReceiver(broadcastReceiver) }
        }

    companion object {
        private val TAG = RfcommTest::class.java.getSimpleName()
        private val GRPC_TIMEOUT = Duration.ofSeconds(10)
        private val CONNECT_TIMEOUT = Duration.ofSeconds(7)
        private const val TEST_UUID = "2ac5d8f1-f58d-48ac-a16b-cdeba0892d65"
        private const val SERIAL_PORT_UUID = "00001101-0000-1000-8000-00805F9B34FB"
        private const val TEST_SERVER_NAME = "RFCOMM Server"
    }
}