use aarch64_cpu::registers::*;
use tock_registers::interfaces::Readable;

use super::TrapFrame;
use crate::trap::PageFaultFlags;

#[repr(u8)]
#[derive(Debug)]
pub(super) enum TrapKind {
    Synchronous = 0,
    Irq         = 1,
    Fiq         = 2,
    SError      = 3,
}

#[repr(u8)]
#[derive(Debug)]
enum TrapSource {
    CurrentSpEl0 = 0,
    CurrentSpElx = 1,
    LowerAArch64 = 2,
    LowerAArch32 = 3,
}

core::arch::global_asm!(
    #[cfg(not(feature = "arm-el2"))]
    include_str!("trap.S"),
    #[cfg(feature = "arm-el2")]
    concat!(".equ arm_el2, 1\n", include_str!("trap.S")),
    trapframe_size = const core::mem::size_of::<TrapFrame>(),
    TRAP_KIND_SYNC = const TrapKind::Synchronous as u8,
    TRAP_KIND_IRQ = const TrapKind::Irq as u8,
    TRAP_KIND_FIQ = const TrapKind::Fiq as u8,
    TRAP_KIND_SERROR = const TrapKind::SError as u8,
    TRAP_SRC_CURR_EL0 = const TrapSource::CurrentSpEl0 as u8,
    TRAP_SRC_CURR_ELX = const TrapSource::CurrentSpElx as u8,
    TRAP_SRC_LOWER_AARCH64 = const TrapSource::LowerAArch64 as u8,
    TRAP_SRC_LOWER_AARCH32 = const TrapSource::LowerAArch32 as u8,
);

#[inline(always)]
pub(super) fn is_valid_page_fault(iss: u64) -> bool {
    // Only handle Translation fault and Permission fault
    matches!(iss & 0b111100, 0b0100 | 0b1100) // IFSC or DFSC bits
}

#[inline(always)]
fn fault_addr() -> usize {
    #[cfg(not(feature = "arm-el2"))]
    {
        FAR_EL1.get() as usize
    }

    #[cfg(feature = "arm-el2")]
    {
        FAR_EL2.get() as usize
    }
}

#[inline(always)]
fn esr_value() -> u64 {
    #[cfg(not(feature = "arm-el2"))]
    {
        ESR_EL1.get()
    }

    #[cfg(feature = "arm-el2")]
    {
        ESR_EL2.get()
    }
}

fn handle_breakpoint(tf: &mut TrapFrame) {
    if crate::trap::breakpoint_handler(tf) {
        return;
    }
    tf.elr += 4;
}

fn handle_page_fault(tf: &mut TrapFrame, access_flags: PageFaultFlags) {
    let vaddr = va!(fault_addr());
    if crate::trap::call_page_fault_handler_with_parent_irqs(
        vaddr,
        access_flags,
        tf.spsr & (1 << 7) == 0,
    ) {
        return;
    }
    #[cfg(feature = "exception-table")]
    if tf.fixup_exception() {
        return;
    }
    let bt = tf.backtrace();
    panic!(
        "Unhandled Page Fault @ {:#x}, fault_vaddr={:#x}, ESR={:#x} ({:?}):\n{:#x?}\n{}",
        tf.elr,
        vaddr,
        esr_value(),
        access_flags,
        tf,
        bt.kind("trap")
    );
}

#[unsafe(no_mangle)]
fn aarch64_trap_handler(tf: &mut TrapFrame, kind: TrapKind, source: TrapSource) {
    if matches!(
        source,
        TrapSource::CurrentSpEl0 | TrapSource::LowerAArch64 | TrapSource::LowerAArch32
    ) {
        let bt = tf.backtrace();
        panic!(
            "Invalid exception {:?} from {:?}:\n{:#x?}\n{}",
            kind,
            source,
            tf,
            bt.kind("trap")
        );
    }
    match kind {
        TrapKind::Fiq | TrapKind::SError => {
            let bt = tf.backtrace();
            panic!(
                "Unhandled exception {:?}:\n{:#x?}\n{}",
                kind,
                tf,
                bt.kind("trap")
            );
        }
        TrapKind::Irq => {
            crate::trap::dispatch_irq(0);
        }
        TrapKind::Synchronous => {
            #[cfg(not(feature = "arm-el2"))]
            let esr = ESR_EL1.extract();
            #[cfg(feature = "arm-el2")]
            let esr = ESR_EL2.extract();

            #[cfg(not(feature = "arm-el2"))]
            let iss = esr.read(ESR_EL1::ISS);
            #[cfg(feature = "arm-el2")]
            let iss = esr.read(ESR_EL2::ISS);

            #[cfg(not(feature = "arm-el2"))]
            let ec = esr.read_as_enum(ESR_EL1::EC);
            #[cfg(feature = "arm-el2")]
            let ec = esr.read_as_enum(ESR_EL2::EC);

            match ec {
                #[cfg(not(feature = "arm-el2"))]
                Some(ESR_EL1::EC::Value::InstrAbortCurrentEL) if is_valid_page_fault(iss) => {
                    handle_page_fault(tf, PageFaultFlags::EXECUTE);
                }
                #[cfg(feature = "arm-el2")]
                Some(ESR_EL2::EC::Value::InstrAbortCurrentEL) if is_valid_page_fault(iss) => {
                    handle_page_fault(tf, PageFaultFlags::EXECUTE);
                }
                #[cfg(not(feature = "arm-el2"))]
                Some(ESR_EL1::EC::Value::DataAbortCurrentEL) if is_valid_page_fault(iss) => {
                    let wnr = (iss & (1 << 6)) != 0; // WnR: Write not Read
                    let cm = (iss & (1 << 8)) != 0; // CM: Cache maintenance
                    handle_page_fault(
                        tf,
                        if wnr & !cm {
                            PageFaultFlags::WRITE
                        } else {
                            PageFaultFlags::READ
                        },
                    );
                }
                #[cfg(feature = "arm-el2")]
                Some(ESR_EL2::EC::Value::DataAbortCurrentEL) if is_valid_page_fault(iss) => {
                    let wnr = (iss & (1 << 6)) != 0; // WnR: Write not Read
                    let cm = (iss & (1 << 8)) != 0; // CM: Cache maintenance
                    handle_page_fault(
                        tf,
                        if wnr & !cm {
                            PageFaultFlags::WRITE
                        } else {
                            PageFaultFlags::READ
                        },
                    );
                }
                #[cfg(not(feature = "arm-el2"))]
                Some(ESR_EL1::EC::Value::Brk64) => {
                    debug!("BRK #{:#x} @ {:#x} ", iss, tf.elr);
                    handle_breakpoint(tf);
                }
                #[cfg(feature = "arm-el2")]
                Some(ESR_EL2::EC::Value::Brk64) => {
                    debug!("BRK #{:#x} @ {:#x} ", iss, tf.elr);
                    handle_breakpoint(tf);
                }
                e => {
                    let vaddr = va!(fault_addr());

                    #[cfg(not(feature = "arm-el2"))]
                    let ec_bits = esr.read(ESR_EL1::EC);
                    #[cfg(feature = "arm-el2")]
                    let ec_bits = esr.read(ESR_EL2::EC);

                    let bt = tf.backtrace();
                    panic!(
                        "Unhandled synchronous exception {:?} @ {:#x}: ESR={:#x} (EC {:#08b}, \
                         FAR: {:#x} ISS {:#x})\n{}",
                        e,
                        tf.elr,
                        esr.get(),
                        ec_bits,
                        vaddr,
                        iss,
                        bt.kind("trap")
                    );
                }
            }
        }
    }
}
