// -*- coding: utf-8 -*-
//
// pwman-f2e
//
// Copyright 2023 Michael Büsch <m@bues.ch>
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License along
// with this program; if not, write to the Free Software Foundation, Inc.,
// 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
//

use anyhow::{self as ah, Context as _};
use rand::{thread_rng, Rng};
use std::collections::{HashMap, VecDeque};
use std::fs::File;
use std::time::{Duration, Instant};
use uhid_virt::{Bus, CreateParams, OutputEvent, StreamError, UHIDDevice};

use crate::ctap::{CtapCmd, CtapResp};

const PACK_SIZE: usize = 64;
const HDR_INIT_SIZE: usize = 7;
const HDR_CONT_SIZE: usize = 5;
pub const MAX_MSG_PAYLOAD_SIZE: usize =
    (PACK_SIZE - HDR_INIT_SIZE) + (128 * (PACK_SIZE - HDR_CONT_SIZE));

const CID_BROADCAST: u32 = 0xFFFFFFFF;
const CID_RESERVED: u32 = 0;

const CTAPHID_PING: u8 = 0x01;
const CTAPHID_MSG: u8 = 0x03;
const CTAPHID_LOCK: u8 = 0x04;
const CTAPHID_INIT: u8 = 0x06;
const CTAPHID_WINK: u8 = 0x08;
const CTAPHID_CBOR: u8 = 0x10;
const CTAPHID_CANCEL: u8 = 0x11;
const CTAPHID_ERROR: u8 = 0x3F;
const CTAPHID_KEEPALIVE: u8 = 0x3B;

const CAPABILITY_WINK: u8 = 0x01;
const CAPABILITY_CBOR: u8 = 0x04;
const CAPABILITY_NMSG: u8 = 0x08;

const TIMEOUT_MS: u64 = 10000;

fn u16_to_bytes(b: &mut [u8], v: u16) {
    b[0..2].copy_from_slice(&v.to_be_bytes());
}

fn u32_to_bytes(b: &mut [u8], v: u32) {
    b[0..4].copy_from_slice(&v.to_be_bytes());
}

fn u64_to_bytes(b: &mut [u8], v: u64) {
    b[0..8].copy_from_slice(&v.to_be_bytes());
}

fn bytes_to_u16(b: &[u8]) -> u16 {
    u16::from_be_bytes([b[0], b[1]])
}

fn bytes_to_u32(b: &[u8]) -> u32 {
    u32::from_be_bytes([b[0], b[1], b[2], b[3]])
}

fn bytes_to_u64(b: &[u8]) -> u64 {
    u64::from_be_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]])
}

struct Channel {
    cid: u32,
    timeout: Instant,
    rx_transaction: bool,
    rx_cmd: u8,
    rx_payload: [u8; MAX_MSG_PAYLOAD_SIZE],
    rx_msg_bcnt: usize,
    rx_count: usize,
    rx_seq_next: u8,
    rx_commands: VecDeque<CtapCmd>,
}

impl Channel {
    fn new(cid: u32) -> Self {
        Channel {
            cid,
            timeout: Instant::now(),
            rx_transaction: false,
            rx_cmd: 0,
            rx_payload: [0; MAX_MSG_PAYLOAD_SIZE],
            rx_msg_bcnt: 0,
            rx_count: 0,
            rx_seq_next: 0xFF,
            rx_commands: VecDeque::new(),
        }
    }

    fn kill_rx_transaction(&mut self) {
        self.rx_transaction = false;
        self.rx_msg_bcnt = 0;
        self.rx_count = 0;
    }
}

/// Emulated FIDO2 HID device.
pub struct CtapHid {
    dev: UHIDDevice<File>,
    chans: HashMap<u32, Channel>,
    next_cid: u32,
}

impl CtapHid {
    pub fn new() -> ah::Result<Self> {
        #[rustfmt::skip]
        let hid_report_descriptor = vec![
            0x06, 0xd0, 0xf1,       // USAGE_PAGE (FIDO Alliance Page 0xF1D0)
            0x09, 0x01,             // USAGE (U2F Authenticator Device)
            0xa1, 0x01,             // COLLECTION (Application)
            0x09, 0x20,             // USAGE (Input Report Data)
            0x15, 0x00,             // LOGICAL_MINIMUM (0)
            0x26, 0xff, 0x00,       // LOGICAL_MAXIMUM (255)
            0x75, 0x08,             // REPORT_SIZE (8)
            0x95, PACK_SIZE as u8,  // REPORT_COUNT (64)
            0x81, 0x02,             // INPUT (Data, Var, Abs)
            0x09, 0x21,             // USAGE (Output Report Data)
            0x15, 0x00,             // LOGICAL_MINIMUM (0)
            0x26, 0xff, 0x00,       // LOGICAL_MAXIMUM (255)
            0x75, 0x08,             // REPORT_SIZE (8)
            0x95, PACK_SIZE as u8,  // REPORT_COUNT (64)
            0x91, 0x02,             // OUTPUT (Data, Var, Abs)
            0xc0,                   // END_COLLECTION
        ];
        let serial_number: u64 = thread_rng().gen();
        let dev = UHIDDevice::create(CreateParams {
            name: "pwman-f2e".to_string(),
            phys: "<emulated device>".to_string(),
            uniq: serial_number.to_string(),
            bus: Bus::USB,
            vendor: 0x6666,
            product: 0xF1D0,
            version: 0x0111,
            country: 0,
            rd_data: hid_report_descriptor,
        })
        .context("uHID-device open failed.")?;
        let mut chans = HashMap::new();
        chans.insert(CID_BROADCAST, Channel::new(CID_BROADCAST));
        Ok(CtapHid {
            dev,
            chans,
            next_cid: 1,
        })
    }

    fn handle_timeouts(&mut self, now: Instant) {
        // Close channels that timed out.
        self.chans.retain(|_, chan| {
            if chan.cid == CID_BROADCAST {
                true // Broadcast never times out.
            } else {
                now < chan.timeout
            }
        });

        // Broadcast RX transaction timeout.
        if let Some(bcast_chan) = self.chans.get_mut(&CID_BROADCAST) {
            if now >= bcast_chan.timeout {
                bcast_chan.kill_rx_transaction();
            }
        }
    }

    fn alloc_cid(&mut self) -> u32 {
        while self.next_cid == CID_RESERVED
            || self.next_cid == CID_BROADCAST
            || self.chans.contains_key(&self.next_cid)
        {
            self.next_cid = self.next_cid.wrapping_add(1);
        }
        let cid = self.next_cid;
        self.next_cid = self.next_cid.wrapping_add(1);
        cid
    }

    fn send(&mut self, cid: u32, cmd: u8, payload: &[u8]) -> ah::Result<()> {
        assert!(payload.len() <= MAX_MSG_PAYLOAD_SIZE);
        let mut pack = [0_u8; PACK_SIZE];
        let mut seq = 0xFF;
        let mut i = 0;
        while i < payload.len() {
            let hdrlen = if seq == 0xFF {
                HDR_INIT_SIZE
            } else {
                HDR_CONT_SIZE
            };
            let plen = (payload.len() - i).min(PACK_SIZE - hdrlen);
            if seq == 0xFF {
                // initialization packet
                u32_to_bytes(&mut pack[0..4], cid); // CID
                pack[4] = cmd | 0x80; // CMD
                u16_to_bytes(&mut pack[5..7], payload.len() as u16); // BCNT
                pack[7..7 + plen].copy_from_slice(&payload[i..i + plen]); // DATA
                seq = 0;
            } else {
                // continuation packet
                assert!(seq <= 0x7F);
                u32_to_bytes(&mut pack[0..4], cid); // CID
                pack[4] = seq; // SEQ
                pack[5..5 + plen].copy_from_slice(&payload[i..i + plen]); // DATA
                seq += 1;
            }
            self.dev.write(&pack).context("Failed to send packet.")?;
            i += plen;
        }
        Ok(())
    }

    fn handle_broadcast(&mut self) -> ah::Result<()> {
        let chan = self.chans.get_mut(&CID_BROADCAST).unwrap();
        let cmd = chan.rx_cmd;
        let payload = &chan.rx_payload;
        match cmd {
            CTAPHID_PING => {
                println!("RX: CTAP broadcast PING");
                let reply = *payload;
                self.send(CID_BROADCAST, CTAPHID_PING, &reply)
                    .context("Failed to send CTAPHID_PING reply")?;
            }
            CTAPHID_MSG => {
                println!("RX: CTAP broadcast unexpected MSG");
            }
            CTAPHID_LOCK => {
                println!("RX: CTAP broadcast LOCK");
                //TODO
            }
            CTAPHID_INIT => {
                println!("RX: CTAP broadcast INIT");
                if payload.len() < 8 {
                    return Err(ah::format_err!("INIT: Payload too short."));
                }
                let nonce = bytes_to_u64(&payload[0..8]);

                let cid = self.alloc_cid();
                let capabilities = CAPABILITY_WINK | CAPABILITY_CBOR | CAPABILITY_NMSG;

                let mut reply = [0_u8; 17];
                u64_to_bytes(&mut reply[0..8], nonce);
                u32_to_bytes(&mut reply[8..12], cid);
                reply[12] = 2; // protocol version identifier
                reply[13] = 1; // major device version number
                reply[14] = 0; // minor device version number
                reply[15] = 0; // build device version number
                reply[16] = capabilities;
                self.send(CID_BROADCAST, CTAPHID_INIT, &reply)
                    .context("Failed to send CTAPHID_INIT reply")?;
            }
            CTAPHID_WINK => {
                println!("RX: CTAP broadcast WINK");
                //TODO
            }
            CTAPHID_CBOR => {
                println!("RX: CTAP broadcast unexpected CBOR");
            }
            CTAPHID_CANCEL => {
                println!("RX: CTAP broadcast unexpected CANCEL");
            }
            CTAPHID_ERROR => {
                println!("RX: CTAP broadcast unexpected ERROR");
            }
            CTAPHID_KEEPALIVE => {
                println!("RX: CTAP broadcast KEEPALIVE");
                //TODO
            }
            _ => {
                return Err(ah::format_err!(
                    "RX: CTAP broadcast unknown command: {cmd}."
                ));
            }
        }
        Ok(())
    }

    fn handle_cid_message(&mut self, cid: u32) -> ah::Result<()> {
        let chan = self.chans.get_mut(&cid).unwrap();
        let cid = chan.cid;
        let cmd = chan.rx_cmd;
        let payload = &chan.rx_payload;
        match cmd {
            CTAPHID_PING => {
                println!("RX: CTAP channel {cid} PING");
                let reply = *payload;
                self.send(cid, CTAPHID_PING, &reply)
                    .context("Failed to send CTAPHID_PING reply")?;
            }
            CTAPHID_MSG => {
                println!("RX: CTAP channel {cid} unexpected MSG");
            }
            CTAPHID_LOCK => {
                println!("RX: CTAP channel {cid} LOCK");
                //TODO
            }
            CTAPHID_INIT => {
                println!("RX: CTAP channel {cid} unexpected INIT");
            }
            CTAPHID_WINK => {
                println!("RX: CTAP channel {cid} WINK");
                //TODO
            }
            CTAPHID_CBOR => {
                println!("RX: CTAP channel {cid} CBOR");
                if payload.is_empty() {
                    return Err(ah::format_err!("CBOR: Payload too short."));
                }
                let msg_cmd = payload[0];
                let msg = CtapCmd::parse_cbor(msg_cmd, &payload[1..])
                    .context("Failed to parse CTAP CBOR command.")?;
                chan.rx_commands.push_back(msg);
            }
            CTAPHID_CANCEL => {
                println!("RX: CTAP channel {cid} CANCEL");
                chan.kill_rx_transaction();
            }
            CTAPHID_ERROR => {
                println!("RX: CTAP channel {cid} unexpected ERROR");
            }
            CTAPHID_KEEPALIVE => {
                println!("RX: CTAP channel {cid} KEEPALIVE");
                //TODO
            }
            _ => {
                return Err(ah::format_err!(
                    "RX: CTAP channel {cid} unknown command: {cmd}."
                ));
            }
        }
        Ok(())
    }

    fn handle_output_data(&mut self, data: &[u8]) -> ah::Result<()> {
        let now = Instant::now();
        self.handle_timeouts(now);

        if data.len() < HDR_CONT_SIZE.min(HDR_INIT_SIZE) {
            return Err(ah::format_err!("Invalid header size."));
        }
        let cid = bytes_to_u32(&data[0..4]);
        if cid == CID_RESERVED {
            return Err(ah::format_err!("Invalid CID (reserved)."));
        }
        let cmd = data[4];

        if cmd & 0x80 != 0 {
            // initialization packet
            let cmd = cmd & 0x7F;
            if data.len() < HDR_INIT_SIZE {
                return Err(ah::format_err!("Invalid header (init)."));
            }
            let bcnt = bytes_to_u16(&data[5..7]) as usize;
            if bcnt > MAX_MSG_PAYLOAD_SIZE {
                return Err(ah::format_err!("Message payload size (BCNT) is too large."));
            }

            let chan = self.chans.entry(cid).or_insert_with(|| Channel::new(cid));
            chan.kill_rx_transaction();

            let chunklen = bcnt.min(PACK_SIZE - HDR_INIT_SIZE);
            chan.rx_payload[0..chunklen]
                .copy_from_slice(&data[HDR_INIT_SIZE..HDR_INIT_SIZE + chunklen]);
            chan.rx_count = chunklen;
            chan.rx_msg_bcnt = bcnt;
            chan.rx_cmd = cmd;
            chan.rx_seq_next = 0;
            chan.rx_transaction = true;
            chan.timeout = now + Duration::from_millis(TIMEOUT_MS);
        } else {
            // continuation packet
            let seq = cmd;
            if let Some(chan) = self.chans.get_mut(&cid) {
                if chan.rx_transaction && seq == chan.rx_seq_next {
                    let chunklen =
                        (chan.rx_msg_bcnt - chan.rx_count).min(PACK_SIZE - HDR_CONT_SIZE);
                    chan.rx_payload[chan.rx_count..chan.rx_count + chunklen]
                        .copy_from_slice(&data[HDR_CONT_SIZE..HDR_CONT_SIZE + chunklen]);
                    chan.rx_count += chunklen;
                    chan.rx_seq_next = (chan.rx_seq_next + 1) & 0x7F;
                }
            }
        }

        // Get the rx transactions that are ready for handling.
        let rx_trans_ready_cids: Vec<u32> = self
            .chans
            .values()
            .filter_map(|chan| {
                if chan.rx_transaction && chan.rx_count >= chan.rx_msg_bcnt {
                    Some(chan.cid)
                } else {
                    None
                }
            })
            .collect();

        // Handle all rx transactions.
        for cid in &rx_trans_ready_cids {
            if *cid == CID_BROADCAST {
                self.handle_broadcast()?;
            } else {
                self.handle_cid_message(*cid)?;
            }
            let chan = self.chans.get_mut(cid).unwrap();
            // Transaction done.
            chan.kill_rx_transaction();
        }
        Ok(())
    }

    fn handle_output_event(&mut self, ev: OutputEvent) -> ah::Result<()> {
        match ev {
            OutputEvent::Start { dev_flags: _ } => {
                println!("OutputEvent::Start()");
            }
            OutputEvent::Stop => {
                println!("OutputEvent::Stop");
            }
            OutputEvent::Open => {
                println!("OutputEvent::Open");
            }
            OutputEvent::Close => {
                println!("OutputEvent::Close");
            }
            OutputEvent::Output { data } => {
                println!("OutputEvent::Output({data:?})");
                if data.is_empty() {
                    return Err(ah::format_err!("Zero OutputEvent::Output data length."));
                }
                self.handle_output_data(&data[1..])
                    .context("Error parsing HID output data.")?;
            }
            OutputEvent::GetReport {
                id: _,
                report_number: _,
                report_type: _,
            } => {
                println!("OutputEvent::GetReport");
            }
            OutputEvent::SetReport {
                id: _,
                report_number: _,
                report_type: _,
                data: _,
            } => {
                println!("OutputEvent::SetReport");
            }
        }
        Ok(())
    }

    pub fn read(&mut self) -> ah::Result<()> {
        match self.dev.read() {
            Ok(ev) => {
                self.handle_output_event(ev)?;
            }
            Err(StreamError::Io(e)) => {
                return Err(ah::format_err!("HID I/O Error: {e}"));
            }
            Err(StreamError::UnknownEventType(t)) => {
                return Err(ah::format_err!("Unknown HID event type: {t}"));
            }
        }
        Ok(())
    }

    pub fn get_commands(&mut self) -> Vec<(u32, CtapCmd)> {
        let mut ret = vec![];
        for (cid, chan) in &mut self.chans {
            while let Some(cmd) = chan.rx_commands.pop_front() {
                ret.push((*cid, cmd));
            }
        }
        ret
    }

    pub fn send_response(&mut self, cid: u32, mut resp: CtapResp) -> ah::Result<()> {
        let mut data = Vec::with_capacity(resp.cbor.len() + 1);
        data.push(resp.status);
        data.append(&mut resp.cbor);
        self.send(cid, CTAPHID_CBOR, &data)
            .context("Failed to send response.")?;
        Ok(())
    }
}

// vim: ts=4 sw=4 expandtab