//! PS/2 controller, see:
//! - https://wiki.osdev.org/I8042_PS/2_Controller
//! - http://www.mcamafia.de/pdf/ibm_hitrc07.pdf

use common::{
    io::{Io, ReadOnly, WriteOnly},
    timeout::Timeout,
};

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
use common::io::Pio;

#[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
use common::io::Mmio;

use log::{debug, error, info, trace, warn};

use std::fmt;

#[derive(Debug)]
pub enum Error {
    CommandRetry,
    NoMoreTries,
    ReadTimeout,
    WriteTimeout,
    CommandTimeout(Command),
    WriteConfigTimeout(ConfigFlags),
    KeyboardCommandFail(KeyboardCommand),
    KeyboardCommandDataFail(KeyboardCommandData),
}

bitflags! {
    pub struct StatusFlags: u8 {
        const OUTPUT_FULL = 1;
        const INPUT_FULL = 1 << 1;
        const SYSTEM = 1 << 2;
        const COMMAND = 1 << 3;
        // Chipset specific
        const KEYBOARD_LOCK = 1 << 4;
        // Chipset specific
        const SECOND_OUTPUT_FULL = 1 << 5;
        const TIME_OUT = 1 << 6;
        const PARITY = 1 << 7;
    }
}

bitflags! {
    #[derive(Clone, Copy, Debug)]
    pub struct ConfigFlags: u8 {
        const FIRST_INTERRUPT = 1 << 0;
        const SECOND_INTERRUPT = 1 << 1;
        const POST_PASSED = 1 << 2;
        // 1 << 3 should be zero
        const CONFIG_RESERVED_3 = 1 << 3;
        const FIRST_DISABLED = 1 << 4;
        const SECOND_DISABLED = 1 << 5;
        const FIRST_TRANSLATE = 1 << 6;
        // 1 << 7 should be zero
        const CONFIG_RESERVED_7 = 1 << 7;
    }
}

#[derive(Clone, Copy, Debug)]
#[repr(u8)]
#[allow(dead_code)]
enum Command {
    ReadConfig = 0x20,
    WriteConfig = 0x60,
    DisableSecond = 0xA7,
    EnableSecond = 0xA8,
    TestSecond = 0xA9,
    TestController = 0xAA,
    TestFirst = 0xAB,
    Diagnostic = 0xAC,
    DisableFirst = 0xAD,
    EnableFirst = 0xAE,
    WriteSecond = 0xD4,
}

#[derive(Clone, Copy, Debug)]
#[repr(u8)]
#[allow(dead_code)]
enum KeyboardCommand {
    EnableReporting = 0xF4,
    SetDefaultsDisable = 0xF5,
    SetDefaults = 0xF6,
    Reset = 0xFF,
}

#[derive(Clone, Copy, Debug)]
#[repr(u8)]
enum KeyboardCommandData {
    ScancodeSet = 0xF0,
}

// Default timeout in microseconds
const DEFAULT_TIMEOUT: u64 = 50_000;
// Reset timeout in microseconds
const RESET_TIMEOUT: u64 = 1_000_000;
// Maximum bytes to drain during flush (Linux: I8042_BUFFER_SIZE)
const FLUSH_LIMIT: usize = 4096;
// Controller self-test pass value (Linux: I8042_RET_CTL_TEST)
const SELFTEST_PASS: u8 = 0x55;
// Controller self-test retries (Linux: 5 attempts)
const SELFTEST_RETRIES: usize = 5;
// AUX port test pass value (Linux returns 0x00 on success)
const AUX_TEST_PASS: u8 = 0x00;

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub struct Ps2 {
    data: Pio<u8>,
    status: ReadOnly<Pio<u8>>,
    command: WriteOnly<Pio<u8>>,
    //TODO: keep in state instead
    pub mouse_resets: usize,
}

#[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
pub struct Ps2 {
    data: Mmio<u8>,
    status: ReadOnly<Mmio<u8>>,
    command: WriteOnly<Mmio<u8>>,
    //TODO: keep in state instead
    pub mouse_resets: usize,
}

impl Ps2 {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    pub fn new() -> Self {
        Ps2 {
            data: Pio::new(0x60),
            status: ReadOnly::new(Pio::new(0x64)),
            command: WriteOnly::new(Pio::new(0x64)),
            mouse_resets: 0,
        }
    }

    #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
    pub fn new() -> Self {
        // The PS/2 controller is an x86/x86_64 legacy device. No other
        // architecture exposes the 0x60/0x64 I/O ports.
        panic!("ps2d: only supported on x86 and x86_64");
    }

    fn status(&mut self) -> StatusFlags {
        StatusFlags::from_bits_truncate(self.status.read())
    }

    fn wait_read(&mut self, micros: u64) -> Result<(), Error> {
        let timeout = Timeout::from_micros(micros);
        loop {
            if self.status().contains(StatusFlags::OUTPUT_FULL) {
                return Ok(());
            }
            timeout.run().map_err(|()| Error::ReadTimeout)?
        }
    }

    fn wait_write(&mut self, micros: u64) -> Result<(), Error> {
        let timeout = Timeout::from_micros(micros);
        loop {
            if !self.status().contains(StatusFlags::INPUT_FULL) {
                return Ok(());
            }
            timeout.run().map_err(|()| Error::WriteTimeout)?
        }
    }

    fn command(&mut self, command: Command) -> Result<(), Error> {
        self.wait_write(DEFAULT_TIMEOUT)
            .map_err(|_| Error::CommandTimeout(command))?;
        self.command.write(command as u8);
        Ok(())
    }

    fn read(&mut self) -> Result<u8, Error> {
        self.read_timeout(DEFAULT_TIMEOUT)
    }

    fn read_timeout(&mut self, micros: u64) -> Result<u8, Error> {
        self.wait_read(micros)?;
        let data = self.data.read();
        Ok(data)
    }

    fn write(&mut self, data: u8) -> Result<(), Error> {
        self.wait_write(DEFAULT_TIMEOUT)?;
        self.data.write(data);
        Ok(())
    }

    fn retry<T, F: Fn(&mut Self) -> Result<T, Error>>(
        &mut self,
        name: fmt::Arguments,
        retries: usize,
        f: F,
    ) -> Result<T, Error> {
        trace!("{}", name);
        let mut res = Err(Error::NoMoreTries);
        for retry in 0..retries {
            res = f(self);
            match res {
                Ok(ok) => {
                    return Ok(ok);
                }
                Err(ref err) => {
                    debug!("{}: retry {}/{}: {:?}", name, retry + 1, retries, err);
                }
            }
        }
        res
    }

    fn config(&mut self) -> Result<ConfigFlags, Error> {
        self.retry(format_args!("read config"), 4, |x| {
            x.command(Command::ReadConfig)?;
            x.read()
        })
        .map(ConfigFlags::from_bits_truncate)
    }

    fn set_config(&mut self, config: ConfigFlags) -> Result<(), Error> {
        self.retry(format_args!("write config {:?}", config), 4, |x| {
            x.command(Command::WriteConfig)?;
            x.write(config.bits())
                .map_err(|_| Error::WriteConfigTimeout(config))?;
            Ok(0)
        })?;
        Ok(())
    }

    fn keyboard_command_inner(&mut self, command: u8) -> Result<u8, Error> {
        self.write(command)?;
        match self.read()? {
            0xFE => Err(Error::CommandRetry),
            value => Ok(value),
        }
    }

    fn keyboard_command(&mut self, command: KeyboardCommand) -> Result<u8, Error> {
        self.retry(format_args!("keyboard command {:?}", command), 4, |x| {
            x.keyboard_command_inner(command as u8)
                .map_err(|_| Error::KeyboardCommandFail(command))
        })
    }

    fn keyboard_command_data(
        &mut self,
        command: KeyboardCommandData,
        data: u8,
    ) -> Result<u8, Error> {
        self.retry(
            format_args!("keyboard command {:?} {:#x}", command, data),
            4,
            |x| {
                let res = x
                    .keyboard_command_inner(command as u8)
                    .map_err(|_| Error::KeyboardCommandDataFail(command))?;
                if res != 0xFA {
                    warn!("keyboard incorrect result of set command: {command:?} {res:02X}");
                    return Ok(res);
                }
                x.write(data)?;
                x.read()
            },
        )
    }

    pub fn mouse_command_async(&mut self, command: u8) -> Result<(), Error> {
        self.command(Command::WriteSecond)?;
        self.write(command as u8)
    }

    pub fn set_leds(&mut self, caps: bool, num: bool, scroll: bool) {
        let mut led_byte = 0u8;
        if scroll { led_byte |= 1; }
        if num { led_byte |= 2; }
        if caps { led_byte |= 4; }
        if let Err(err) = self.keyboard_command_inner(0xED) {
            log::debug!("ps2d: LED command 0xED not supported: {:?}", err);
            return;
        }
        match self.read_timeout(DEFAULT_TIMEOUT) {
            Ok(0xFA) => {
                if let Err(err) = self.write(led_byte) {
                    log::debug!("ps2d: failed to send LED byte {:02X}: {:?}", led_byte, err);
                }
            }
            Ok(val) => {
                log::debug!("ps2d: LED command ACK expected 0xFA, got {:02X}", val);
            }
            Err(err) => {
                log::debug!("ps2d: LED command ACK timeout: {:?}", err);
            }
        }
    }

    pub fn next(&mut self) -> Option<(bool, u8)> {
        let status = self.status();
        if status.contains(StatusFlags::OUTPUT_FULL) {
            let data = self.data.read();
            Some((!status.contains(StatusFlags::SECOND_OUTPUT_FULL), data))
        } else {
            None
        }
    }

    /// Drain all pending bytes from the controller output buffer.
    /// Borrowed from Linux i8042_flush(): stale firmware/BIOS bytes can be
    /// misinterpreted as device responses during initialization.
    fn flush(&mut self) -> usize {
        let mut count = 0;
        while self.status().contains(StatusFlags::OUTPUT_FULL) {
            if count >= FLUSH_LIMIT {
                warn!("flush: exceeded limit, controller may be stuck");
                break;
            }
            let data = self.data.read();
            trace!("flush: discarded {:02X}", data);
            count += 1;
        }
        if count > 0 {
            debug!("flushed {} stale bytes from controller", count);
        }
        count
    }

    /// Test the AUX (mouse) port via controller command 0xA9.
    /// Borrowed from Linux: verifies electrical connectivity before
    /// attempting to talk to the mouse. Returns true if the port passed.
    fn test_aux_port(&mut self) -> bool {
        if let Err(err) = self.command(Command::TestSecond) {
            warn!("aux port test command failed: {:?}", err);
            return false;
        }
        match self.read() {
            Ok(AUX_TEST_PASS) => {
                debug!("aux port test passed");
                true
            }
            Ok(val) => {
                warn!("aux port test failed: {:02X}", val);
                false
            }
            Err(err) => {
                warn!("aux port test read timeout: {:?}", err);
                false
            }
        }
    }

    pub fn init_keyboard(&mut self) -> Result<(), Error> {
        let mut b;

        {
            // Enable first device
            self.command(Command::EnableFirst)?;
        }

        {
            // Reset keyboard
            b = self.keyboard_command(KeyboardCommand::Reset)?;
            if b == 0xFA {
                b = self.read().unwrap_or(0);
                if b != 0xAA {
                    error!("keyboard failed self test: {:02X}", b);
                }
            } else {
                error!("keyboard failed to reset: {:02X}", b);
            }
        }

        {
            // Set scancode set to 2
            let scancode_set = 2;
            b = self.keyboard_command_data(KeyboardCommandData::ScancodeSet, scancode_set)?;
            if b != 0xFA {
                error!(
                    "keyboard failed to set scancode set {}: {:02X}",
                    scancode_set, b
                );
            }
        }

        Ok(())
    }

    pub fn init(&mut self) -> Result<(), Error> {
        // Linux i8042_controller_check(): verify controller is present by
        // flushing any stale data. A stuck output buffer means no controller.
        self.flush();

        // Bare-metal controllers may be slow after firmware handoff.
        // Give the controller a moment to finish POST before sending commands.
        std::thread::sleep(std::time::Duration::from_millis(50));

        {
            // Disable both ports first — use retry because the controller
            // may still be settling or temporarily unresponsive.
            // Failure here is non-fatal: we continue and attempt the rest
            // of initialization.  A truly absent controller will fail later
            // at self-test or keyboard reset.
            if let Err(err) = self.retry(
                format_args!("disable first port"),
                3,
                |x| x.command(Command::DisableFirst),
            ) {
                warn!("disable first port failed: {:?}", err);
            }
            if let Err(err) = self.retry(
                format_args!("disable second port"),
                3,
                |x| x.command(Command::DisableSecond),
            ) {
                warn!("disable second port failed: {:?}", err);
            }
        }

        // Flush again after disabling — firmware may have queued more bytes
        self.flush();

        // Linux i8042_controller_init() step 1: write a known-safe config
        // (interrupts off, both ports disabled) so stale config can't cause
        // spurious interrupts during the rest of init.
        {
            let config = ConfigFlags::POST_PASSED
                | ConfigFlags::FIRST_DISABLED
                | ConfigFlags::SECOND_DISABLED;
            self.set_config(config)?;
        }

        // Linux i8042_controller_selftest(): retry up to 5 times with delay.
        // "On some really fragile systems this does not take the first time."
        {
            let mut passed = false;
            for attempt in 0..SELFTEST_RETRIES {
                if let Err(err) = self.command(Command::TestController) {
                    warn!("self-test command failed (attempt {}): {:?}", attempt + 1, err);
                    continue;
                }
                match self.read() {
                    Ok(SELFTEST_PASS) => {
                        passed = true;
                        break;
                    }
                    Ok(val) => {
                        warn!(
                            "self-test unexpected value {:02X} (attempt {}/{})",
                            val,
                            attempt + 1,
                            SELFTEST_RETRIES
                        );
                    }
                    Err(err) => {
                        warn!("self-test read timeout (attempt {}): {:?}", attempt + 1, err);
                    }
                }
                // Linux: msleep(50) between retries
                std::thread::sleep(std::time::Duration::from_millis(50));
            }
            if !passed {
                // Linux on x86: "giving up on controller selftest, continuing anyway"
                warn!("controller self-test did not pass after {} attempts, continuing", SELFTEST_RETRIES);
            }
        }

        // Flush any bytes the self-test may have left behind
        self.flush();

        // Linux i8042_controller_init() step 2: set keyboard defaults
        // (disable scanning so keyboard doesn't send scancodes during init)
        self.retry(format_args!("keyboard defaults"), 4, |x| {
            let b = x.keyboard_command(KeyboardCommand::SetDefaultsDisable)?;
            if b != 0xFA {
                error!("keyboard failed to set defaults: {:02X}", b);
                return Err(Error::CommandRetry);
            }
            Ok(b)
        })?;

        // Initialize keyboard
        if let Err(err) = self.init_keyboard() {
            error!("failed to initialize keyboard: {:?}", err);
            return Err(err);
        }

        // Linux: test AUX port (command 0xA9) before enabling.
        // Skips mouse init entirely if the port is not electrically present.
        let aux_ok = self.test_aux_port();

        // Enable second device (mouse) only if AUX port tested OK
        let enable_mouse = if aux_ok {
            match self.command(Command::EnableSecond) {
                Ok(()) => true,
                Err(err) => {
                    warn!("failed to enable aux port after test passed: {:?}", err);
                    false
                }
            }
        } else {
            info!("skipping mouse init: aux port test did not pass");
            false
        };

        {
            if let Err(err) = self.keyboard_command_inner(KeyboardCommand::EnableReporting as u8) {
                error!("failed to initialize keyboard reporting: {:?}", err);
            }
        }

        // Enable clocks and interrupts
        {
            let config = ConfigFlags::POST_PASSED
                | ConfigFlags::FIRST_INTERRUPT
                | ConfigFlags::FIRST_TRANSLATE
                | if enable_mouse {
                    ConfigFlags::SECOND_INTERRUPT
                } else {
                    ConfigFlags::SECOND_DISABLED
                };
            self.set_config(config)?;
        }

        Ok(())
    }
}