// Isolated copy of the frozen Context host fixture; only the retained Wait
// storage uses RefCell for this backend's serialized actual Table wake hook.
use std::{cell::RefCell, rc::Rc};
use std::{vec, vec::Vec};
use turkos_os_context_switch::context::*;
use turkos_os_exceptions::exceptions::{FRAME_BYTES, FRAME_CANARY, Frame};
use turkos_os_threads::threads::{
    self as thread, StackLayout, StackObservation, Table, ThreadId, ThreadRecord,
};
pub(crate) const ROOT: u64 = 0x41000000;
pub(crate) const ENTRY: u64 = 0x40080000;
pub(crate) const RETURN: u64 = 0x40081000;
pub(crate) struct Stack {
    pub(crate) layout: StackLayout,
    pub(crate) bytes: Rc<RefCell<Vec<u8>>>,
}
#[derive(Default)]
pub(crate) struct Hardware {
    pub(crate) cpu: u32,
    pub(crate) daif: u64,
    pub(crate) in_irq: bool,
    pub(crate) sp: u64,
    pub(crate) current: Option<ThreadId>,
    pub(crate) critical: u32,
    pub(crate) safe: bool,
    pub(crate) fp: bool,
    pub(crate) corrupt_guard: bool,
    pub(crate) corrupt_canary: bool,
    pub(crate) capture: Option<Capture>,
    pub(crate) capture_episode: u64,
    pub(crate) prepare_fail: bool,
    pub(crate) commit_fail: Option<FailurePhase>,
    pub(crate) bad_registers: u8,
    pub(crate) finish_fail: bool,
    pub(crate) abort_fail: bool,
    pub(crate) release_fail: bool,
    pub(crate) prepares: usize,
    pub(crate) commits: usize,
    pub(crate) finishes: usize,
    pub(crate) aborts: usize,
    pub(crate) releases: usize,
    pub(crate) pins: usize,
    pub(crate) stack_pages: u64,
    pub(crate) wake_epoch: u64,
    pub(crate) work: bool,
    pub(crate) boundary_wake: bool,
    pub(crate) wfi: usize,
    pub(crate) idle_fail: Option<FailurePhase>,
    pub(crate) idle_receipt_bad: u8,
}
pub(crate) struct ThreadBackend {
    pub(crate) hardware: Rc<RefCell<Hardware>>,
}
unsafe impl thread::Backend for ThreadBackend {
    type Stack = Stack;
    fn execution(&self) -> thread::Execution {
        let h = self.hardware.borrow();
        thread::Execution {
            cpu: h.cpu,
            irq_masked: h.daif & 128 != 0,
            in_irq: h.in_irq,
            sp: h.sp,
        }
    }
    fn layout(&self, s: &Stack) -> Result<StackLayout, thread::Error> {
        Ok(s.layout)
    }
    fn entry_executable(&self, pc: u64) -> bool {
        (ENTRY..ENTRY + 0x10000).contains(&pc)
    }
    fn prepare(&mut self, s: &mut Stack) -> Result<(), thread::Error> {
        let mut b = s.bytes.borrow_mut();
        b.fill(0);
        b[..8].copy_from_slice(&thread::LOW_CANARY.to_le_bytes());
        let n = b.len();
        b[n - 8..].copy_from_slice(&thread::HIGH_CANARY.to_le_bytes());
        Ok(())
    }
    fn observe(&self, s: &Stack) -> Result<StackObservation, thread::Error> {
        let b = s.bytes.borrow();
        let h = self.hardware.borrow();
        Ok(StackObservation {
            layout: s.layout,
            guard_descriptor: u64::from(h.corrupt_guard),
            low_canary: u64::from_le_bytes(b[..8].try_into().unwrap())
                ^ u64::from(h.corrupt_canary),
            high_canary: u64::from_le_bytes(b[b.len() - 8..].try_into().unwrap()),
            initially_zeroed: b[8..b.len() - 8].iter().all(|v| *v == 0),
        })
    }
    fn release(&mut self, s: Stack) -> Result<(), (thread::Error, Stack)> {
        let mut h = self.hardware.borrow_mut();
        s.bytes.borrow_mut().fill(0);
        h.stack_pages -= s.layout.bytes / 4096;
        Ok(())
    }
}
pub(crate) struct Binding {
    pub(crate) id: ThreadId,
    pub(crate) nonce: u64,
}
pub(crate) struct Prepared {
    pub(crate) nonce: u64,
}
pub(crate) struct IdlePrepared {
    pub(crate) epoch: u64,
    pub(crate) wait: Option<thread::Wait>,
}
pub(crate) struct Backend {
    pub(crate) table: RefCell<Table<ThreadBackend, 4>>,
    pub(crate) hardware: Rc<RefCell<Hardware>>,
    pub(crate) registers: Registers,
    pub(crate) stacks: Vec<Rc<RefCell<Vec<u8>>>>,
    pub(crate) sleep_wait: RefCell<Option<thread::Wait>>,
}
impl Backend {
    fn map(_: thread::Error) -> Error {
        Error::State
    }
    pub(crate) fn expected(image: Image) -> Registers {
        let (q, fpcr, fpsr) = image.simd();
        let (mut x, mut argument, mut timer) = ([0; 12], 0, None);
        match image {
            Image::Initial {
                argument: a,
                return_pc,
                ..
            } => {
                argument = a;
                x[11] = return_pc;
            }
            Image::Cooperative(c) => x = c.x19_x30,
            Image::Timer(f) => {
                x.copy_from_slice(&f.x[19..31]);
                timer = Some(f);
            }
        }
        Registers {
            sp: image.sp(),
            x19_x30: x,
            q,
            fpcr,
            fpsr,
            pc: image.pc(),
            argument,
            timer,
        }
    }
}
unsafe impl turkos_os_context_switch::context::Backend for Backend {
    type Binding = Binding;
    type Prepared = Prepared;
    type IdlePrepared = IdlePrepared;
    fn execution(&self) -> Execution {
        let h = self.hardware.borrow();
        Execution {
            cpu: h.cpu,
            daif: h.daif,
            in_irq: h.in_irq,
            sp: h.sp,
            root: ROOT,
            thread: h.current,
            fp_enabled: h.fp,
            critical_depth: h.critical,
            safe_point: h.safe,
            tco: 0,
        }
    }
    fn record(&self, b: &Binding) -> Result<ThreadRecord, Error> {
        if b.nonce != b.id.generation() {
            return Err(Error::Stale);
        }
        self.table.borrow().record(b.id).map_err(Self::map)
    }
    fn stack(&self, b: &Binding) -> Result<StackObservation, Error> {
        let r = self.record(b)?;
        let h = self.hardware.borrow();
        let bytes = self.stacks[b.id.slot()].borrow();
        Ok(StackObservation {
            layout: r.layout,
            guard_descriptor: u64::from(h.corrupt_guard),
            low_canary: u64::from_le_bytes(bytes[..8].try_into().unwrap())
                ^ u64::from(h.corrupt_canary),
            high_canary: u64::from_le_bytes(bytes[bytes.len() - 8..].try_into().unwrap()),
            initially_zeroed: false,
        })
    }
    fn root(&self, _: &Binding) -> Result<u64, Error> {
        Ok(ROOT)
    }
    fn executable(&self, pc: u64) -> bool {
        (ENTRY..ENTRY + 0x10000).contains(&pc)
    }
    fn initial_return_pc(&self, _: &Binding) -> Result<u64, Error> {
        Ok(RETURN)
    }
    fn capture_matches(&self, binding: &Binding, c: &Capture) -> bool {
        let h = self.hardware.borrow();
        let record = match self.record(binding) {
            Ok(record) => record,
            Err(_) => return false,
        };
        let bytes = self.stacks[binding.id.slot()].borrow();
        let encoded = capture_bytes(c);
        let offset = match c.address.checked_sub(record.layout.base) {
            Some(offset) => offset as usize,
            None => return false,
        };
        h.current == Some(binding.id)
            && h.capture_episode == c.episode
            && h.capture == Some(*c)
            && bytes.get(offset..offset + encoded.len()) == Some(encoded.as_slice())
    }
    fn registers(&self) -> Registers {
        self.registers
    }
    fn prepare(
        &mut self,
        _: Option<&Binding>,
        _: &Binding,
        _: &Image,
        _: Disposition,
    ) -> Result<Prepared, Error> {
        let mut h = self.hardware.borrow_mut();
        if h.prepare_fail {
            return Err(Error::Backend);
        }
        h.prepares += 1;
        Ok(Prepared {
            nonce: h.prepares as u64,
        })
    }
    fn commit(
        &mut self,
        p: &mut Prepared,
        source: Option<&Binding>,
        target: &Binding,
        image: &Image,
        disposition: Disposition,
    ) -> Result<(), Failure> {
        assert!(p.nonce > 0);
        let fail = self.hardware.borrow().commit_fail;
        if fail == Some(FailurePhase::BeforeSwitch) {
            return Err(Failure {
                error: Error::Backend,
                phase: FailurePhase::BeforeSwitch,
            });
        }
        {
            let mut h = self.hardware.borrow_mut();
            h.in_irq = false;
            h.commits += 1;
        }
        if let Some(s) = source {
            if disposition == Disposition::Sleeping {
                *self.sleep_wait.borrow_mut() = Some(
                    self.table
                        .borrow_mut()
                        .sleep_current(s.id)
                        .map_err(|_| Failure {
                            error: Error::Backend,
                            phase: FailurePhase::AfterSwitch,
                        })?,
                );
            } else {
                self.table
                    .borrow_mut()
                    .yield_current(s.id)
                    .map_err(|_| Failure {
                        error: Error::Backend,
                        phase: FailurePhase::AfterSwitch,
                    })?;
            }
        }
        {
            let mut h = self.hardware.borrow_mut();
            h.sp = image.sp();
            h.current = Some(target.id);
        }
        self.table
            .borrow_mut()
            .activate(target.id)
            .map_err(|_| Failure {
                error: Error::Backend,
                phase: FailurePhase::AfterSwitch,
            })?;
        self.registers = Self::expected(*image);
        let bad = self.hardware.borrow().bad_registers;
        match bad {
            1 => self.registers.q[31][1] ^= 1,
            2 => self.registers.x19_x30[10] ^= 1,
            3 => self.registers.fpcr ^= 1,
            4 => {
                if let Some(f) = &mut self.registers.timer {
                    f.esr ^= 1;
                }
            }
            5 => self.registers.fpsr ^= 1,
            6 => self.registers.pc ^= 4,
            7 => self.registers.sp ^= 16,
            _ => {}
        }
        if fail == Some(FailurePhase::AfterSwitch) {
            return Err(Failure {
                error: Error::Backend,
                phase: FailurePhase::AfterSwitch,
            });
        }
        Ok(())
    }
    fn finish(&mut self, p: Prepared) -> Result<(), (Error, Prepared)> {
        let mut h = self.hardware.borrow_mut();
        if h.finish_fail {
            return Err((Error::Backend, p));
        }
        h.finishes += 1;
        Ok(())
    }
    fn abort(&mut self, p: Prepared) -> Result<(), (Error, Prepared)> {
        let mut h = self.hardware.borrow_mut();
        if h.abort_fail {
            return Err((Error::Backend, p));
        }
        h.aborts += 1;
        Ok(())
    }
    fn release(&mut self, b: Binding) -> Result<(), (Error, Binding)> {
        let mut h = self.hardware.borrow_mut();
        if h.release_fail {
            return Err((Error::Backend, b));
        }
        h.pins -= 1;
        h.releases += 1;
        Ok(())
    }
    fn ready_empty(&self) -> bool {
        !self.hardware.borrow().work && self.table.borrow().snapshot().ready == 0
    }
    fn wake_epoch(&self) -> u64 {
        self.hardware.borrow().wake_epoch
    }
    fn idle_prepare(&mut self, epoch: u64) -> Result<IdlePrepared, Error> {
        Ok(IdlePrepared {
            epoch,
            wait: self.sleep_wait.borrow_mut().take(),
        })
    }
    fn idle_commit(&mut self, p: &mut IdlePrepared, epoch: u64) -> Result<IdleReceipt, Failure> {
        assert_eq!(p.epoch, epoch);
        if self.hardware.borrow().idle_fail == Some(FailurePhase::BeforeSwitch) {
            return Err(Failure {
                error: Error::Backend,
                phase: FailurePhase::BeforeSwitch,
            });
        }
        let boundary = self.hardware.borrow().boundary_wake;
        if boundary {
            let mut h = self.hardware.borrow_mut();
            h.wake_epoch += 1;
            h.work = true;
            h.boundary_wake = false;
        }
        if boundary && let Some(wait) = p.wait.take() {
            match self.table.borrow_mut().wake(wait) {
                Ok(_) => {}
                Err((_, wait)) => {
                    p.wait = Some(wait);
                    return Err(Failure {
                        error: Error::Backend,
                        phase: FailurePhase::AfterSwitch,
                    });
                }
            }
        }
        let checked = self.hardware.borrow().wake_epoch;
        let work = self.hardware.borrow().work || self.table.borrow().snapshot().ready > 0;
        let executed = !work && checked == epoch;
        if executed {
            {
                let mut h = self.hardware.borrow_mut();
                h.wfi += 1;
                h.wake_epoch += 1;
                h.work = true;
            }
            if let Some(wait) = p.wait.take() {
                match self.table.borrow_mut().wake(wait) {
                    Ok(_) => {}
                    Err((_, wait)) => {
                        p.wait = Some(wait);
                        return Err(Failure {
                            error: Error::Backend,
                            phase: FailurePhase::AfterSwitch,
                        });
                    }
                }
            }
        }
        let h = self.hardware.borrow();
        let mut r = IdleReceipt {
            checked_epoch: checked,
            work_at_handoff: work,
            executed_wfi: executed,
            wake_epoch: h.wake_epoch,
        };
        match h.idle_receipt_bad {
            1 => r.checked_epoch = epoch - 1,
            2 => r.wake_epoch = checked - 1,
            3 => r.executed_wfi = true,
            4 => {
                r.checked_epoch = epoch + 1;
                r.wake_epoch = epoch + 1;
                r.work_at_handoff = false;
                r.executed_wfi = true;
            }
            _ => {}
        }
        if h.idle_fail == Some(FailurePhase::AfterSwitch) {
            return Err(Failure {
                error: Error::Backend,
                phase: FailurePhase::AfterSwitch,
            });
        }
        Ok(r)
    }
    fn idle_finish(&mut self, mut p: IdlePrepared) -> Result<(), (Error, IdlePrepared)> {
        if self.hardware.borrow().finish_fail {
            Err((Error::Backend, p))
        } else {
            assert!(self.sleep_wait.borrow().is_none());
            *self.sleep_wait.borrow_mut() = p.wait.take();
            Ok(())
        }
    }
    fn idle_abort(&mut self, mut p: IdlePrepared) -> Result<(), (Error, IdlePrepared)> {
        if self.hardware.borrow().abort_fail {
            Err((Error::Backend, p))
        } else {
            assert!(self.sleep_wait.borrow().is_none());
            *self.sleep_wait.borrow_mut() = p.wait.take();
            Ok(())
        }
    }
}
pub(crate) fn backend(count: usize) -> (Backend, Vec<Binding>, Rc<RefCell<Hardware>>) {
    let hardware = Rc::new(RefCell::new(Hardware {
        daif: 0x3c0,
        sp: 0x90001000,
        safe: true,
        fp: true,
        wake_epoch: 1,
        ..Hardware::default()
    }));
    let mut table = Table::new(ThreadBackend {
        hardware: hardware.clone(),
    })
    .ok()
    .unwrap();
    let mut bindings = Vec::new();
    let mut stacks = Vec::new();
    for i in 0..count {
        let bytes = Rc::new(RefCell::new(vec![0xa5; 16384]));
        let stack = Stack {
            layout: StackLayout {
                base: 0x80001000 + i as u64 * 0x20000,
                physical: 0x42000000 + i as u64 * 0x20000,
                bytes: 16384,
            },
            bytes: bytes.clone(),
        };
        let id = table
            .create(stack, ENTRY + i as u64 * 4, i as u64 + 11)
            .ok()
            .unwrap();
        bindings.push(Binding {
            id,
            nonce: id.generation(),
        });
        stacks.push(bytes);
        hardware.borrow_mut().pins += 1;
        hardware.borrow_mut().stack_pages += 4;
    }
    let registers = Backend::expected(Image::Initial {
        entry: ENTRY,
        argument: 0,
        sp: 0x90001000,
        return_pc: RETURN,
    });
    (
        Backend {
            table: RefCell::new(table),
            hardware: hardware.clone(),
            registers,
            stacks,
            sleep_wait: RefCell::new(None),
        },
        bindings,
        hardware,
    )
}
pub(crate) fn setup() -> (ContextBank<Backend, 4>, [Handle; 2], Rc<RefCell<Hardware>>) {
    let (b, mut pins, h) = backend(2);
    let mut bank = ContextBank::new(b).ok().unwrap();
    let a = bank.attach(pins.remove(0)).ok().unwrap();
    let c = bank.attach(pins.remove(0)).ok().unwrap();
    (bank, [a, c], h)
}
pub(crate) fn finish<const N: usize>(d: Dispatch<'_, Backend, N>) {
    assert!(d.complete().is_ok());
}
pub(crate) fn capture(
    bank: &ContextBank<Backend, 4>,
    handle: Handle,
    h: &Rc<RefCell<Hardware>>,
    value: u64,
    timer: bool,
) -> Capture {
    let r = bank.with_backend_observation(|b| {
        b.table
            .borrow()
            .record(bank.record(handle).unwrap().thread)
            .unwrap()
    });
    let sp = r.layout.usable_top() - 1024 * (1 + value % 4);
    let episode = h.borrow().capture_episode.checked_add(1).unwrap();
    let c = if timer {
        let mut f = Frame::EMPTY;
        f.x = core::array::from_fn(|i| value + i as u64);
        f.interrupted_sp = sp;
        f.sp_el0 = 0x80004000;
        f.elr = ENTRY + 0x200 + value % 256 * 4;
        f.spsr = 0x305;
        f.esr = 0x123456;
        f.far = 0xabcdef;
        f.vector_slot = 5;
        f.q = core::array::from_fn(|i| [value * 100 + i as u64, value * 200 + i as u64]);
        f.fpcr = value as u32;
        f.fpsr = value as u32 + 3;
        f.current_el = 4;
        f.canary = FRAME_CANARY;
        Capture {
            image: Image::Timer(f),
            address: sp - FRAME_BYTES as u64,
            abi: Abi::COOPERATIVE,
            episode,
        }
    } else {
        let mut f = Cooperative::EMPTY;
        f.sp = sp;
        f.x19_x30 = core::array::from_fn(|i| value + i as u64);
        f.x19_x30[11] = ENTRY + 0x100 + value % 256 * 4;
        f.q = core::array::from_fn(|i| [value * 100 + i as u64, value * 200 + i as u64]);
        f.fpcr = value as u32;
        f.fpsr = value as u32 + 3;
        Capture {
            image: Image::Cooperative(f),
            address: sp - COOP_BYTES as u64,
            abi: Abi::COOPERATIVE,
            episode,
        }
    };
    let mut hw = h.borrow_mut();
    hw.sp = sp - COOP_BYTES as u64;
    hw.capture = Some(c);
    hw.capture_episode = episode;
    hw.in_irq = timer;
    drop(hw);
    install_capture(bank, handle, h, c);
    c
}

fn capture_bytes(c: &Capture) -> Vec<u8> {
    let mut bytes = Vec::new();
    let mut word = |v: u64| bytes.extend_from_slice(&v.to_le_bytes());
    match c.image {
        Image::Cooperative(f) => {
            word(f.sp);
            for v in f.x19_x30 {
                word(v);
            }
            word(f.padding);
            for q in f.q {
                for v in q {
                    word(v);
                }
            }
            word(u64::from(f.fpcr) | (u64::from(f.fpsr) << 32));
            word(f.canary);
        }
        Image::Timer(f) => {
            for v in f.x {
                word(v);
            }
            for v in [
                f.interrupted_sp,
                f.sp_el0,
                f.elr,
                f.spsr,
                f.esr,
                f.far,
                f.vector_slot,
            ] {
                word(v);
            }
            for q in f.q {
                for v in q {
                    word(v);
                }
            }
            word(u64::from(f.fpcr) | (u64::from(f.fpsr) << 32));
            for v in [f.current_el, f.canary, f.padding] {
                word(v);
            }
        }
        Image::Initial { .. } => panic!("an initial image is not a physical capture"),
    }
    bytes
}
pub(crate) fn install_capture(
    bank: &ContextBank<Backend, 4>,
    handle: Handle,
    h: &Rc<RefCell<Hardware>>,
    c: Capture,
) {
    bank.with_backend_observation(|b| {
        let record = b
            .table
            .borrow()
            .record(bank.record(handle).unwrap().thread)
            .unwrap();
        let offset = (c.address - record.layout.base) as usize;
        let encoded = capture_bytes(&c);
        b.stacks[record.id.slot()].borrow_mut()[offset..offset + encoded.len()]
            .copy_from_slice(&encoded);
    });
    let mut actual = h.borrow_mut();
    actual.capture = Some(c);
    actual.capture_episode = c.episode;
}

pub(crate) fn setup_idle() -> (ContextBank<Backend, 4>, [Handle; 2], Rc<RefCell<Hardware>>) {
    let (mut bank, [a, b], h) = setup();
    finish(bank.start(a).unwrap());
    let c = capture(&bank, a, &h, 17, false);
    finish(bank.cooperate(a, b, c).unwrap());
    let c = capture(&bank, b, &h, 43, false);
    finish(bank.sleep_to(b, a, c).unwrap());
    assert_eq!(
        bank.with_backend_observation(|back| back.table.borrow().snapshot().sleeping),
        1
    );
    (bank, [a, b], h)
}