/*
 * Copyright (C) 2023 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.platform.test.rule

import android.hardware.input.InputManager
import android.os.ParcelFileDescriptor
import android.platform.helpers.uinput.UInputDevice
import android.platform.uiautomatorhelpers.DeviceHelpers
import android.util.Log
import android.view.InputDevice
import androidx.core.content.getSystemService
import androidx.test.platform.app.InstrumentationRegistry
import com.google.common.util.concurrent.MoreExecutors
import com.google.gson.stream.JsonReader
import java.io.InputStreamReader
import java.util.concurrent.CountDownLatch
import java.util.concurrent.Executors
import java.util.concurrent.TimeUnit
import org.junit.runner.Description

/**
 * This rule allows end-to-end tests to add input devices through Uinput more easily. Additionally
 * it will wait for registration to complete and unregister devices after the test is complete.
 *
 * Sample usage:
 * ```
 * class InputDeviceTest {
 *     @get:Rule
 *     val inputDeviceRule = InputDeviceRule()
 *
 *     @Test
 *     fun testWithInputDevice() {
 *         inputDeviceRule.registerDevice(UinputKeyboard())
 *         // Continue test with input device added
 *     }
 * }
 * ```
 */
class InputDeviceRule : TestWatcher(), UInputDevice.EventInjector {

    private val inputManager = DeviceHelpers.context.getSystemService<InputManager>()!!
    private val deviceAddedMap = mutableMapOf<UInputDevice, CountDownLatch>()
    private val inputManagerDevices = mutableMapOf<DeviceId, UInputDevice>()

    private lateinit var inputStream: ParcelFileDescriptor.AutoCloseInputStream
    private lateinit var outputStream: ParcelFileDescriptor.AutoCloseOutputStream

    private val inputDeviceListenerDelegate =
        object : InputManager.InputDeviceListener {
            override fun onInputDeviceAdded(deviceId: Int) = updateInputDevice(DeviceId(deviceId))

            override fun onInputDeviceChanged(deviceId: Int) = updateInputDevice(DeviceId(deviceId))

            override fun onInputDeviceRemoved(deviceId: Int) {
                val deviceIdWrapped = DeviceId(deviceId)
                inputManagerDevices[deviceIdWrapped]?.let {
                    deviceAddedMap.remove(it)
                    inputManagerDevices.remove(deviceIdWrapped)
                }
            }
        }

    override fun starting(description: Description?) {
        super.starting(description)

        val (stdOut, stdIn) =
            InstrumentationRegistry.getInstrumentation()
                .uiAutomation
                .executeShellCommandRw("uinput -")

        inputStream = ParcelFileDescriptor.AutoCloseInputStream(stdOut)
        outputStream = ParcelFileDescriptor.AutoCloseOutputStream(stdIn)

        inputManager.registerInputDeviceListener(
            inputDeviceListenerDelegate,
            DeviceHelpers.context.mainThreadHandler,
        )
    }

    override fun finished(description: Description?) {
        inputStream.close()
        outputStream.close()

        inputManager.unregisterInputDeviceListener(inputDeviceListenerDelegate)
        deviceAddedMap.clear()
        inputManagerDevices.clear()
    }

    /**
     * Registers the provided device with Uinput. This call waits for
     * [InputManager.InputDeviceListener.onInputDeviceAdded] to be called before returning
     *
     * @throws RuntimeException if the device did not register successfully.
     */
    fun registerDevice(device: UInputDevice) {
        deviceAddedMap.putIfAbsent(device, CountDownLatch(1))

        writeCommand(device.getRegisterCommand())

        deviceAddedMap[device]!!.let { latch ->
            latch.await(20, TimeUnit.SECONDS)
            if (latch.count != 0L) {
                throw RuntimeException(
                    "Did not receive added notification for device ${device.name}"
                )
            }
        }
    }

    /** Send the [keycode] event, both key down and key up events, for the provided [deviceId]. */
    override fun sendKeyEvent(deviceId: Int, keycode: Int) {
        sendEventWithValues(deviceId, EV_KEY, keycode, KEY_DOWN)
        sendEventWithValues(deviceId, EV_KEY, keycode, KEY_UP)
    }

    /** Send the [eventType], [keycode], [value] for the provided [deviceId]. */
    override fun sendEventWithValues(deviceId: Int, eventType: Int, keycode: Int, value: Int) {
        injectEvdevEvents(deviceId, listOf(eventType, keycode, value, EV_SYN, SYN_REPORT, 0))
    }

    /**
     * The provided file should contain Uinput events in json format. The file path should be
     * relative to the assets root.
     */
    override fun sendEventsFromInputFile(deviceId: Int, inputFile: String) {
        InstrumentationRegistry.getInstrumentation().context.assets.open(inputFile).use {
            writeCommand(it.readBytes())
        }

        val executor = MoreExecutors.listeningDecorator(Executors.newCachedThreadPool())
        val future = executor.submit { waitForEmptyUinputEventQueue() }
        try {
            future.get(waitForSyncTokenDurationSeconds, TimeUnit.SECONDS)
        } catch (e: Exception) {
            error("Did not receive '$EVENT_REPLAY_COMPLETE_SYNC_TOKEN' within timeout")
        }
    }

    /**
     * Update the time base from which the following events are scheduled to the current time.
     *
     * See frameworks/base/cmds/uinput/README.md for a more thorough explanation.
     *
     * @param deviceId The id corresponding to [UInputDevice] to update the time base of
     */
    override fun updateTimeBase(deviceId: Int) {
        writeCommand("""{"command": "updateTimeBase", "id": $deviceId}""")
    }

    /**
     * Inject array of uinput events for a device. The following is an example of events: [[EV_KEY],
     * [KEY_UP], [KEY_DOWN], [EV_SYN], [SYN_REPORT], 0]. The number of entries in the provided
     * [evdevEvents] has to be a multiple of 3.
     *
     * @param deviceId The id corresponding to [UInputDevice] to associate with the [evdevEvents]
     * @param evdevEvents The uinput events to be injected
     */
    private fun injectEvdevEvents(deviceId: Int, evdevEvents: List<Int>) {
        assert(evdevEvents.size % 3 == 0) { "Number of injected events should be a multiple of 3" }

        writeCommand("""{"command": "inject","id": $deviceId,"events": $evdevEvents}""")
    }

    private fun writeCommand(command: String) {
        writeCommand(command.toByteArray())
    }

    private fun writeCommand(command: ByteArray) {
        outputStream.write(command)
        outputStream.flush()
    }

    private fun updateInputDevice(deviceId: DeviceId) {
        val device: InputDevice = inputManager.getInputDevice(deviceId.deviceId) ?: return
        val uinputDevice = deviceAddedMap.keys.firstOrNull { it.isInputDevice(device) } ?: return

        inputManagerDevices[deviceId] = uinputDevice
        deviceAddedMap[uinputDevice]!!.countDown()
    }

    private fun waitForEmptyUinputEventQueue() {
        writeCommand(EVENT_REPLAY_COMPLETE_SYNC_TOKEN_EVENT)

        Log.d(
            "InputDeviceRule",
            "'$EVENT_REPLAY_COMPLETE_SYNC_TOKEN' sync token sent, waiting for it to be processed...",
        )
        var nextSyncToken: String? = null
        while (nextSyncToken != EVENT_REPLAY_COMPLETE_SYNC_TOKEN) {
            nextSyncToken = readSyncTokenFromInputStream()
        }
        Log.d("InputDeviceRule", "'$EVENT_REPLAY_COMPLETE_SYNC_TOKEN' sync token processed")
    }

    /**
     * Returns the syncToken from the inputStream or null if the inputStream contains content other
     * than a syncToken.
     */
    private fun readSyncTokenFromInputStream(): String? {
        val reader = JsonReader(InputStreamReader(inputStream, Charsets.UTF_8))
        var reason: String? = null
        var syncToken: String? = null

        reader.beginObject()
        while (reader.hasNext()) {
            when (reader.nextName()) {
                "reason" -> reason = reader.nextString()
                "syncToken" -> syncToken = reader.nextString()
                else -> reader.skipValue()
            }
        }
        reader.endObject()

        if (reason != "sync" || syncToken == null) return null

        return syncToken
    }

    @JvmInline value class DeviceId(val deviceId: Int)

    private companion object {
        // See
        // https://cs.android.com/android/kernel/superproject/+/common-android-mainline:common/include/uapi/linux/input-event-codes.h
        // for these mappings.
        const val EV_KEY = 1
        const val EV_SYN = 0
        const val SYN_REPORT = 0
        const val KEY_UP = 0
        const val KEY_DOWN = 1

        const val EVENT_REPLAY_COMPLETE_SYNC_TOKEN = "event_replay_complete"
        val EVENT_REPLAY_COMPLETE_SYNC_TOKEN_EVENT =
            """
            {
                "id": 1,
                "command": "sync",
                "syncToken": "$EVENT_REPLAY_COMPLETE_SYNC_TOKEN"
            }

            {
              "id": 1,
              "command": "delay",
              "duration": 100
            }
            """
                .trimIndent()
        const val waitForSyncTokenDurationSeconds = 20L
    }
}