use clap::Parser;
use clap_num::maybe_hex;
use regex::Regex;
use std::fs;
use std::os::fd::AsRawFd;
use std::vec::Vec;
use nix::{self, request_code_write};
use nix::libc;

const DATA_MAX: usize = 4096;

const MON_IOC_MAGIC: u8 = 0x92;
const MON_IOCQ_URB_LEN: u8 = 1;
const MON_IOCG_STATS: u8 = 3;
const MON_IOCX_GETX: u8 = 10;

#[derive(Parser, Debug, Default)]
#[command(author = "Trevor Bentley", version, about, long_about = None)]
#[command(help_template = "\
{name} v{version}, by {author-with-newline}
{about-with-newline}
{usage-heading} {usage}

{all-args}{after-help}
")]
pub struct CliArgs {
    /// Which USB bus to monitor (optional.  Default: all)
    #[arg(short, long)]
    pub bus: Option<u16>,
    /// USB device Vendor ID to monitor (optional.  Default: all)
    #[arg(short, long, value_parser=maybe_hex::<u16>)]
    pub vid: Option<u16>,
    /// USB device Product ID to monitor (optional.  Default: all)
    #[arg(short, long, value_parser=maybe_hex::<u16>)]
    pub pid: Option<u16>,
    /// Suppress packet output
    #[arg(short, long)]
    pub quiet: bool,
    /// Show field header
    #[arg(long)]
    pub header: bool,
}
impl CliArgs {
    pub fn args() -> CliArgs {
        CliArgs::parse()
    }
    pub fn dev_filter() -> DeviceFilter {
        let cli = CliArgs::args();
        DeviceFilter {
            bus: cli.bus,
            vid: cli.vid,
            pid: cli.pid,
        }
    }
}

pub type U8Cmp = fn(x: u8) -> bool;
pub type U16Cmp = fn(x: u16) -> bool;
pub type CIntCmp = fn(x: libc::c_int) -> bool;
pub type CUintCmp = fn(x: libc::c_uint) -> bool;
pub type SliceCmp = fn(x: &[u8]) -> bool;
pub type BoolCmp = fn(x: bool) -> bool;

#[derive(Default, Copy, Clone)]
pub struct DeviceFilter {
    pub bus: Option<u16>,
    pub vid: Option<u16>,
    pub pid: Option<u16>,
}

#[derive(Default, Copy, Clone)]
#[allow(dead_code)]
pub struct UsbmonFilter {
    // Capture behavior
    pub pkts_after: usize,    // packets to output after this packet

    // USB packet
    pub pkt_type: Option<U8Cmp>,
    pub xfer_type: Option<U8Cmp>,
    pub epnum: Option<U8Cmp>,
    pub devnum: Option<U8Cmp>,
    pub busnum: Option<U16Cmp>,
    pub flag_setup: Option<U8Cmp>,
    pub flag_data: Option<U8Cmp>,
    pub status: Option<CIntCmp>,
    pub data_len: Option<CUintCmp>,
    pub dir_out: Option<BoolCmp>,

    // USB data
    pub data: Option<SliceCmp>,

    // Setup packet
    pub setup_request_type: Option<U8Cmp>,
    pub setup_request: Option<U8Cmp>,
    pub setup_value: Option<U16Cmp>,
    pub setup_index: Option<U16Cmp>,
    pub setup_length: Option<U16Cmp>,
}


#[derive(Default)]
pub struct UsbmonDevice {
    bus: u16,
    dev: u8,

    vendor_id: u16,
    product_id: u16,
    manufacturer: String,
    product: String,
}
impl UsbmonDevice {
    pub fn device_match(&self, bus: u16, devnum: Option<u8>,
                        vid: Option<u16>, pid: Option<u16>) -> bool {
        let mut is_match = true;
        is_match &= bus == 0 || self.bus == bus;
        is_match &= devnum.map_or(true, |x| x == self.dev);
        is_match &= vid.map_or(true, |x| x == self.vendor_id);
        is_match &= pid.map_or(true, |x| x == self.product_id);
        is_match
    }
}

pub fn read_devices() -> Vec<UsbmonDevice> {
    let mut cur_dev: UsbmonDevice = UsbmonDevice::default();
    let mut all_devs: Vec<UsbmonDevice> = Vec::new();
    let mon_devices: &str = "/sys/kernel/debug/usb/devices";
    let contents = fs::read_to_string(mon_devices)
        .expect("debugfs devices file not found.  Is usbmod kernel loaded?");

    let t_re = Regex::new(r"^T:  Bus=([0-9]*) Lev=([0-9]*) Prnt=([0-9]*) Port=([0-9]*) Cnt=([0-9]*) Dev#=(...) .*$").unwrap();
    let p_re = Regex::new(r"^P:  Vendor=(....) ProdID=(....).*$").unwrap();
    let manu_re = Regex::new(r"^S:  Manufacturer=(.*)$").unwrap();
    let prod_re = Regex::new(r"^S:  Product=(.*)$").unwrap();

    for line in contents.lines() {
        if let Some(t_caps) = t_re.captures(line) {
            cur_dev.bus = t_caps.get(1).map_or("0", |s| s.as_str().trim()).parse::<u16>().unwrap_or(0);
            cur_dev.dev = t_caps.get(6).map_or("0", |s| s.as_str().trim()).parse::<u8>().unwrap_or(0);
        }

        if let Some(p_caps) = p_re.captures(line) {
            cur_dev.vendor_id = p_caps.get(1).map_or(0, |s| u16::from_str_radix(s.as_str(), 16).unwrap_or(0));
            cur_dev.product_id = p_caps.get(2).map_or(0, |s| u16::from_str_radix(s.as_str(), 16).unwrap_or(0));
        }

        if let Some(manu_caps) = manu_re.captures(line) {
            cur_dev.manufacturer = manu_caps.get(1).map_or("", |s| s.as_str().trim()).to_owned();
        }

        if let Some(prod_caps) = prod_re.captures(line) {
            cur_dev.product = prod_caps.get(1).map_or("", |s| s.as_str().trim()).to_owned();
            all_devs.push(cur_dev);
            cur_dev = Default::default();
        }
    }
    all_devs.sort_by_key(|x| x.bus);
    all_devs
}

pub fn print_devices(all_devs: &Vec<UsbmonDevice>) {
    println!("bus  dev   vid:pid    product   [manufacturer]");
    println!("-----------------------------------------------");
    for cur_dev in all_devs {
        println!("{:02}   {:03}   {:04x}:{:04x}   {}   [{}]",
                 cur_dev.bus, cur_dev.dev,
                 cur_dev.vendor_id, cur_dev.product_id,
                 cur_dev.product, cur_dev.manufacturer);
    }
}

//struct usbmon_packet {
//      u64 id;                 /*  0: URB ID - from submission to callback */
//      unsigned char type;     /*  8: Same as text; extensible. */
//      unsigned char xfer_type; /*    ISO (0), Intr, Control, Bulk (3) */
//      unsigned char epnum;    /*     Endpoint number and transfer direction */
//      unsigned char devnum;   /*     Device address */
//      u16 busnum;             /* 12: Bus number */
//      char flag_setup;        /* 14: Same as text */
//      char flag_data;         /* 15: Same as text; Binary zero is OK. */
//      s64 ts_sec;             /* 16: gettimeofday */
//      s32 ts_usec;            /* 24: gettimeofday */
//      int status;             /* 28: */
//      unsigned int length;    /* 32: Length of data (submitted or actual) */
//      unsigned int len_cap;   /* 36: Delivered length */
//      union {                 /* 40: */
//              unsigned char setup[SETUP_LEN]; /* Only for Control S-type */
//              struct iso_rec {                /* Only for ISO */
//                      int error_count;
//                      int numdesc;
//              } iso;
//      } s;
//      int interval;           /* 48: Only for Interrupt and ISO */
//      int start_frame;        /* 52: For ISO */
//      unsigned int xfer_flags; /* 56: copy of URB's transfer_flags */
//      unsigned int ndesc;     /* 60: Actual number of ISO descriptors */
//};                            /* 64 total length */

#[derive(Default, Copy, Clone)]
#[repr(C, packed)]
struct IsoRec {
    error_count: libc::c_int,
    numdesc: libc::c_int,
}

#[derive(Default, Copy, Clone)]
#[allow(non_snake_case)]
#[repr(C, packed)]
struct SetupPacket {
    bmRequestType: u8,
    bRequest: u8,
    wValue: u16,
    wIndex: u16,
    wLength: u16,
}
impl SetupPacket {
    fn request_type(&self) -> u8 {
        let x = self.bmRequestType;
        x
    }
    fn request(&self) -> u8 {
        let x = self.bRequest;
        x
    }
    fn value(&self) -> u16 {
        let x = self.wValue;
        x
    }
    fn index(&self) -> u16 {
        let x = self.wIndex;
        x
    }
    fn length(&self) -> u16 {
        let x = self.wLength;
        x
    }
}

#[derive(Copy, Clone)]
#[repr(C, packed)]
union SetupUnion {
    setup: SetupPacket,
    iso: IsoRec,
}
impl Default for SetupUnion {
    fn default() -> Self { SetupUnion { setup: Default::default() } }
}

#[derive(Default, Copy, Clone)]
#[repr(C, packed)]
struct UsbmonPacket {
    id: u64,
    pkt_type: u8,  // "S" (submit), "C" (complete)
    xfer_type: u8, // ISO (0), Interrupt (1), Control (2), Bulk (3)
    epnum: u8,     // direction encoded in highest bit, 0x80 == in
    devnum: u8,
    busnum: u16,
    flag_setup: u8, // 0 or "-" (setup packet present)
    flag_data: u8,  // 0, ">" (out no data), "=" (out data), or "<" (in)
    ts_sec: i64,
    ts_usec: i32,
    // status always -115 (EINPROGRESS) for submit events
    status: libc::c_int, // USB URB errno.  -71 == EPROTO, -115 == EINPROGRESS
    length: libc::c_uint,
    len_cap: libc::c_uint,
    s: SetupUnion,
    interval: libc::c_int,    // IRQ and ISO
    start_frame: libc::c_int, // ISO

    // from kernel's include/linux/usb.h:
    //
    // #define URB_SHORT_NOT_OK        0x0001  /* report short reads as errors */
    // #define URB_ISO_ASAP            0x0002  /* iso-only; use the first unexpired
    //                                          * slot in the schedule */
    // #define URB_NO_TRANSFER_DMA_MAP 0x0004  /* urb->transfer_dma valid on submit */
    // #define URB_ZERO_PACKET         0x0040  /* Finish bulk OUT with short packet */
    // #define URB_NO_INTERRUPT        0x0080  /* HINT: no non-error interrupt
    //                                          * needed */
    // #define URB_FREE_BUFFER         0x0100  /* Free transfer buffer with the URB */
    //
    // /* The following flags are used internally by usbcore and HCDs */
    // #define URB_DIR_IN              0x0200  /* Transfer from device to host */
    // #define URB_DIR_OUT             0
    // #define URB_DIR_MASK            URB_DIR_IN
    //
    // #define URB_DMA_MAP_SINGLE      0x00010000      /* Non-scatter-gather mapping */
    // #define URB_DMA_MAP_PAGE        0x00020000      /* HCD-unsupported S-G */
    // #define URB_DMA_MAP_SG          0x00040000      /* HCD-supported S-G */
    // #define URB_MAP_LOCAL           0x00080000      /* HCD-local-memory mapping */
    // #define URB_SETUP_MAP_SINGLE    0x00100000      /* Setup packet DMA mapped */
    // #define URB_SETUP_MAP_LOCAL     0x00200000      /* HCD-local setup packet */
    // #define URB_DMA_SG_COMBINED     0x00400000      /* S-G entries were combined */
    // #define URB_ALIGNED_TEMP_BUFFER 0x00800000      /* Temp buffer was alloc'd */
    xfer_flags: libc::c_uint, // bit field of above URB flags
    ndesc: libc::c_uint, // number of ISO descriptors
}

impl UsbmonPacket {
    fn setup_packet(&self) -> Option<&SetupPacket> {
        match self.flag_setup {
            b'-' => None,
            _ => unsafe {
                let setup = &self.s.setup;
                Some(setup)
            },
        }
    }
    fn pkt_type(&self) -> u8 {
        let x = self.pkt_type;
        x
    }
    fn xfer_type(&self) -> u8 {
        let x = self.xfer_type;
        x
    }
    fn busnum(&self) -> u16 {
        let x = self.busnum;
        x
    }
    fn devnum(&self) -> u8 {
        let x = self.devnum;
        x
    }
    fn epnum(&self) -> u8 {
        let x = self.epnum;
        x & 0x7f
    }
    fn dir_out(&self) -> bool {
        let epnum = self.epnum;
        match epnum & 0x80 { 0 => true, _ => false }
    }
    fn dir_chr(&self) -> char {
        let dir_out = self.dir_out();
        match dir_out { true => '>', _ => '<' }
    }
    fn flag_setup(&self) -> u8 {
        let x = self.flag_setup;
        x
    }
    fn flag_data(&self) -> u8 {
        let x = self.flag_data;
        x
    }
    fn status(&self) -> libc::c_int {
        let x = self.status;
        match self.pkt_type() {
            b'S' => 0,
            _ => x
        }
    }
    fn status_chr(&self) -> char {
        match self.pkt_type() {
            b'S' => ' ', // submit packets always have status EINPROGRESS
            _ => match self.status {
                0 => ' ',
                _ => 'e',
            },
        }
    }
    fn length(&self) -> libc::c_uint {
        let x = self.len_cap;
        x
    }
    fn length_available(&self) -> libc::c_uint {
        std::cmp::min(self.length(), DATA_MAX as u32)
    }
    fn ts_sec(&self) -> i64 {
        let x = self.ts_sec;
        x
    }
    fn ts_usec(&self) -> i64 {
        let x = self.ts_usec;
        x as i64
    }
    fn usec_since(&self, since: i64) -> i64 {
        let ts_usec = (self.ts_sec() * 1000 * 1000) + (self.ts_usec());
        ts_usec - since
    }

    fn is_error(&self) -> bool {
        // 'S' packets always return EINPROGRESS
        self.pkt_type != b'S' && self.status != 0
    }

    fn match_filter(&self, f: Option<&UsbmonFilter>, data: &[u8]) -> bool {
        let mut matched = true;
        if f.is_none() {
            return true;
        }
        let f = f.unwrap();

        matched &= f.pkt_type.map(|x| x(self.pkt_type())).unwrap_or(true);
        matched &= f.xfer_type.map(|x| x(self.xfer_type())).unwrap_or(true);
        matched &= f.epnum.map(|x| x(self.epnum())).unwrap_or(true);
        matched &= f.devnum.map(|x| x(self.devnum())).unwrap_or(true);
        matched &= f.busnum.map(|x| x(self.busnum())).unwrap_or(true);
        matched &= f.dir_out.map(|x| x(self.dir_out())).unwrap_or(true);
        matched &= f.flag_setup.map(|x| x(self.flag_setup())).unwrap_or(true);
        matched &= f.flag_data.map(|x| x(self.flag_data())).unwrap_or(true);
        matched &= f.status.map(|x| x(self.status())).unwrap_or(true);
        matched &= f.data_len.map(|x| x(self.length())).unwrap_or(true);

        matched &= f.data.map(|x| x(data)).unwrap_or(true);

        matched &= f.setup_request_type
            .map(|x| x(self.setup_packet()
                       .map(|s| s.request_type()).unwrap_or(0))).unwrap_or(true);
        matched &= f.setup_request
            .map(|x| x(self.setup_packet()
                       .map(|s| s.request()).unwrap_or(0))).unwrap_or(true);
        matched &= f.setup_value
            .map(|x| x(self.setup_packet()
                       .map(|s| s.value()).unwrap_or(0))).unwrap_or(true);
        matched &= f.setup_index
            .map(|x| x(self.setup_packet()
                       .map(|s| s.index()).unwrap_or(0))).unwrap_or(true);
        matched &= f.setup_length
            .map(|x| x(self.setup_packet()
                       .map(|s| s.length()).unwrap_or(0))).unwrap_or(true);

        matched
    }

    fn print_hdr() {
        println!("   usec    bs:dv:ep dir pkt:xfr  [len]: |setup| data");
        println!("------------------------------------------------------------------------");
    }
    fn print(&self, start_time: i64, pkt_data: &[u8; DATA_MAX]) {
        print!("{:010} {:02}:{:02}:{:02}  {}   {}:{:02} {} [{:03}]: ",
               self.usec_since(start_time),
               self.busnum(),
               self.devnum(),
               self.epnum(),
               self.dir_chr(),
               self.pkt_type(),
               self.xfer_type(),
               self.status_chr(),
               self.length());
        if let Some(setup_pkt) = self.setup_packet() {
            print!("|{:02x} {:02x} {:04x} {:04x} {:02}| ",
                   setup_pkt.request_type(),
                   setup_pkt.request(),
                   setup_pkt.value(),
                   setup_pkt.index(),
                   setup_pkt.length());
        }
        // print data
        if self.length() > 0 {
            for i in (0..self.length_available()).step_by(16) {
                if i != 0 || self.setup_packet().is_some() {
                    println!();
                    print!("                                        ")
                }
                let max = std::cmp::min(16, self.length_available() - i);
                for j in 0..max {
                    print!("{:02x}", pkt_data[(i + j) as usize]);
                }
            }
        }
        println!();
    }
}

#[derive(Default, Copy, Clone)]
#[repr(C, packed)]
struct MonStats {
    queued: u32,
    dropped: u32,
}

#[derive(Copy, Clone)]
#[repr(C, packed(1))]
struct MonGetArg {
    hdr: *mut UsbmonPacket,
    data: *mut u8,
    alloc: libc::size_t,
}
impl Default for MonGetArg {
    fn default() -> Self { MonGetArg {
        hdr: std::ptr::null::<UsbmonPacket>() as *mut UsbmonPacket,
        data: std::ptr::null::<u8>() as *mut u8,
        alloc: 0,
    } }
}


pub fn monitor(dev_filter: Option<&DeviceFilter>,
               start_filter: Option<&UsbmonFilter>,
               active_filter: Option<&UsbmonFilter>,
               end_filter: Option<&UsbmonFilter>) -> i64 {
    let cli_args = CliArgs::args();
    let bus: u16 = match dev_filter {
        None => 0,
        Some(f) => f.bus.unwrap_or(0),
    };
    let bus_devfile: String = format!("/dev/usbmon{}", bus);
    let bus_dev = std::fs::File::open(&bus_devfile)
        .expect(&format!("device {} not found.  is usbmon module loaded?", &bus_devfile));
    let bus_dev_fd = bus_dev.as_raw_fd();

    let vid = match dev_filter {
        None => None,
        Some(f) => f.vid,
    };
    let pid = match dev_filter {
        None => None,
        Some(f) => f.pid,
    };

    nix::ioctl_none!(usbmon_read_len, MON_IOC_MAGIC, MON_IOCQ_URB_LEN);
    nix::ioctl_read!(usbmon_read_stats, MON_IOC_MAGIC, MON_IOCG_STATS, MonStats);

    // Due to a bug in either the packing or the size_of operator, the
    // size of 24 bytes has to be explicitly set here to calculate the
    // ioctl because the struct reports as 32 bytes even when packed.
    //
    //nix::ioctl_write_ptr!(usbmon_read_data, MON_IOC_MAGIC, MON_IOCX_GETX, MonGetArg);
    // request_code_write!(MON_IOC_MAGIC, MON_IOCX_GETX, std::mem::size_of::<MonGetArg>())
    nix::ioctl_write_ptr_bad!(usbmon_read_data, request_code_write!(MON_IOC_MAGIC, MON_IOCX_GETX, 24), MonGetArg);

    let mut _next_len: i32;
    let mut mon_stats = Default::default();
    let mut pkt_data: [u8; DATA_MAX] = [0; DATA_MAX];
    let mut mon_pkt: UsbmonPacket = Default::default();
    let mon_arg = MonGetArg {
        hdr: &mut mon_pkt,
        data: &mut pkt_data as *mut u8,
        alloc: DATA_MAX,
    };

    let mut all_devs: Vec<UsbmonDevice> = read_devices();
    let mut start_time: i64 = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default().as_micros() as i64;
    let last_time: i64;

    let mut capturing: bool = false;
    let mut stop_capturing_count: Option<usize> = None;

    if cli_args.header {
        UsbmonPacket::print_hdr();
    }
    loop {
        unsafe { usbmon_read_stats(bus_dev_fd, &mut mon_stats).expect("stats ioctl failed"); }
        let queued = mon_stats.queued;
        let _dropped = mon_stats.dropped;

        if queued == 0 {
            std::thread::sleep(std::time::Duration::from_millis(20));
            continue;
        }

        unsafe {
            // unused, just to prove that ioctl is working
            _next_len = usbmon_read_len(bus_dev_fd).expect("len ioctl failed");

            // get the actual event
            usbmon_read_data(bus_dev_fd, &mon_arg).expect("pkt ioctl failed");

            let hdr = mon_arg.hdr;

            // Check if this packet's bus/dev combo has been enumerated yet.  Re-enumerate
            // if it hasn't or if the packet indicates a bus error (which can be device
            // connects and disconnects).
            if !all_devs.iter().any(|d| d.device_match((*hdr).busnum(),
                                                       Some((*hdr).devnum()),
                                                       None, None)) ||
                (*hdr).is_error() {
                    all_devs = read_devices();
                }

            // See if we have enumerated a device matching the registered capture filter,
            // and if this packet matches it.
            let matching: Vec<&UsbmonDevice> = all_devs.iter().filter(
                |d| d.device_match(bus, None, vid, pid)).collect();
            match matching.len() {
                0 => {
                    // No known device matching the capture filter.  Maybe it isn't
                    // connected yet.
                    continue;
                },
                n if n == all_devs.len() => {
                    // All devices matched means no filter, continue
                },
                _ => {
                    // Multiple devices matched
                    let mut matched = false;
                    for dev in matching {
                        // Device is enumerated, did it produce this packet?
                        if dev.device_match((*hdr).busnum(),
                                            Some((*hdr).devnum()),
                                            dev_filter.map_or(None, |f| f.vid),
                                            dev_filter.map_or(None, |f| f.pid)) {
                            matched = true;
                            break;
                        }
                    }

                    if !matched {
                        // This packet isn't from a matching device.
                        continue;
                    }

                    // Packet matches filter, continue.
                }
            }

            if !capturing && (*hdr).match_filter(start_filter, &pkt_data) {
                // If the start filter matches, begin capturing packets
                capturing = true;
                start_time = (*hdr).usec_since(0);
                if !cli_args.quiet {
                    (*hdr).print(start_time, &pkt_data);
                }
            }
            else if capturing {
                if end_filter.is_some() && (*hdr).match_filter(end_filter, &pkt_data) {
                    if !cli_args.quiet {
                        (*hdr).print(start_time, &pkt_data);
                    }
                    if let Some(f) = end_filter {
                        stop_capturing_count = Some(f.pkts_after);
                    }
                }
                else if (*hdr).match_filter(active_filter, &pkt_data) {
                    if !cli_args.quiet {
                        (*hdr).print(start_time, &pkt_data);
                    }
                }
            }

            if let Some(cnt) = stop_capturing_count {
                if cnt == 0 {
                    last_time = (*hdr).usec_since(start_time);
                    break;
                }
                let cnt = cnt - 1;
                stop_capturing_count = Some(cnt);
            }
        }
    }
    last_time
}