use core::sync::atomic::Ordering;

use crate::cap::cnode;
use crate::cap::object::ObjectTag;
use crate::cap::ops;
use crate::cap::pool::POOL;
use crate::cap::table::Rights;
use crate::error::KernelError;
use crate::ipc::message;
use crate::ipc::{endpoint, notification};
use crate::mem::addr;
use crate::proc::context::IpcMessage;
use crate::proc::{BlockedReason, PROCESSES};
use crate::ring::{
    CompletionEntry, MAX_CQ_ENTRIES, MAX_SQ_ENTRIES, RingHeader, RingIndex, RingOpcode,
    SubmissionEntry,
};
use crate::types::Pid;

const MAX_RING_BATCH: u32 = 16;

fn _assert_from_bytes<T: zerocopy::FromBytes>() {}
#[allow(dead_code)]
fn _assert_ring_types_are_zerocopy() {
    _assert_from_bytes::<SubmissionEntry>();
    _assert_from_bytes::<CompletionEntry>();
}

struct UserSnapshot {
    sq_tail: RingIndex,
    cq_head: RingIndex,
}

unsafe fn snapshot_user_fields(ring_base: *const u8) -> UserSnapshot {
    let header = unsafe { &*(ring_base as *const RingHeader) };
    UserSnapshot {
        sq_tail: RingIndex::new(header.sq_tail.load(Ordering::Acquire)),
        cq_head: RingIndex::new(header.cq_head.load(Ordering::Acquire)),
    }
}

unsafe fn commit_header(ring_base: *mut u8, new_sq_head: RingIndex, new_cq_tail: RingIndex) {
    let header = unsafe { &*(ring_base as *const RingHeader) };
    header.sq_head.store(new_sq_head.raw(), Ordering::Release);
    header.cq_tail.store(new_cq_tail.raw(), Ordering::Release);
}

unsafe fn read_sq_entry(ring_base: *const u8, index: u32) -> SubmissionEntry {
    let sq_base = unsafe { ring_base.add(super::ring_sq_offset()) };
    let entry_ptr =
        unsafe { sq_base.add((index as usize) * core::mem::size_of::<SubmissionEntry>()) };
    unsafe { core::ptr::read_volatile(entry_ptr as *const SubmissionEntry) }
}

unsafe fn write_cq_entry(ring_base: *mut u8, index: u32, entry: CompletionEntry) {
    let cq_base = unsafe { ring_base.add(super::ring_cq_offset()) };
    let entry_ptr =
        unsafe { cq_base.add((index as usize) * core::mem::size_of::<CompletionEntry>()) };
    unsafe { core::ptr::write_volatile(entry_ptr as *mut CompletionEntry, entry) };
}


struct CqeResult {
    value: i64,
    extra: u64,
}

impl CqeResult {
    const fn ok() -> Self {
        Self { value: 0, extra: 0 }
    }

    const fn success(value: i64) -> Self {
        Self { value, extra: 0 }
    }

    const fn with_extra(value: i64, extra: u64) -> Self {
        Self { value, extra }
    }
}

fn ring_cap_create(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let address = sqe.cap_slot as u64;
    let tag = ObjectTag::try_from(sqe.arg0)?;

    let ptable = PROCESSES.lock();
    let (cnode_id, cnode_gen, depth) = cnode::cnode_coords(pid, &ptable)?;
    let mut pool = POOL.lock_after(&ptable);
    ops::create_via_cnode(&mut pool, cnode_id, cnode_gen, address, depth, tag)
        .map(|id| CqeResult::success(id.raw() as i64))
}

fn ring_cap_derive(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let src_addr = sqe.cap_slot as u64;
    let dest_addr = sqe.arg0;
    let rights_mask = Rights::from_bits(sqe.arg1 as u16);

    let ptable = PROCESSES.lock();
    let (cnode_id, cnode_gen, depth) = cnode::cnode_coords(pid, &ptable)?;
    let mut pool = POOL.lock_after(&ptable);
    ops::derive_via_cnode(&mut pool, cnode_id, cnode_gen, src_addr, dest_addr, depth, rights_mask)
        .map(|()| CqeResult::ok())
}

fn ring_ipc_send(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let address = sqe.cap_slot as u64;
    let msg = IpcMessage::from_regs([sqe.arg0, sqe.arg1, sqe.arg2, 0, 0, 0]);

    let mut ptable = PROCESSES.lock();
    let cap = {
        let pool = POOL.lock_after(&ptable);
        cnode::resolve_caller_validate(
            pid, address, ObjectTag::Endpoint, Rights::WRITE, &ptable, &pool,
        )?
    };

    ptable[pid].ipc_message = msg;

    let recv_pid = {
        let mut pool = POOL.lock_after(&ptable);
        let ep = pool
            .get_mut(cap.object_id(), cap.generation())?
            .as_endpoint_mut()?;
        let dequeued = endpoint::dequeue_genuine_receiver(&mut ep.receivers, &mut ptable);
        if let Some(ref br) = dequeued {
            ep.holder = Some(br.pid());
        }
        dequeued
    };

    match recv_pid {
        Some(blocked_recv) => {
            let recv_pid_val = blocked_recv.pid();
            let sender_prio = crate::sched::effective_priority_of(&ptable[pid]);

            let recv = &mut ptable[recv_pid_val];
            recv.ipc_message = msg;
            recv.ipc_badge = pid.raw() as u64;
            recv.unblock(blocked_recv)?;
            message::inject_into_context(&mut recv.saved_context, &msg);
            recv.saved_context.rax = pid.raw() as u64;
            recv.seal_context();

            crate::sched::boost_effective(&mut ptable, recv_pid_val, sender_prio);

            Ok(CqeResult::ok())
        }
        None => Err(KernelError::WouldBlock),
    }
}

fn ring_ipc_recv(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let address = sqe.cap_slot as u64;

    let mut ptable = PROCESSES.lock();
    let cap = {
        let pool = POOL.lock_after(&ptable);
        cnode::resolve_caller_validate(
            pid, address, ObjectTag::Endpoint, Rights::READ, &ptable, &pool,
        )?
    };

    let sender_pid = {
        let mut pool = POOL.lock_after(&ptable);
        let ep = pool
            .get_mut(cap.object_id(), cap.generation())?
            .as_endpoint_mut()?;
        let dequeued = endpoint::dequeue(&mut ep.senders, &mut ptable);
        if dequeued.is_some() {
            ep.holder = Some(pid);
        }
        dequeued
    };

    match sender_pid {
        Some(blocked_sender) => {
            let sender_pid = blocked_sender.pid();
            let sender_msg = ptable[sender_pid].ipc_message;
            ptable[pid].ipc_message = sender_msg;
            ptable[pid].ipc_badge = sender_pid.raw() as u64;

            match ptable[sender_pid].blocked_reason() {
                Some(BlockedReason::Calling(_, _)) => {
                    let mut pool = POOL.lock_after(&ptable);
                    match pool
                        .get_mut(cap.object_id(), cap.generation())
                        .and_then(|d| d.as_endpoint_mut())
                    {
                        Ok(ep) => {
                            match endpoint::enqueue(&mut ep.receivers, blocked_sender, &mut ptable)
                            {
                                Ok(()) => {
                                    ptable[pid].reply_target = Some(sender_pid);
                                }
                                Err(_) => {
                                    ptable[sender_pid].saved_context.rax =
                                        crate::error::KernelError::ResourceExhausted.to_errno()
                                            as u64;
                                    ptable[sender_pid].seal_context();
                                    {
                                        let proof = ptable[sender_pid].blocked_proof();
                                        let r = ptable[sender_pid].unblock(proof);
                                        debug_assert!(r.is_ok());
                                    }
                                }
                            }
                        }
                        Err(e) => {
                            ptable[sender_pid].saved_context.rax = e.to_errno() as u64;
                            ptable[sender_pid].seal_context();
                            {
                                let r = ptable[sender_pid].unblock(blocked_sender);
                                debug_assert!(r.is_ok());
                            }
                        }
                    }
                }
                _ => {
                    ptable[sender_pid].unblock(blocked_sender)?;
                }
            }

            let sender_prio = crate::sched::effective_priority_of(&ptable[sender_pid]);
            crate::sched::reset_effective(&mut ptable, pid);
            crate::sched::boost_effective(&mut ptable, pid, sender_prio);

            Ok(CqeResult::with_extra(
                sender_pid.raw() as i64,
                sender_msg.regs[0],
            ))
        }
        None => Err(KernelError::WouldBlock),
    }
}

fn ring_notify_signal(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let address = sqe.cap_slot as u64;
    let bits = sqe.arg0;

    let mut ptable = PROCESSES.lock();
    let cap = {
        let pool = POOL.lock_after(&ptable);
        cnode::resolve_caller_validate(
            pid, address, ObjectTag::Notification, Rights::WRITE, &ptable, &pool,
        )?
    };

    notification::do_signal(&cap, bits, &mut ptable).map(|_| CqeResult::ok())
}

fn ring_notify_poll(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let address = sqe.cap_slot as u64;

    let ptable = PROCESSES.lock();
    let cap = {
        let pool = POOL.lock_after(&ptable);
        cnode::resolve_caller_validate(
            pid, address, ObjectTag::Notification, Rights::READ, &ptable, &pool,
        )?
    };

    drop(ptable);
    notification::do_poll(&cap).map(|val| CqeResult::with_extra(0, val))
}

fn process_submission(sqe: &SubmissionEntry, pid: Pid) -> CompletionEntry {
    let result = match RingOpcode::from_u8(sqe.opcode) {
        Some(RingOpcode::Nop) => Ok(CqeResult::ok()),
        Some(RingOpcode::CapCreate) => ring_cap_create(sqe, pid),
        Some(RingOpcode::CapDerive) => ring_cap_derive(sqe, pid),
        Some(RingOpcode::IpcSend) => ring_ipc_send(sqe, pid),
        Some(RingOpcode::IpcRecv) => ring_ipc_recv(sqe, pid),
        Some(RingOpcode::NotifySignal) => ring_notify_signal(sqe, pid),
        Some(RingOpcode::NotifyPoll) => ring_notify_poll(sqe, pid),
        None => Err(KernelError::InvalidParameter),
    };

    match result {
        Ok(cqe) => CompletionEntry {
            result: cqe.value,
            user_data: sqe.user_data as u64,
            extra: cqe.extra,
        },
        Err(e) => CompletionEntry {
            result: e.to_errno(),
            user_data: sqe.user_data as u64,
            extra: 0,
        },
    }
}

pub fn ring_enter(
    ring_phys_base: x86_64::PhysAddr,
    pid: Pid,
    _min_complete: u32,
) -> Result<i64, KernelError> {
    if !ring_phys_base.as_u64().is_multiple_of(4096) {
        return Err(KernelError::InvalidAddress);
    }

    let max_phys = (crate::mem::phys::BitmapFrameAllocator::total_frames() as u64) * 4096;
    let ring_end = ring_phys_base
        .as_u64()
        .checked_add(super::ring_total_size() as u64)
        .ok_or(KernelError::InvalidAddress)?;
    if ring_end > max_phys {
        return Err(KernelError::InvalidAddress);
    }

    let hhdm = addr::hhdm_offset();
    let ring_virt = (ring_phys_base.as_u64() + hhdm) as *mut u8;

    let mut ptable = PROCESSES.lock();
    let proc = ptable.get_mut(pid).ok_or(KernelError::InvalidObject)?;

    let sq_head = proc.ring_sq_head;
    let cq_tail = proc.ring_cq_tail;
    let user_snap = unsafe { snapshot_user_fields(ring_virt) };

    let pending = user_snap.sq_tail.distance(sq_head);
    if pending > MAX_SQ_ENTRIES {
        return Err(KernelError::InvalidParameter);
    }

    let cq_used = cq_tail.distance(user_snap.cq_head);
    let cq_available = MAX_CQ_ENTRIES.saturating_sub(cq_used);
    let to_process = pending.min(cq_available).min(MAX_RING_BATCH);

    if to_process == 0 {
        return Ok(0);
    }
    drop(ptable);

    let (completed, cq_write) = (0..to_process).fold((0u32, cq_tail), |(done, cq_w), i| {
        let sq_slot = sq_head.advance(i).slot(MAX_SQ_ENTRIES);
        let sqe = unsafe { read_sq_entry(ring_virt, sq_slot) };
        let cqe = process_submission(&sqe, pid);

        let cq_slot = cq_w.slot(MAX_CQ_ENTRIES);
        unsafe { write_cq_entry(ring_virt, cq_slot, cqe) };
        (done + 1, cq_w.advance(1))
    });

    let new_sq_head = sq_head.advance(to_process);
    unsafe { commit_header(ring_virt, new_sq_head, cq_write) };

    let mut ptable = PROCESSES.lock();
    if let Some(proc) = ptable.get_mut(pid) {
        proc.ring_sq_head = new_sq_head;
        proc.ring_cq_tail = cq_write;
    }

    Ok(completed as i64)
}