// Host model of the frozen AsyncBank and actual frozen Table/stack bytes.
// Derived from the frozen context-handoff fixture; no guest or assembly claim.
#[path = "../../os-context-switch/tests/support/mod.rs"]
#[allow(dead_code)]
pub(crate) mod saved;
use std::{
    cell::{Cell, RefCell},
    rc::Rc,
    vec::Vec,
};
use turkos_os_context_switch::context::*;
use turkos_os_threads::threads::{StackObservation, ThreadRecord};

#[derive(Default)]
pub(crate) struct Cpu {
    pub(crate) nonce: u64,
    pub(crate) registers: Option<Registers>,
    pub(crate) depart_fail: Option<FailurePhase>,
    pub(crate) arrive_fail: bool,
    pub(crate) departures: u64,
    pub(crate) arrivals: u64,
}
pub(crate) struct Device {
    pub(crate) inner: saved::Backend,
    pub(crate) cpu: Rc<RefCell<Cpu>>,
    pub(crate) root: Rc<Cell<u64>>,
}
unsafe impl Backend for Device {
    type Binding = saved::Binding;
    type Prepared = saved::Prepared;
    type IdlePrepared = saved::IdlePrepared;
    fn execution(&self) -> Execution {
        Execution {
            root: self.root.get(),
            ..self.inner.execution()
        }
    }
    fn record(&self, b: &Self::Binding) -> Result<ThreadRecord, Error> {
        self.inner.record(b)
    }
    fn stack(&self, b: &Self::Binding) -> Result<StackObservation, Error> {
        self.inner.stack(b)
    }
    fn root(&self, b: &Self::Binding) -> Result<u64, Error> {
        self.inner.root(b)
    }
    fn executable(&self, pc: u64) -> bool {
        self.inner.executable(pc)
    }
    fn initial_return_pc(&self, b: &Self::Binding) -> Result<u64, Error> {
        self.inner.initial_return_pc(b)
    }
    fn capture_matches(&self, b: &Self::Binding, c: &Capture) -> bool {
        self.inner.capture_matches(b, c)
    }
    fn registers(&self) -> Registers {
        self.cpu.borrow().registers.unwrap()
    }
    fn prepare(
        &mut self,
        s: Option<&Self::Binding>,
        t: &Self::Binding,
        i: &Image,
        d: Disposition,
    ) -> Result<Self::Prepared, Error> {
        self.inner.prepare(s, t, i, d)
    }
    fn commit(
        &mut self,
        _: &mut Self::Prepared,
        _: Option<&Self::Binding>,
        _: &Self::Binding,
        _: &Image,
        _: Disposition,
    ) -> Result<(), Failure> {
        panic!("synchronous commit must never span an asynchronous handoff")
    }
    fn finish(&mut self, p: Self::Prepared) -> Result<(), (Error, Self::Prepared)> {
        self.inner.finish(p)
    }
    fn abort(&mut self, p: Self::Prepared) -> Result<(), (Error, Self::Prepared)> {
        self.inner.abort(p)
    }
    fn release(&mut self, p: Self::Binding) -> Result<(), (Error, Self::Binding)> {
        self.inner.release(p)
    }
    fn ready_empty(&self) -> bool {
        self.inner.ready_empty()
    }
    fn wake_epoch(&self) -> u64 {
        self.inner.wake_epoch()
    }
    fn idle_prepare(&mut self, e: u64) -> Result<Self::IdlePrepared, Error> {
        self.inner.idle_prepare(e)
    }
    fn idle_commit(&mut self, p: &mut Self::IdlePrepared, e: u64) -> Result<IdleReceipt, Failure> {
        self.inner.idle_commit(p, e)
    }
    fn idle_finish(&mut self, p: Self::IdlePrepared) -> Result<(), (Error, Self::IdlePrepared)> {
        self.inner.idle_finish(p)
    }
    fn idle_abort(&mut self, p: Self::IdlePrepared) -> Result<(), (Error, Self::IdlePrepared)> {
        self.inner.idle_abort(p)
    }
}
unsafe impl HandoffBackend for Device {
    fn depart(
        &mut self,
        p: &mut Self::Prepared,
        source: Option<&Self::Binding>,
        _: u64,
        d: Disposition,
    ) -> Result<(), Failure> {
        assert!(p.nonce > 0);
        let failure = self.cpu.borrow().depart_fail;
        if failure == Some(FailurePhase::BeforeSwitch) {
            return Err(Failure {
                error: Error::Backend,
                phase: FailurePhase::BeforeSwitch,
            });
        }
        if let Some(source) = source {
            match d {
                Disposition::Ready => self
                    .inner
                    .table
                    .borrow_mut()
                    .yield_current(source.id)
                    .unwrap(),
                Disposition::Sleeping => {
                    self.inner.sleep_wait = Some(
                        self.inner
                            .table
                            .borrow_mut()
                            .sleep_current(source.id)
                            .unwrap(),
                    );
                }
            }
        }
        self.inner.hardware.borrow_mut().current = None;
        self.root.set(saved::ROOT);
        self.cpu.borrow_mut().departures += 1;
        if failure == Some(FailurePhase::AfterSwitch) {
            return Err(Failure {
                error: Error::Backend,
                phase: FailurePhase::AfterSwitch,
            });
        }
        Ok(())
    }
    fn arrive(
        &mut self,
        _: &mut Self::Prepared,
        target: &Self::Binding,
        _: &Image,
        nonce: u64,
    ) -> Result<(), Error> {
        let cpu = self.cpu.borrow();
        if cpu.nonce != nonce || cpu.registers.is_none() || cpu.arrive_fail {
            return Err(Error::Capture);
        }
        drop(cpu);
        self.inner
            .table
            .borrow_mut()
            .activate(target.id)
            .map_err(|_| Error::State)?;
        self.inner.hardware.borrow_mut().current = Some(target.id);
        self.cpu.borrow_mut().arrivals += 1;
        Ok(())
    }
}
pub(crate) type Setup = (
    AsyncBank<Device, 4>,
    Vec<Handle>,
    Rc<RefCell<saved::Hardware>>,
    Rc<RefCell<Cpu>>,
);
pub(crate) fn setup(count: usize) -> Setup {
    let (b, pins, h) = saved::backend(count);
    let cpu = Rc::new(RefCell::new(Cpu::default()));
    let mut bank = ContextBank::<_, 4>::new(Device {
        inner: b,
        cpu: cpu.clone(),
        root: Rc::new(Cell::new(BOOTSTRAP_ROOT)),
    })
    .ok()
    .unwrap();
    let handles = pins
        .into_iter()
        .map(|p| bank.attach(p).ok().unwrap())
        .collect();
    (AsyncBank::new(bank).ok().unwrap(), handles, h, cpu)
}

pub(crate) fn physical_step(bank: &AsyncBank<Device, 4>, token: &OwnedHandoff) {
    let image = bank.target_image(token).unwrap();
    bank.bank().with_backend_observation(|d| {
        let actual = saved::Backend::expected(image);
        d.inner.hardware.borrow_mut().sp = image.sp() - 64;
        let mut cpu = d.cpu.borrow_mut();
        cpu.nonce = token.nonce();
        cpu.registers = Some(actual);
    });
}
#[allow(dead_code)]
fn land(bank: &mut AsyncBank<Device, 4>, token: OwnedHandoff) {
    let token = bank.depart(token).ok().unwrap();
    physical_step(bank, &token);
    assert!(bank.landed(token).is_ok());
}
pub(crate) fn capture(
    bank: &AsyncBank<Device, 4>,
    source: Handle,
    h: &Rc<RefCell<saved::Hardware>>,
    value: u64,
    timer: bool,
) -> Capture {
    let r = bank.bank().record(source).unwrap();
    let image = if timer {
        let mut f = turkos_os_exceptions::exceptions::Frame::EMPTY;
        f.interrupted_sp = 0x80004000 + r.thread.slot() as u64 * 0x20000;
        f.elr = saved::ENTRY + 0x100;
        f.spsr = 0x305;
        f.vector_slot = 5;
        f.current_el = 4;
        f.canary = turkos_os_exceptions::exceptions::FRAME_CANARY;
        Image::Timer(f)
    } else {
        let mut f = Cooperative::EMPTY;
        f.sp = 0x80003000 + r.thread.slot() as u64 * 0x20000;
        f.x19_x30 = core::array::from_fn(|i| value + i as u64);
        f.x19_x30[11] = saved::ENTRY + 0x200 + value % 64 * 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 + 1;
        Image::Cooperative(f)
    };
    let episode = h.borrow().capture_episode + 1;
    let c = Capture {
        image,
        address: image.sp() - if timer { 848 } else { 640 },
        abi: Abi::COOPERATIVE,
        episode,
    };
    let mut bytes = Vec::new();
    let mut word = |v: u64| bytes.extend_from_slice(&v.to_le_bytes());
    match 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)
            }
        }
        _ => panic!(),
    }
    bank.bank().with_backend_observation(|d| {
        let record = d.inner.table.borrow().record(r.thread).unwrap();
        let offset = (c.address - record.layout.base) as usize;
        d.inner.stacks[r.thread.slot()].borrow_mut()[offset..offset + bytes.len()]
            .copy_from_slice(&bytes);
    });
    let mut hw = h.borrow_mut();
    hw.sp = c.address - 64;
    hw.capture = Some(c);
    hw.capture_episode = episode;
    hw.in_irq = timer;
    c
}

pub(crate) const BOOTSTRAP_ROOT: u64 = 0x44000000;
pub(crate) fn bootstrap() -> Registers {
    Registers {
        sp: 0x90004000,
        x19_x30: core::array::from_fn(|i| 0x1040 + i as u64),
        q: core::array::from_fn(|i| [0x104000 + i as u64, 0x204000 + i as u64]),
        fpcr: 0x400000,
        fpsr: 0x8000000,
        pc: saved::ENTRY + 0x900,
        argument: 0,
        timer: None,
    }
}
pub(crate) struct ReturnPrepared {
    pub(crate) source: turkos_os_threads::threads::ThreadId,
    pub(crate) expected: Registers,
    pub(crate) nonce: u64,
}
unsafe impl CompletionBackend for Device {
    type ReturnPrepared = ReturnPrepared;
    fn prepare_return(
        &mut self,
        source: &Self::Binding,
        _: &Capture,
    ) -> Result<ReturnPrepared, Error> {
        let mut h = self.inner.hardware.borrow_mut();
        if h.prepare_fail {
            return Err(Error::Backend);
        }
        h.prepares += 1;
        Ok(ReturnPrepared {
            source: source.id,
            expected: bootstrap(),
            nonce: h.prepares as u64,
        })
    }
    fn leave_return(
        &mut self,
        p: &mut ReturnPrepared,
        source: &Self::Binding,
    ) -> Result<(), Failure> {
        let fail = self.inner.hardware.borrow().commit_fail;
        if fail == Some(FailurePhase::BeforeSwitch) {
            return Err(Failure {
                error: Error::Backend,
                phase: FailurePhase::BeforeSwitch,
            });
        }
        assert_eq!(p.source, source.id);
        assert!(p.nonce > 0);
        if self.inner.hardware.borrow().bad_registers != 9 {
            self.inner
                .table
                .borrow_mut()
                .exit_current(source.id)
                .unwrap();
        }
        self.inner.hardware.borrow_mut().current = None;
        self.root.set(BOOTSTRAP_ROOT);
        if fail == Some(FailurePhase::AfterSwitch) {
            return Err(Failure {
                error: Error::Backend,
                phase: FailurePhase::AfterSwitch,
            });
        }
        Ok(())
    }
    fn returned(&self, p: &ReturnPrepared) -> Result<(), Error> {
        let h = self.inner.hardware.borrow();
        if h.current.is_some()
            || h.sp >= p.expected.sp
            || h.sp < p.expected.sp - 1024
            || self.cpu.borrow().registers != Some(p.expected)
        {
            return Err(Error::Capture);
        }
        Ok(())
    }
    fn return_root(&self, _: &ReturnPrepared) -> Result<u64, Error> {
        Ok(BOOTSTRAP_ROOT)
    }
    fn finish_return(&mut self, p: ReturnPrepared) -> Result<(), (Error, ReturnPrepared)> {
        let mut h = self.inner.hardware.borrow_mut();
        if h.finish_fail {
            return Err((Error::Backend, p));
        }
        h.finishes += 1;
        Ok(())
    }
    fn abort_return(&mut self, p: ReturnPrepared) -> Result<(), (Error, ReturnPrepared)> {
        let mut h = self.inner.hardware.borrow_mut();
        if h.abort_fail {
            return Err((Error::Backend, p));
        }
        h.aborts += 1;
        Ok(())
    }
}