// Copyright (c) 2022 Huawei Technologies Co.,Ltd. All rights reserved.
//
// StratoVirt is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan
// PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//         http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

use std::sync::{Arc, Mutex, Weak};

use anyhow::{Context, Result};
use clap::Parser;
use log::{debug, info, warn};
use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};

use super::descriptor::{
    UsbConfigDescriptor, UsbDescConfig, UsbDescDevice, UsbDescEndpoint, UsbDescIface, UsbDescOther,
    UsbDescriptorOps, UsbDeviceDescriptor, UsbEndpointDescriptor, UsbInterfaceDescriptor,
};
use super::hid::{Hid, HidDevState, HidType, QUEUE_LENGTH, QUEUE_MASK};
use super::xhci::xhci_controller::{endpoint_number_to_id, XhciDevice};
use super::{config::*, USB_DEVICE_BUFFER_DEFAULT_LEN};
use super::{
    notify_controller, UsbDevState, UsbDevice, UsbDeviceBase, UsbDeviceRequest, UsbPacket,
    UsbPacketStatus,
};
use machine_manager::config::valid_id;
use migration::{DeviceStateDesc, MigrationHook, MigrationManager, StateTransfer};
use migration_derive::DescSerde;
use ui::input::{register_keyboard, unregister_keyboard, KeyboardOpts};
use util::gen_base_func;

/// Keyboard device descriptor
static DESC_DEVICE_KEYBOARD: Lazy<Arc<UsbDescDevice>> = Lazy::new(|| {
    Arc::new(UsbDescDevice {
        device_desc: UsbDeviceDescriptor {
            bLength: USB_DT_DEVICE_SIZE,
            bDescriptorType: USB_DT_DEVICE,
            idVendor: 0x0627,
            idProduct: USB_PRODUCT_ID_KEYBOARD,
            bcdDevice: 0,
            iManufacturer: STR_MANUFACTURER_INDEX,
            iProduct: STR_PRODUCT_KEYBOARD_INDEX,
            iSerialNumber: STR_SERIAL_KEYBOARD_INDEX,
            bcdUSB: 0x0100,
            bDeviceClass: 0,
            bDeviceSubClass: 0,
            bDeviceProtocol: 0,
            bMaxPacketSize0: 8,
            bNumConfigurations: 1,
        },
        configs: vec![Arc::new(UsbDescConfig {
            config_desc: UsbConfigDescriptor {
                bLength: USB_DT_CONFIG_SIZE,
                bDescriptorType: USB_DT_CONFIGURATION,
                wTotalLength: 0,
                bNumInterfaces: 1,
                bConfigurationValue: 1,
                iConfiguration: STR_CONFIG_KEYBOARD_INDEX,
                bmAttributes: USB_CONFIGURATION_ATTR_ONE | USB_CONFIGURATION_ATTR_REMOTE_WAKEUP,
                bMaxPower: 50,
            },
            iad_desc: vec![],
            interfaces: vec![DESC_IFACE_KEYBOARD.clone()],
        })],
    })
});
/// Keyboard interface descriptor
static DESC_IFACE_KEYBOARD: Lazy<Arc<UsbDescIface>> = Lazy::new(|| {
    Arc::new(UsbDescIface {
        interface_desc: UsbInterfaceDescriptor {
            bLength: USB_DT_INTERFACE_SIZE,
            bDescriptorType: USB_DT_INTERFACE,
            bInterfaceNumber: 0,
            bAlternateSetting: 0,
            bNumEndpoints: 1,
            bInterfaceClass: USB_CLASS_HID,
            bInterfaceSubClass: USB_SUBCLASS_BOOT,
            bInterfaceProtocol: USB_IFACE_PROTOCOL_KEYBOARD,
            iInterface: 0,
        },
        other_desc: vec![Arc::new(UsbDescOther {
            // HID descriptor
            data: vec![0x09, 0x21, 0x11, 0x01, 0x00, 0x01, 0x22, 0x3f, 0],
        })],
        endpoints: vec![Arc::new(UsbDescEndpoint {
            endpoint_desc: UsbEndpointDescriptor {
                bLength: USB_DT_ENDPOINT_SIZE,
                bDescriptorType: USB_DT_ENDPOINT,
                bEndpointAddress: USB_DIRECTION_DEVICE_TO_HOST | 0x1,
                bmAttributes: USB_ENDPOINT_ATTR_INT,
                wMaxPacketSize: 8,
                bInterval: 0xa,
            },
            extra: Vec::new(),
        })],
    })
});

/// String descriptor index
const STR_MANUFACTURER_INDEX: u8 = 1;
const STR_PRODUCT_KEYBOARD_INDEX: u8 = 2;
const STR_CONFIG_KEYBOARD_INDEX: u8 = 3;
const STR_SERIAL_KEYBOARD_INDEX: u8 = 4;

// Up flag.
const SCANCODE_UP: u16 = 0x80;
// Grey keys.
const SCANCODE_GREY: u16 = 0x80;
// Used to expand Grey keys.
const SCANCODE_EMUL0: u16 = 0xe0;

/// String descriptor
const DESC_STRINGS: [&str; 5] = [
    "",
    "StratoVirt",
    "StratoVirt USB Keyboard",
    "HID Keyboard",
    "1",
];

#[derive(Parser, Clone, Debug, Default)]
#[command(no_binary_name(true))]
pub struct UsbKeyboardConfig {
    #[arg(long)]
    pub classtype: String,
    #[arg(long, value_parser = valid_id)]
    id: String,
    #[arg(long)]
    bus: Option<String>,
    #[arg(long)]
    port: Option<String>,
}

/// USB keyboard device.
pub struct UsbKeyboard {
    base: UsbDeviceBase,
    hid: Hid,
    /// USB controller used to notify controller to transfer data.
    cntlr: Option<Weak<Mutex<XhciDevice>>>,
}

#[derive(DescSerde, Deserialize, Serialize)]
#[desc_version(current_version = "0.1.0")]
struct UsbKeyboardDevState {
    usb_state: UsbDevState,
    hid_state: HidDevState,
}

pub struct UsbKeyboardAdapter {
    usb_kbd: Arc<Mutex<UsbKeyboard>>,
}

impl KeyboardOpts for UsbKeyboardAdapter {
    fn do_key_event(&mut self, keycode: u16, down: bool) -> Result<()> {
        trace::usb_keyboard_event(&keycode, &down);

        let mut scan_codes = Vec::new();
        let mut keycode = keycode;
        if keycode & SCANCODE_GREY != 0 {
            scan_codes.push(u32::from(SCANCODE_EMUL0));
            keycode &= !SCANCODE_GREY;
        }

        if !down {
            keycode |= SCANCODE_UP;
        }
        scan_codes.push(u32::from(keycode));

        let mut locked_kbd = self.usb_kbd.lock().unwrap();
        if scan_codes.len() as u32 + locked_kbd.hid.num > QUEUE_LENGTH {
            trace::usb_keyboard_queue_full();
            // Return ok to ignore the request.
            return Ok(());
        }
        for code in scan_codes {
            let index = ((locked_kbd.hid.head + locked_kbd.hid.num) & QUEUE_MASK) as usize;
            locked_kbd.hid.num += 1;
            locked_kbd.hid.keyboard.keycodes[index] = code;
        }
        drop(locked_kbd);
        let clone_kbd = self.usb_kbd.clone();
        // Wakeup endpoint.
        let ep_id = endpoint_number_to_id(true, 1);
        notify_controller(&(clone_kbd as Arc<Mutex<dyn UsbDevice>>), ep_id)
    }
}

impl UsbKeyboard {
    pub fn new(config: UsbKeyboardConfig) -> Self {
        Self {
            base: UsbDeviceBase::new(config.id, USB_DEVICE_BUFFER_DEFAULT_LEN),
            hid: Hid::new(HidType::Keyboard),
            cntlr: None,
        }
    }
}

impl UsbDevice for UsbKeyboard {
    gen_base_func!(usb_device_base, usb_device_base_mut, UsbDeviceBase, base);

    fn realize(mut self) -> Result<Arc<Mutex<dyn UsbDevice>>> {
        self.base.reset_usb_endpoint();
        self.base.speed = USB_SPEED_FULL;
        let mut s: Vec<String> = DESC_STRINGS.iter().map(|&s| s.to_string()).collect();
        let prefix = &s[STR_SERIAL_KEYBOARD_INDEX as usize];
        s[STR_SERIAL_KEYBOARD_INDEX as usize] = self.base.generate_serial_number(prefix);
        self.base.init_descriptor(DESC_DEVICE_KEYBOARD.clone(), s)?;
        let id = self.device_id().to_string();
        let kbd = Arc::new(Mutex::new(self));
        let kbd_adapter = Arc::new(Mutex::new(UsbKeyboardAdapter {
            usb_kbd: kbd.clone(),
        }));
        register_keyboard(&id, kbd_adapter);

        MigrationManager::register_device_instance(
            UsbKeyboardDevState::descriptor(),
            kbd.clone(),
            &id,
        );

        Ok(kbd)
    }

    fn unrealize(&mut self) -> Result<()> {
        unregister_keyboard(self.device_id());

        MigrationManager::unregister_device_instance(
            UsbKeyboardDevState::descriptor(),
            self.device_id(),
        );

        Ok(())
    }

    fn cancel_packet(&mut self, _packet: &Arc<Mutex<UsbPacket>>) {}

    fn reset(&mut self) {
        info!("Keyboard device reset");
        self.base.remote_wakeup = 0;
        self.base.addr = 0;
        self.hid.reset();
    }

    fn handle_control(&mut self, packet: &Arc<Mutex<UsbPacket>>, device_req: &UsbDeviceRequest) {
        let mut locked_packet = packet.lock().unwrap();
        match self
            .base
            .handle_control_for_descriptor(&mut locked_packet, device_req)
        {
            Ok(handled) => {
                if handled {
                    debug!("Keyboard control handled by descriptor, return directly.");
                    return;
                }
            }
            Err(e) => {
                warn!(
                    "Received incorrect USB Keyboard descriptor message: {:?}",
                    e
                );
                locked_packet.status = UsbPacketStatus::Stall;
                return;
            }
        }
        self.hid
            .handle_control_packet(&mut locked_packet, device_req, &mut self.base.data_buf);
    }

    fn handle_data(&mut self, p: &Arc<Mutex<UsbPacket>>) {
        let mut locked_p = p.lock().unwrap();
        self.hid.handle_data_packet(&mut locked_p);
    }

    fn set_controller(&mut self, cntlr: Weak<Mutex<XhciDevice>>) {
        self.cntlr = Some(cntlr);
    }

    fn get_controller(&self) -> Option<Weak<Mutex<XhciDevice>>> {
        self.cntlr.clone()
    }
}

impl StateTransfer for UsbKeyboard {
    fn get_state_vec(&self) -> Result<Vec<u8>> {
        let state = UsbKeyboardDevState {
            usb_state: self.base.get_usb_state(),
            hid_state: self.hid.get_state()?,
        };

        Ok(serde_json::to_vec(&state)?)
    }

    fn set_state_mut(&mut self, state: &[u8], _version: u32) -> Result<()> {
        let usb_kbd_state: UsbKeyboardDevState =
            serde_json::from_slice(state).with_context(|| "Failed to get usb kbd state")?;

        self.base.set_usb_state(&usb_kbd_state.usb_state);
        self.hid.set_state(&usb_kbd_state.hid_state)?;

        Ok(())
    }

    fn get_device_alias(&self) -> u64 {
        MigrationManager::get_desc_alias(&UsbKeyboardDevState::descriptor().name).unwrap_or(!0)
    }
}

impl MigrationHook for UsbKeyboard {}
#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_keyboard_interface() {
        let interface_descriptor = &DESC_IFACE_KEYBOARD.clone().interface_desc;
        // bInterfaceClass: 3(HID)
        assert_eq!(interface_descriptor.bInterfaceClass, 3);
        // bInterfaceSubClass: 0(No Subclass) 1(Boot Interface Subclass) 2-255(Reserved)
        assert_eq!(interface_descriptor.bInterfaceSubClass, 1);
        // bInterfaceProtocol: 0(None) 1(Keyboard) 2(Mouse) 3-255(Reserved)
        assert_eq!(interface_descriptor.bInterfaceProtocol, 1);
    }

    #[test]
    fn test_usb_device_method() {
        let mut keyboard = UsbKeyboard::new(UsbKeyboardConfig {
            classtype: "usb-keyboard".to_string(),
            id: "keyboard".to_string(),
            bus: None,
            port: None,
        });

        let _ = &keyboard.reset();
        let _ = &keyboard.unrealize();
        let _ = &keyboard.get_controller();
        let device_req = UsbDeviceRequest {
            request_type: USB_DEVICE_OUT_REQUEST,
            request: USB_REQUEST_SET_ADDRESS,
            value: 0,
            index: 0,
            length: 0,
        };
        let target_dev =
            Arc::downgrade(&Arc::new(Mutex::new(keyboard))) as Weak<Mutex<dyn UsbDevice>>;
        let packet = Arc::new(Mutex::new(UsbPacket::new(
            1,
            u32::from(USB_TOKEN_OUT),
            0,
            0,
            Vec::new(),
            None,
            Some(target_dev),
        )));
        let mut keyboard = UsbKeyboard::new(UsbKeyboardConfig {
            classtype: "usb-keyboard".to_string(),
            id: "keyboard".to_string(),
            bus: None,
            port: None,
        });
        let _ = &keyboard.handle_control(&packet, &device_req);
        let _ = &keyboard.handle_data(&packet);
        let _ = &keyboard.cancel_packet(&packet);
        let _ = &keyboard.realize();
    }

    #[test]
    fn test_key_event() {
        let mut usb_adapter = UsbKeyboardAdapter {
            usb_kbd: Arc::new(Mutex::new(UsbKeyboard::new(UsbKeyboardConfig {
                classtype: "usb-keyboard".to_string(),
                id: "keyboard".to_string(),
                bus: None,
                port: None,
            }))),
        };
        // 0x0057: scancode of F11
        let _ = usb_adapter.do_key_event(0x0057, true);
        let _ = usb_adapter.do_key_event(0x0057, false);
    }

    #[test]
    fn test_usb_kbd_snapshot() {
        let config = UsbKeyboardConfig {
            classtype: "usb-keyboard".to_string(),
            id: "test_usb_keyboard".to_string(),
            bus: None,
            port: None,
        };

        let mut kbd_device = UsbKeyboard::new(config);

        let dev_state = kbd_device.get_state_vec();
        assert!(dev_state.is_ok());

        let ret = kbd_device.set_state_mut(dev_state.unwrap().as_ref(), 0_u32);
        assert!(ret.is_ok());
    }
}