//! Exception and interrupt management functions

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

use aarch64_cpu::{
    asm::barrier,
    registers::{
        ELR_EL1, ESR_EL1, FAR_EL1, SCTLR_EL1, SP_EL0, TCR_EL1, TPIDR_EL0, TPIDR_EL1, TTBR0_EL1,
        TTBR1_EL1, VBAR_EL1,
    },
};
use abi::{process::Signal, SyscallFunction};
use kernel_arch::{sync::hack_locks, Architecture, ArchitectureImpl};
use kernel_arch_aarch64::{
    context::ExceptionFrame,
    mem::table::{EntryType, PageTable, L1, L2, L3},
};
use libk::{device::external_interrupt_controller, task::thread::Thread};
use libk_mm::{address::PhysicalAddress, table::EntryLevelExt, PageFaultKind};
use tock_registers::interfaces::{Readable, Writeable};

use crate::{
    mem::KERNEL_VIRT_OFFSET,
    syscall::{self, raw_syscall_handler},
};

/// Initializes the exception/interrupt vectors. May be called repeatedly (though that makes no
/// sense).
pub fn init_exceptions() {
    extern "C" {
        static __aarch64_el1_vectors: u8;
    }
    let vbar = unsafe { &__aarch64_el1_vectors as *const _ };
    VBAR_EL1.set(vbar as u64);
    barrier::isb(barrier::SY);
}

unsafe fn perform_ptw<F: Fn(u32, EntryType)>(virt: usize, handler: F) {
    let l1i = virt.page_index::<L1>();
    let l2i = virt.page_index::<L2>();
    let l3i = virt.page_index::<L3>();

    let ttbr_phys = if virt > KERNEL_VIRT_OFFSET {
        TTBR1_EL1.get()
    } else {
        TTBR0_EL1.get() & !1
    };

    // Strip ASID
    let ttbr_phys = PhysicalAddress::from_u64(ttbr_phys & !(0xFFFF << 48));
    let Some(l1) = PageTable::<L1>::from_physical(ttbr_phys) else {
        handler(0, EntryType::Invalid);
        return;
    };
    handler(0, EntryType::Table(ttbr_phys));

    let l2 = l1.walk(l1i);
    handler(1, l2);
    let l2 = match l2 {
        EntryType::Table(l2) => PageTable::<L2>::from_physical(l2).unwrap(),
        _ => return,
    };

    let l3 = l2.walk(l2i);
    handler(2, l3);
    let l3 = match l3 {
        EntryType::Table(l3) => PageTable::<L3>::from_physical(l3).unwrap(),
        _ => return,
    };

    let l3e = match l3[l3i].as_page() {
        Some(page) => EntryType::Page(page),
        None => EntryType::Invalid,
    };
    handler(3, l3e);
}

fn dump_irrecoverable_exception(frame: &ExceptionFrame, ec: u64, iss: u64) {
    let far = FAR_EL1.get() as usize;
    let ttbr0 = TTBR0_EL1.get();

    log::error!(target: ":raw", "SYNC exception:\n");
    log::error!(target: ":raw", "FAR: {:#x}\n", far);
    log::error!(target: ":raw", "ELR: {:#x}\n", ELR_EL1.get());
    log::error!(target: ":raw", "ESR: {:#x}\n", ESR_EL1.get());
    log::error!(target: ":raw", "SP_EL0: {:#x}\n", SP_EL0.get());
    log::error!(target: ":raw", "TTBR0_EL1: {:#x}\n", ttbr0);
    log::error!(target: ":raw", "TTBR1_EL1: {:#x}\n", TTBR1_EL1.get());
    log::error!(target: ":raw", "Register dump:\n");
    log::error!(target: ":raw", "{:?}\n", frame);

    log::error!(target: ":raw", "FAR_EL1 page table walk:\n");
    unsafe {
        perform_ptw(far, |level, result| {
            log::error!(target: ":raw", "[{level}] {result}\n");
        });
    }

    match ec {
        // Data abort from lower level
        0b100100 => {
            log::error!(target: ":raw", "Exception kind: Data Abort from EL0\n");
            let dfsc = iss & 0x3F;

            if iss & (1 << 24) != 0 {
                let access_size_str = match (iss >> 22) & 0x3 {
                    0 => "i8",
                    1 => "i16",
                    2 => "i32",
                    3 => "i64",
                    _ => unreachable!(),
                };
                let access_type_str = if iss & (1 << 6) != 0 { "write" } else { "read" };

                log::error!(
                    target: ":raw",
                    "Invalid {} of a {} to/from {:#x}\n",
                    access_type_str,
                    access_size_str,
                    FAR_EL1.get()
                );
            }

            log::error!(target: ":raw", "DFSC = {:#x}\n", dfsc);
        }
        // Instruction abort from lower level
        0b100000 => {
            log::error!(
                target: ":raw",
                "Exception kind: Instruction Abort from EL0\n"
            );
            let ifsc = iss & 0x3F;
            log::error!(target: ":raw", "IFSC = {:#x}\n", ifsc);
        }

        _ => (),
    }

    log::error!(target: ":raw", "System register dump:\n");
    log::error!(target: ":raw", "SCTLR_EL1 = {:#x}\n", SCTLR_EL1.get());
    log::error!(target: ":raw", "TCR_EL1 = {:#x}\n", TCR_EL1.get());
    log::error!(target: ":raw", "TPIDR_EL1 = {:#x}\n", TPIDR_EL1.get());
    log::error!(target: ":raw", "TPIDR_EL0 = {:#x}\n", TPIDR_EL0.get());
}

#[no_mangle]
extern "C" fn __aa64_el0_sync_handler(frame: *mut ExceptionFrame) {
    assert!(ArchitectureImpl::interrupt_mask());
    let frame = unsafe { &mut *frame };

    el0_sync_inner(frame);

    unsafe {
        let thread = Thread::current();
        thread.handle_pending_signals(frame);
    }
}

#[no_mangle]
extern "C" fn __aa64_el0_irq_handler(frame: *mut ExceptionFrame) {
    assert!(ArchitectureImpl::interrupt_mask());
    let frame = unsafe { &mut *frame };

    irq_common();

    unsafe {
        let thread = Thread::current();
        thread.handle_pending_signals(frame);
    }
}

#[no_mangle]
extern "C" fn __aa64_el0_fiq_handler() {
    unimplemented!()
}

#[no_mangle]
extern "C" fn __aa64_el0_serror_handler() {
    unimplemented!()
}

// EL1
static EL1_FAULT_TAKEN: AtomicBool = AtomicBool::new(false);

#[no_mangle]
extern "C" fn __aa64_el1_sync_handler(frame: *mut ExceptionFrame) {
    if EL1_FAULT_TAKEN.swap(true, Ordering::Release) {
        // Fault already taken
        ArchitectureImpl::halt();
    }

    let frame = unsafe { &*frame };
    let esr_el1 = ESR_EL1.get();
    let ec = (esr_el1 >> 26) & 0x3F;
    let iss = esr_el1 & 0x1FFFFFF;

    unsafe { hack_locks() };

    dump_irrecoverable_exception(frame, ec, iss);

    panic!("Irrecoverable exception in kernel mode");
}

#[no_mangle]
extern "C" fn __aa64_el1_irq_handler(_frame: *mut ExceptionFrame) {
    irq_common();
}

#[no_mangle]
extern "C" fn __aa64_el1_fiq_handler() {
    unimplemented!()
}

#[no_mangle]
extern "C" fn __aa64_el1_serror_handler() {
    unimplemented!()
}

fn el0_sync_inner(frame: &mut ExceptionFrame) {
    let esr_el1 = ESR_EL1.get();
    let ec = (esr_el1 >> 26) & 0x3F;
    let iss = esr_el1 & 0x1FFFFFF;

    let dump = match ec {
        // SVC in AArch64
        0b010101 => {
            let seq = Thread::current().trace_syscall_entry(frame);
            let func = frame.r[8];
            if func == usize::from(SyscallFunction::ExitSignal) {
                unsafe {
                    syscall::handle_signal_exit(frame);
                }
                return;
            }

            let args = &frame.r[0..6];
            let result = raw_syscall_handler(func, args) as _;
            frame.r[0] = result;
            Thread::current().trace_syscall_exit(frame, seq);
            false
        }
        // Software Step from lower Exception Level
        0b110010 => {
            let thread = Thread::current();

            if thread.handle_single_step(frame) {
                // Make the PE actually step the instruction
                frame.spsr_el1 |= 1 << 21;
                false
            } else {
                thread.raise_signal(Signal::Aborted);
                true
            }
        }
        // BRK in AArch64
        0b111100 => {
            let thread = Thread::current();
            log::warn!(
                "Thread {} {:?} hit a breakpoint",
                thread.id,
                *thread.name.read()
            );
            thread.raise_signal(Signal::Aborted);
            false
        }
        // Data abort from a lower EL
        0b100100 => {
            let thread = Thread::current();
            let dfsc = esr_el1 & 0x3F;
            let wnr = esr_el1 & (1 << 6) != 0;
            let fnv = esr_el1 & (1 << 10) != 0;
            let far = FAR_EL1.get() as usize;
            let kind = match dfsc {
                // Translation fault, level 3
                0b000111 => Some(PageFaultKind::TranslationFault { write: wnr }),
                // Permission fault, level 3
                0b001111 if wnr => Some(PageFaultKind::WriteFault),
                _ => None,
            };
            if let Some(kind) = kind
                && !fnv
                && thread.handle_page_fault(far, kind).is_ok()
            {
                false
            } else {
                log::warn!(
                    "Data abort in {} {:?} at {:#x} with address {:#x}",
                    thread.id,
                    *thread.name.read(),
                    ELR_EL1.get(),
                    FAR_EL1.get()
                );
                thread.raise_signal(Signal::MemoryAccessViolation);
                true
            }
        }
        _ => {
            let thread = Thread::current();
            thread.raise_signal(Signal::MemoryAccessViolation);
            true
        }
    };

    if dump {
        dump_irrecoverable_exception(frame, ec, iss);
    }
}

fn irq_common() {
    // TODO some concept of "root" interrupt controller
    external_interrupt_controller()
        .unwrap()
        .handle_pending_irqs();
}

global_asm!(include_str!("vectors.S"));
