use core::sync::atomic::{AtomicBool, Ordering};

use alloc::{boxed::Box, collections::btree_map::BTreeMap, sync::Arc};
use async_trait::async_trait;
use libk::{
    block,
    error::Error,
    fs::devfs,
    vfs::{Terminal, TerminalInput, TerminalOutput},
};
use libk_util::sync::{IrqSafeSpinlock, spin_rwlock::IrqSafeRwLock};
use yggdrasil_abi::{
    io::{FileMode, TerminalOptions, TerminalOutputOptions},
    process::ProcessId,
};

use crate::{
    class_driver::UsbDeviceDriver,
    communication::UsbDirection,
    device::{UsbDeviceAccess, UsbDeviceDetachHandler},
    error::UsbError,
    info::UsbEndpointType,
    pipe::{control::ControlTransferSetup, normal::UsbBulkOutPipeAccess},
};

pub struct FT232Driver;
struct FT232Device {
    device: Arc<UsbDeviceAccess>,
    bulk_out: UsbBulkOutPipeAccess,
    baud_rate: IrqSafeRwLock<u32>,
}

pub struct UsbSerialDeviceWrapper {
    device: Arc<dyn UsbSerialDevice>,
    index: u32,
    disconnected: AtomicBool,
}

#[async_trait]
pub trait UsbSerialDevice: Send + Sync + 'static {
    async fn setup(&self) -> Result<(), UsbError>;

    async fn set_baud_rate(&self, baud: u32) -> Result<(), UsbError>;
    fn baud_rate(&self) -> u32;

    async fn write(&self, buffer: &[u8], options: TerminalOutputOptions)
    -> Result<usize, UsbError>;

    fn display_name(&self) -> &str;
    fn device(&self) -> &Arc<UsbDeviceAccess>;
}

#[async_trait]
impl UsbDeviceDetachHandler for UsbSerialDeviceWrapper {
    async fn handle_device_detach(&self) {
        self.disconnected.store(true, Ordering::Release);
        log::info!("USB serial #{} disconnected", self.index);
        remove_usb_serial(self.index);
    }
}

impl TerminalOutput for UsbSerialDeviceWrapper {
    fn write(&self, byte: u8, options: &TerminalOutputOptions) -> Result<(), Error> {
        self.write_multiple(&[byte], options)?;
        Ok(())
    }

    fn write_multiple(
        &self,
        bytes: &[u8],
        options: &TerminalOutputOptions,
    ) -> Result<usize, Error> {
        if self.disconnected.load(Ordering::Acquire) {
            return Err(Error::InvalidOperation);
        }
        let result = block!(self.device.write(bytes, *options).await)?;
        match result {
            Ok(len) => Ok(len),
            Err(UsbError::MemoryError(err) | UsbError::SystemError(err)) => Err(err),
            Err(error) => {
                log::warn!("{}: write error: {:?}", self.device.display_name(), error);
                Err(Error::InvalidOperation)
            }
        }
    }

    fn baud_rate(&self) -> u32 {
        self.device.baud_rate()
    }

    fn set_baud_rate(&self, baud: u32) -> Result<(), Error> {
        if self.disconnected.load(Ordering::Acquire) {
            return Err(Error::InvalidOperation);
        }
        let result = block!(self.device.set_baud_rate(baud).await)?;
        match result {
            Ok(()) => Ok(()),
            Err(UsbError::MemoryError(err) | UsbError::SystemError(err)) => Err(err),
            Err(error) => {
                log::warn!(
                    "{}: baud rate set error: {:?}",
                    self.device.display_name(),
                    error
                );
                Err(Error::InvalidOperation)
            }
        }
    }

    fn open(&self, pid: ProcessId) -> Result<(), Error> {
        if self.disconnected.load(Ordering::Acquire) {
            return Err(Error::InvalidOperation);
        }
        let _ = pid;
        let result = block! { self.device.setup().await }?;
        match result {
            Ok(()) => Ok(()),
            Err(UsbError::MemoryError(err) | UsbError::SystemError(err)) => Err(err),
            Err(error) => {
                log::warn!("{}: setup error: {:?}", self.device.display_name(), error);
                Err(Error::InvalidOperation)
            }
        }
    }
}

impl FT232Device {
    const FTDI_RESET: u8 = 0;
    const FTDI_SET_BAUD_RATE: u8 = 3;

    async fn ftdi_reset(&self, port: u16) -> Result<(), UsbError> {
        self.device
            .control_pipe()
            .control_transfer(ControlTransferSetup {
                bm_request_type: 0b01000000,
                b_request: Self::FTDI_RESET,
                w_value: 0,
                w_index: port,
                w_length: 0,
            })
            .await
    }

    async fn ftdi_set_baud_rate(&self, port: u16, baud: u32) -> Result<(), UsbError> {
        let w_value = match baud {
            300 => 0x2710,
            600 => 0x1388,
            1200 => 0x09C4,
            2400 => 0x04E2,
            4800 => 0x0271,
            9600 => 0x4138,
            19200 => 0x809C,
            38400 => 0xC04E,
            57600 => 0x0034,
            115200 => 0x001A,
            230400 => 0x000D,
            460800 => 0x4006,
            921600 => 0x8003,
            _ => {
                log::warn!("ft232: unsupported baud rate {baud}");
                return Err(UsbError::InvalidConfiguration);
            }
        };

        self.device
            .control_pipe()
            .control_transfer(ControlTransferSetup {
                bm_request_type: 0b01000000,
                b_request: Self::FTDI_SET_BAUD_RATE,
                w_value,
                w_index: port,
                w_length: 0,
            })
            .await
    }
}

#[async_trait]
impl UsbSerialDevice for FT232Device {
    async fn setup(&self) -> Result<(), UsbError> {
        log::info!("ft232: setup");
        Ok(())
    }

    async fn set_baud_rate(&self, baud: u32) -> Result<(), UsbError> {
        self.ftdi_set_baud_rate(0, baud).await?;
        // *self.baud_rate.write() = baud;
        log::info!("ft232: set baud rate {baud}");
        Ok(())
    }
    fn baud_rate(&self) -> u32 {
        *self.baud_rate.read()
    }

    async fn write(
        &self,
        buffer: &[u8],
        options: TerminalOutputOptions,
    ) -> Result<usize, UsbError> {
        if options.contains(TerminalOutputOptions::NL_TO_CRNL) {
            for &byte in buffer {
                if byte == b'\n' {
                    self.bulk_out.write(b"\r").await?;
                }
                self.bulk_out.write(&[byte]).await?;
            }
            Ok(buffer.len())
        } else {
            self.bulk_out.write(buffer).await
        }
    }

    fn display_name(&self) -> &str {
        "FT232 Serial Converter"
    }
    fn device(&self) -> &Arc<UsbDeviceAccess> {
        &self.device
    }
}

#[async_trait]
impl UsbDeviceDriver for FT232Driver {
    async fn run(self: Arc<Self>, device: Arc<UsbDeviceAccess>) -> Result<(), UsbError> {
        let interface = device.interface(0);
        let endpoints = interface.endpoints();

        let bulk_in = endpoints
            .iter()
            .find(|ep| ep.ty == UsbEndpointType::Bulk && ep.direction == UsbDirection::In)
            .ok_or(UsbError::InvalidConfiguration)?;
        let bulk_out = endpoints
            .iter()
            .find(|ep| ep.ty == UsbEndpointType::Bulk && ep.direction == UsbDirection::Out)
            .ok_or(UsbError::InvalidConfiguration)?;

        let bulk_in = device
            .open_bulk_in_pipe(bulk_in.number, bulk_in.max_packet_size as _)
            .await?;
        let bulk_out = device
            .open_bulk_out_pipe(bulk_out.number, bulk_out.max_packet_size as _)
            .await?;

        let ft232 = Arc::new(FT232Device {
            device,
            bulk_out,
            baud_rate: IrqSafeRwLock::new(115200),
        });
        ft232.ftdi_reset(0).await?;
        let terminal = register_usb_serial(ft232)?;

        let mut buffer = [0; 3];
        loop {
            let len = bulk_in.read(&mut buffer).await?;
            if len < 2 {
                continue;
            }
            for &byte in &buffer[2..len] {
                terminal.write_to_input(byte);
            }
        }
    }

    fn name(&self) -> &'static str {
        "FT232 Serial Converter"
    }

    fn probe(&self, class: u8, subclass: u8, protocol: u8, vid: u16, pid: u16) -> bool {
        let _ = (class, subclass, protocol);
        vid == 0x0403 && pid == 0x6001
    }
}

static USB_SERIALS: IrqSafeSpinlock<BTreeMap<u32, Arc<Terminal<Arc<UsbSerialDeviceWrapper>>>>> =
    IrqSafeSpinlock::new(BTreeMap::new());

fn register_usb_serial(
    serial: Arc<dyn UsbSerialDevice>,
) -> Result<Arc<Terminal<Arc<UsbSerialDeviceWrapper>>>, UsbError> {
    let mut serials = USB_SERIALS.lock();
    for i in 0..64 {
        if serials.contains_key(&i) {
            continue;
        }
        let wrapper = Arc::new(UsbSerialDeviceWrapper {
            device: serial,
            index: i,
            disconnected: AtomicBool::new(false),
        });
        wrapper.device.device().set_detach_handler(wrapper.clone());
        let input = TerminalInput::with_capacity(64).expect("Couldn't allocate input buffer");
        let terminal = Arc::new(Terminal::from_parts(
            TerminalOptions::const_default(),
            input,
            wrapper,
        ));
        serials.insert(i, terminal.clone());
        let name = alloc::format!("ttyUSB{i}");
        devfs::add_named_char_device(terminal.clone(), name.clone(), FileMode::new(0o600)).ok();
        return Ok(terminal);
    }
    Err(UsbError::DriverError)
}

fn remove_usb_serial(index: u32) {
    let serial = USB_SERIALS.lock().remove(&index);
    if serial.is_none() {
        log::warn!("usb-serial #{index} doesn't exist in the table");
    }
    let name = alloc::format!("ttyUSB{index}");
    devfs::remove_node(name).ok();
}
