//! x86-64 exception and interrupt handling
use core::{arch::global_asm, mem::size_of};

use abi::{bitflags, primitive_enum, process::Signal};
use kernel_arch_x86::registers::{CR2, CR3};
use kernel_arch_x86_64::context::ExceptionFrame;
use libk::{arch::Cpu, task::thread::Thread};
use libk_mm::PageFaultKind;
use libk_util::sync::spin_rwlock::IrqSafeRwLock;
use tock_registers::interfaces::Readable;

use crate::arch::x86_64::{apic, X86_64};

primitive_enum! {
    enum ExceptionKind: u64 {
        DivisionError = 0,
        Debug = 1,
        NonMaskableInterrupt = 2,
        Breakpoint = 3,
        Overflow = 4,
        BoundRangeExceeded = 5,
        InvalidOpcode = 6,
        DeviceNotAvailable = 7,
        DoubleFault = 8,
        InvalidTss = 10,
        SegmentNotPresent = 11,
        StackSegmentFault = 12,
        GeneralProtectionFault = 13,
        PageFault = 14,
        FpuException = 16,
        AlignmentCheck = 17,
        MachineCheck = 18,
        SimdFpuException = 19,
        VirtualizationException = 20,
        ControlProtectionException = 21,

        Unknown = 99,
    }
}

impl ExceptionKind {
    fn ring3_possible(&self) -> bool {
        matches!(
            self,
            Self::DivisionError
                | Self::Debug
                | Self::Breakpoint
                | Self::Overflow
                | Self::BoundRangeExceeded
                | Self::InvalidOpcode
                | Self::GeneralProtectionFault
                | Self::PageFault
                | Self::FpuException
                | Self::AlignmentCheck
                | Self::SimdFpuException
        )
    }
}

/// Exception table entry
#[allow(dead_code)]
#[derive(Clone, Copy)]
#[repr(C, packed)]
pub struct Entry {
    base_lo: u16,
    selector: u16,
    __res0: u8,
    flags: u8,
    base_hi: u16,
    base_ex: u32,
    __res1: u32,
}

#[allow(dead_code)]
#[repr(C, packed)]
struct Pointer {
    limit: u16,
    offset: usize,
}

const SIZE: usize = 256;

impl Entry {
    /// Entry is valid
    pub const PRESENT: u8 = 1 << 7;
    /// Entry is a 32-bit interrupt
    pub const INT32: u8 = 0xE;

    /// Empty entry
    pub const NULL: Self = Self {
        base_lo: 0,
        base_hi: 0,
        base_ex: 0,
        selector: 0,
        flags: 0,
        __res0: 0,
        __res1: 0,
    };

    /// Constructs an interrupt table entry
    pub const fn new(base: usize, selector: u16, flags: u8) -> Self {
        Self {
            base_lo: (base & 0xFFFF) as u16,
            base_hi: ((base >> 16) & 0xFFFF) as u16,
            base_ex: (base >> 32) as u32,
            selector,
            flags,
            __res0: 0,
            __res1: 0,
        }
    }
}

static IDT: IrqSafeRwLock<[Entry; SIZE]> = IrqSafeRwLock::new([Entry::NULL; SIZE]);

bitflags! {
    struct PageFaultFlags: usize {
        /// If clear, fault was caused by a non-present page entry
        const P: bit 0;
        /// If set, fault was caused by a write
        const W: bit 1;
    }
}

fn dump_user_exception(kind: ExceptionKind, frame: &ExceptionFrame) {
    let thread = Thread::current();
    log::warn!("{:?} in {} ({:?})", kind, thread.id, *thread.name.read());
    log::warn!("ip = {:02x}:{:08x}", frame.cs, frame.rip);
    log::warn!("cr3 = {:#010x}", CR3.get());
    if kind == ExceptionKind::PageFault {
        log::warn!("cr2 = {:#010x}", CR2.get());
    }
}

fn user_exception_inner(kind: ExceptionKind, frame: &mut ExceptionFrame) {
    let thread = Thread::current();

    let dump = match kind {
        ExceptionKind::Debug => {
            if thread.handle_single_step(frame) {
                false
            } else {
                thread.raise_signal(Signal::Aborted);
                true
            }
        }
        ExceptionKind::PageFault => {
            let cr2 = CR2.get();
            let flags = PageFaultFlags::new(frame.exc_code as usize);
            let kind = match (
                flags.contains(PageFaultFlags::P),
                flags.contains(PageFaultFlags::W),
            ) {
                (true, true) => PageFaultKind::WriteFault,
                (true, false) => PageFaultKind::Other,
                (false, write) => PageFaultKind::TranslationFault { write },
            };

            if thread.handle_page_fault(cr2, kind).is_ok() {
                false
            } else {
                thread.raise_signal(Signal::MemoryAccessViolation);
                true
            }
        }
        ExceptionKind::GeneralProtectionFault => {
            if thread.handle_breakpoint(frame) {
                false
            } else {
                thread.raise_signal(Signal::MemoryAccessViolation);
                true
            }
        }
        ExceptionKind::FpuException => {
            thread.raise_signal(Signal::Aborted);
            true
        }
        ExceptionKind::InvalidOpcode => {
            // TODO handle ud2 as breakpoint? (it's 2 bytes)
            thread.raise_signal(Signal::Aborted);
            true
        }
        ExceptionKind::Breakpoint => {
            thread.raise_signal(Signal::Aborted);
            true
        }
        ExceptionKind::SimdFpuException => {
            thread.raise_signal(Signal::MemoryAccessViolation);
            true
        }
        _ => {
            log::warn!("No handler for exception: {:?}", kind);
            thread.raise_signal(Signal::Aborted);
            true
        }
    };

    if dump {
        dump_user_exception(kind, frame);
    }
}

fn kernel_exception_inner(kind: ExceptionKind, frame: &ExceptionFrame) -> ! {
    let cr3 = CR3.get();
    log::error!("{:?} in KERNEL, frame {:p}, cr3 = {:#x}", kind, frame, cr3);

    if kind == ExceptionKind::PageFault {
        let cr2 = CR2.get();
        log::error!("cr2 = {:#x}", cr2);
    }

    // XXX
    // frame.dump(debug::LogLevel::Fatal);

    panic!("Irrecoverable exception");
}

extern "C" fn __x86_64_exception_handler(frame: *mut ExceptionFrame) {
    let frame = unsafe { &mut *frame };
    let kind = ExceptionKind::try_from(frame.exc_number).unwrap_or(ExceptionKind::Unknown);

    if kind.ring3_possible() && frame.cs == 0x23 {
        user_exception_inner(kind, frame);

        unsafe {
            Thread::current().handle_pending_signals(frame);
        }
    } else {
        kernel_exception_inner(kind, frame)
    }
}

extern "C" fn __x86_64_nmi_handler() -> ! {
    let ipi = Cpu::local().get_ipi();
    if let Some(ipi) = ipi {
        X86_64::handle_ipi(Some(ipi));
        unreachable!();
    } else {
        panic!("Spurious/unknown NMI received");
    }
}

/// Initializes the interrupt descriptor table for the given CPU.
///
/// # Safety
///
/// Only meant to be called once per each CPU during their init.
pub unsafe fn init_exceptions(cpu_index: usize) {
    if cpu_index == 0 {
        let mut idt = IDT.write();
        extern "C" {
            static __x86_64_exception_vectors: [usize; 32];
        }

        for (i, &entry) in __x86_64_exception_vectors.iter().enumerate() {
            idt[i] = Entry::new(entry, 0x08, Entry::PRESENT | Entry::INT32);
        }

        apic::setup_vectors(&mut idt[32..]);
    }

    let idt = IDT.read();
    let idtr = Pointer {
        limit: (idt.len() * size_of::<Entry>()) as u16 - 1,
        offset: idt.as_ptr().addr(),
    };

    core::arch::asm!("wbinvd; lidt ({0})", in(reg) &idtr, options(att_syntax));
}

global_asm!(
    include_str!("vectors.S"),
    exception_handler = sym __x86_64_exception_handler,
    nmi_handler = sym __x86_64_nmi_handler,
    options(att_syntax)
);
