// Copyright (c) 2026 taitep
// SPDX-License-Identifier: BSD-2-Clause
//
// This file is part of TRVE (https://git.taitep.se/trve.git)
// See LICENSE file in the project root for full license text.

use std::sync::atomic::{AtomicU32, AtomicU64};

use super::illegal;
use crate::{
    core::Core,
    decode::Instruction,
    exceptions::{Exception, ExceptionType},
    mem::{RAM_START, Ram},
};

pub(super) fn find_and_exec(instr: Instruction, core: &mut Core) -> Result<(), Exception> {
    match (instr.funct5(), instr.funct3()) {
        (0b00010, 0b010) if instr.rs2() == 0 => lr_w(core, instr),
        (0b00010, 0b011) if instr.rs2() == 0 => lr_d(core, instr),
        (0b00011, 0b010) => sc_w(core, instr),
        (0b00011, 0b011) => sc_d(core, instr),
        _ => illegal(instr),
    }
}

fn lr_d(core: &mut Core, instr: Instruction) -> Result<(), Exception> {
    core.reservation = None;

    let addr = core.reg_read(instr.rs1());
    if !addr.is_multiple_of(8) {
        return Err(Exception {
            type_: ExceptionType::LoadAddressMisaligned,
            value: addr,
        });
    }
    if addr < RAM_START {
        return Err(Exception {
            type_: ExceptionType::LoadAccessFault,
            value: addr,
        });
    }
    let ram_addr = addr - RAM_START;

    let reservation_data = core
        .mem
        .ram
        .wait_for_even_version(ram_addr)
        .ok_or_else(|| Exception {
            type_: ExceptionType::LoadAccessFault,
            value: addr,
        })?;

    core.reg_write(instr.rd(), unsafe {
        let index = ram_addr as usize / 8;
        core.mem
            .ram
            .buf_transmuted::<AtomicU64>()
            .get(index)
            .ok_or_else(|| Exception {
                type_: ExceptionType::LoadAccessFault,
                value: addr,
            })?
            .load(instr.amo_ordering())
    });

    core.reservation = Some(reservation_data);
    core.advance_pc();

    Ok(())
}
fn sc_d(core: &mut Core, instr: Instruction) -> Result<(), Exception> {
    let res = if let Some((reserved_chunk_id, reserved_version)) = core.reservation {
        let addr = core.reg_read(instr.rs1());
        if !addr.is_multiple_of(8) {
            return Err(Exception {
                type_: ExceptionType::StoreAmoAddressMisaligned,
                value: addr,
            });
        }
        if addr < RAM_START {
            return Err(Exception {
                type_: ExceptionType::StoreAmoAccessFault,
                value: addr,
            });
        }
        let ram_addr = addr - RAM_START;
        let chunk_id = ram_addr as usize / Ram::VERSION_CHUNK_SIZE;

        if chunk_id != reserved_chunk_id {
            // Mismatched reservation location and address
            core.reg_write(instr.rd(), 1);
            core.reservation = None;
            return Ok(());
        }

        let claim_res = core
            .mem
            .ram
            .claim_expected(chunk_id, reserved_version)
            .ok_or_else(|| Exception {
                type_: ExceptionType::StoreAmoAccessFault,
                value: addr,
            })?;

        if claim_res.is_some() {
            core.reservation = None;
            let value = core.reg_read(instr.rs2());
            unsafe {
                let index = ram_addr as usize / 8;
                core.mem
                    .ram
                    .buf_transmuted::<AtomicU64>()
                    .get(index)
                    .ok_or_else(|| Exception {
                        type_: ExceptionType::StoreAmoAccessFault,
                        value: addr,
                    })?
                    .store(value, instr.amo_ordering());

                Ok(0)
            }
        } else {
            core.reservation = None;
            Ok(1)
        }
    } else {
        core.reservation = None;
        Ok(1)
    }
    .map(|s| core.reg_write(instr.rd(), s));

    core.advance_pc();

    res
}

fn lr_w(core: &mut Core, instr: Instruction) -> Result<(), Exception> {
    core.reservation = None;

    let addr = core.reg_read(instr.rs1());
    if !addr.is_multiple_of(4) {
        return Err(Exception {
            type_: ExceptionType::LoadAddressMisaligned,
            value: addr,
        });
    }
    if addr < RAM_START {
        return Err(Exception {
            type_: ExceptionType::LoadAccessFault,
            value: addr,
        });
    }
    let ram_addr = addr - RAM_START;

    let reservation_data = core
        .mem
        .ram
        .wait_for_even_version(ram_addr)
        .ok_or_else(|| Exception {
            type_: ExceptionType::LoadAccessFault,
            value: addr,
        })?;

    core.reg_write(instr.rd(), unsafe {
        let index = ram_addr as usize / 4;
        core.mem
            .ram
            .buf_transmuted::<AtomicU32>()
            .get(index)
            .ok_or_else(|| Exception {
                type_: ExceptionType::LoadAccessFault,
                value: addr,
            })?
            .load(instr.amo_ordering())
    } as i32 as i64 as u64);

    core.reservation = Some(reservation_data);
    core.advance_pc();

    Ok(())
}
fn sc_w(core: &mut Core, instr: Instruction) -> Result<(), Exception> {
    let res = if let Some((reserved_chunk_id, reserved_version)) = core.reservation {
        let addr = core.reg_read(instr.rs1());
        if !addr.is_multiple_of(4) {
            return Err(Exception {
                type_: ExceptionType::StoreAmoAddressMisaligned,
                value: addr,
            });
        }
        if addr < RAM_START {
            return Err(Exception {
                type_: ExceptionType::StoreAmoAccessFault,
                value: addr,
            });
        }
        let ram_addr = addr - RAM_START;
        let chunk_id = ram_addr as usize / Ram::VERSION_CHUNK_SIZE;

        if chunk_id != reserved_chunk_id {
            // Mismatched reservation location and address
            core.reg_write(instr.rd(), 1);
            core.reservation = None;
            return Ok(());
        }

        let claim_res = core
            .mem
            .ram
            .claim_expected(chunk_id, reserved_version)
            .ok_or_else(|| Exception {
                type_: ExceptionType::StoreAmoAccessFault,
                value: addr,
            })?;

        if claim_res.is_some() {
            core.reservation = None;
            let value = core.reg_read(instr.rs2());
            unsafe {
                let index = ram_addr as usize / 4;
                core.mem
                    .ram
                    .buf_transmuted::<AtomicU32>()
                    .get(index)
                    .ok_or_else(|| Exception {
                        type_: ExceptionType::StoreAmoAccessFault,
                        value: addr,
                    })?
                    .store(value as u32, instr.amo_ordering());

                Ok(0)
            }
        } else {
            core.reservation = None;
            Ok(1)
        }
    } else {
        core.reservation = None;
        Ok(1)
    }
    .map(|s| core.reg_write(instr.rd(), s));

    core.advance_pc();

    res
}