use crate::{Error, Task, spin::Context};

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Observation {
    pub cpu: u16,
    pub task: Option<Task>,
    pub in_irq: bool,
    pub daif: u64,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Ticket {
    pub scope: u64,
    pub task: Task,
    pub generation: u64,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Failure {
    pub error: Error,
    pub quarantined: bool,
}

/// Exclusive scheduler access, serialized against all task transitions.
///
/// # Safety
/// Observation and task validation must describe actual live tasks. Prepare
/// must acquire every scheduler resource without publishing a partial park or
/// wake. A failed prepare retains any acquired resources in this backend and
/// reports `quarantined`. Commit cannot fail and publishes the exact prepared
/// task before that task can execute again. The scope identity is stable and
/// globally unique while any token from this backend can remain alive.
/// Cancellation must validate its caller's authority and retire the exact park
/// ticket; it may wake only a still-live matching task generation.
pub unsafe trait Backend {
    type Park;
    type Wake;
    type Cancel;
    fn identity(&self) -> u64;
    fn observation(&self) -> Observation;
    fn validate_task(&self, task: Task) -> bool;
    fn prepare_park(&mut self, ticket: Ticket) -> Result<Self::Park, Failure>;
    fn commit_park(&mut self, prepared: Self::Park);
    fn prepare_wake(&mut self, ticket: Ticket) -> Result<Self::Wake, Failure>;
    fn commit_wake(&mut self, prepared: Self::Wake);
    fn authorize_cancel(&self, caller: Task, waiter: Task) -> bool;
    fn prepare_cancel(&mut self, ticket: Ticket) -> Result<Self::Cancel, Failure>;
    fn commit_cancel(&mut self, prepared: Self::Cancel);
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum State {
    Vacant,
    Queued,
    Granted,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct Slot {
    state: State,
    task: Task,
    generation: u64,
}
const EMPTY: Slot = Slot {
    state: State::Vacant,
    task: Task {
        id: 0,
        generation: 0,
    },
    generation: 0,
};
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
struct Owner {
    task: Task,
    generation: u64,
    wait_slot: Option<usize>,
    claimed: bool,
}

/// Nonclone ownership of the mutex, independent of the current CPU.
#[derive(Debug)]
pub struct Lease {
    scope: u64,
    task: Task,
    generation: u64,
}
impl Lease {
    pub fn task(&self) -> Task {
        self.task
    }
    pub fn generation(&self) -> u64 {
        self.generation
    }
}
/// Nonclone ownership of one bounded waiter record.
#[derive(Debug)]
pub struct Wait {
    scope: u64,
    task: Task,
    slot: usize,
    generation: u64,
}
impl Wait {
    pub fn task(&self) -> Task {
        self.task
    }
    pub fn generation(&self) -> u64 {
        self.generation
    }
}
#[derive(Debug)]
pub enum Request {
    Owned(Lease),
    Waiting(Wait),
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct WaitRecord {
    pub task: Task,
    pub generation: u64,
    pub granted: bool,
}

pub struct Mutex<B: Backend, const N: usize, T = ()> {
    backend: B,
    value: T,
    scope: u64,
    owner: Option<Owner>,
    slots: [Slot; N],
    queue: [usize; N],
    len: usize,
    epoch: u64,
    faulted: bool,
}

impl<B: Backend, const N: usize> Mutex<B, N> {
    pub fn new(backend: B) -> Result<Self, (Error, B)> {
        Self::with_value(backend, ()).map_err(|(error, backend, ())| (error, backend))
    }
    pub fn into_backend(self) -> Result<B, (Error, Self)> {
        self.into_resources().map(|(backend, ())| backend)
    }
}
impl<B: Backend, const N: usize, T> Mutex<B, N, T> {
    pub fn with_value(backend: B, value: T) -> Result<Self, (Error, B, T)> {
        if N == 0 || backend.identity() == 0 {
            return Err((Error::Invalid, backend, value));
        }
        let scope = backend.identity();
        Ok(Self {
            backend,
            value,
            scope,
            owner: None,
            slots: [EMPTY; N],
            queue: [0; N],
            len: 0,
            epoch: 0,
            faulted: false,
        })
    }
    fn current(&self, ctx: &Context) -> Result<Task, Error> {
        if self.faulted {
            return Err(Error::Quarantined);
        }
        if self.backend.identity() != self.scope {
            return Err(Error::Foreign);
        }
        let observed = self.backend.observation();
        if observed.cpu != ctx.cpu() || observed.task != ctx.task() || observed.in_irq {
            return Err(Error::Context);
        }
        let task = ctx.may_sleep(observed.daif)?;
        if !self.backend.validate_task(task) {
            return Err(Error::Stale);
        }
        Ok(task)
    }
    fn failed(&mut self, failure: Failure) -> Error {
        if failure.quarantined {
            self.faulted = true;
        }
        failure.error
    }
    pub fn validate_lease(&self, ctx: &Context, lease: &Lease) -> Result<(), Error> {
        let task = self.current(ctx)?;
        if lease.scope != self.scope {
            return Err(Error::Foreign);
        }
        let owner = self.owner.ok_or(Error::Stale)?;
        if owner.task != lease.task || owner.generation != lease.generation || !owner.claimed {
            return Err(Error::Stale);
        }
        if task != lease.task {
            return Err(Error::Foreign);
        }
        Ok(())
    }
    /// Only the current live task with this mutex's actual private lease can
    /// access the protected payload. No backend mutation escapes this method.
    pub fn with_lease<R>(
        &mut self,
        ctx: &Context,
        lease: &Lease,
        access: impl FnOnce(&mut T) -> R,
    ) -> Result<R, Error> {
        self.validate_lease(ctx, lease)?;
        Ok(access(&mut self.value))
    }
    pub fn request(&mut self, ctx: &Context) -> Result<Request, Error> {
        let task = self.current(ctx)?;
        if self.owner.is_some_and(|owner| owner.task == task) {
            return Err(Error::Recursive);
        }
        if self
            .slots
            .iter()
            .any(|slot| slot.state != State::Vacant && slot.task == task)
        {
            return Err(Error::Busy);
        }
        let generation = self.epoch.checked_add(1).ok_or(Error::Overflow)?;
        if self.owner.is_none() {
            if self.len != 0 {
                return Err(Error::Quarantined);
            }
            self.owner = Some(Owner {
                task,
                generation,
                wait_slot: None,
                claimed: true,
            });
            self.epoch = generation;
            return Ok(Request::Owned(Lease {
                scope: self.scope,
                task,
                generation,
            }));
        }
        let index = self
            .slots
            .iter()
            .position(|slot| slot.state == State::Vacant && slot.generation != u64::MAX)
            .ok_or(Error::Capacity)?;
        if self.len == N {
            return Err(Error::Capacity);
        }
        let ticket = Ticket {
            scope: self.scope,
            task,
            generation,
        };
        let prepared = match self.backend.prepare_park(ticket) {
            Ok(v) => v,
            Err(f) => return Err(self.failed(f)),
        };
        self.slots[index] = Slot {
            state: State::Queued,
            task,
            generation,
        };
        self.queue[self.len] = index;
        self.len += 1;
        self.epoch = generation;
        self.backend.commit_park(prepared);
        Ok(Request::Waiting(Wait {
            scope: self.scope,
            task,
            slot: index,
            generation,
        }))
    }
    fn validate_wait(&self, wait: &Wait) -> Result<Slot, Error> {
        if wait.scope != self.scope {
            return Err(Error::Foreign);
        }
        let slot = *self.slots.get(wait.slot).ok_or(Error::Stale)?;
        if slot.state == State::Vacant
            || slot.task != wait.task
            || slot.generation != wait.generation
        {
            return Err(Error::Stale);
        }
        Ok(slot)
    }
    pub fn claim(&mut self, ctx: &Context, wait: Wait) -> Result<Lease, (Error, Wait)> {
        let check = (|| {
            let task = self.current(ctx)?;
            let slot = self.validate_wait(&wait)?;
            if task != wait.task {
                return Err(Error::Foreign);
            }
            if slot.state != State::Granted {
                return Err(Error::Busy);
            }
            let owner = self.owner.ok_or(Error::Stale)?;
            if owner.task != task
                || owner.generation != wait.generation
                || owner.wait_slot != Some(wait.slot)
                || owner.claimed
            {
                return Err(Error::Stale);
            }
            Ok(owner)
        })();
        match check {
            Err(error) => Err((error, wait)),
            Ok(mut owner) => {
                owner.claimed = true;
                owner.wait_slot = None;
                self.owner = Some(owner);
                self.slots[wait.slot].state = State::Vacant;
                Ok(Lease {
                    scope: self.scope,
                    task: wait.task,
                    generation: wait.generation,
                })
            }
        }
    }
    pub fn release(&mut self, ctx: &Context, lease: Lease) -> Result<(), (Error, Lease)> {
        let check = (|| {
            self.validate_lease(ctx, &lease)?;
            if self.len == 0 {
                return Ok(None);
            }
            let index = self.queue[0];
            let slot = self.slots[index];
            if slot.state != State::Queued || !self.backend.validate_task(slot.task) {
                return Err(Error::Stale);
            }
            Ok(Some((index, slot)))
        })();
        let next = match check {
            Ok(v) => v,
            Err(error) => return Err((error, lease)),
        };
        self.owner=None;
        let prepared = if let Some((_, slot)) = next {
            let ticket = Ticket {
                scope: self.scope,
                task: slot.task,
                generation: slot.generation,
            };
            match self.backend.prepare_wake(ticket) {
                Ok(v) => Some(v),
                Err(f) => return Err((self.failed(f), lease)),
            }
        } else {
            None
        };
        if let Some((index, slot)) = next {
            self.queue.copy_within(1..self.len, 0);
            self.len -= 1;
            self.slots[index].state = State::Granted;
            self.owner = Some(Owner {
                task: slot.task,
                generation: slot.generation,
                wait_slot: Some(index),
                claimed: false,
            });
        } else {
            self.owner = None;
        }
        if let Some(prepared) = prepared {
            self.backend.commit_wake(prepared);
        }
        Ok(())
    }
    pub fn cancel(&mut self, ctx: &Context, wait: Wait) -> Result<(), (Error, Wait)> {
        let check = (|| {
            let task = self.current(ctx)?;
            let slot = self.validate_wait(&wait)?;
            if !self.backend.authorize_cancel(task, slot.task) {
                return Err(Error::Foreign);
            }
            if slot.state != State::Queued {
                return Err(Error::Busy);
            }
            self.queue[..self.len]
                .iter()
                .position(|&index| index == wait.slot)
                .ok_or(Error::Stale)
        })();
        let index = match check {
            Ok(v) => v,
            Err(error) => return Err((error, wait)),
        };
        let ticket = Ticket {
            scope: self.scope,
            task: wait.task,
            generation: wait.generation,
        };
        let prepared = match self.backend.prepare_cancel(ticket) {
            Ok(v) => v,
            Err(f) => return Err((self.failed(f), wait)),
        };
        self.queue.copy_within(index + 1..self.len, index);
        self.len -= 1;
        self.slots[wait.slot].state = State::Vacant;
        self.backend.commit_cancel(prepared);
        Ok(())
    }
    pub fn owner(&self) -> Option<(Task, u64, bool)> {
        self.owner.map(|v| (v.task, v.generation, v.claimed))
    }
    pub fn waiting(&self) -> usize {
        self.len
    }
    pub fn faulted(&self) -> bool {
        self.faulted
    }
    pub fn records(&self) -> [Option<WaitRecord>; N] {
        core::array::from_fn(|i| {
            let s = self.slots[i];
            if s.state == State::Vacant {
                None
            } else {
                Some(WaitRecord {
                    task: s.task,
                    generation: s.generation,
                    granted: s.state == State::Granted,
                })
            }
        })
    }
    pub fn with_backend<R>(&self, observe: impl FnOnce(&B) -> R) -> R {
        observe(&self.backend)
    }
    pub fn into_resources(self) -> Result<(B, T), (Error, Self)> {
        if self.faulted
            || self.owner.is_some()
            || self.slots.iter().any(|s| s.state != State::Vacant)
        {
            Err((Error::Busy, self))
        } else {
            Ok((self.backend, self.value))
        }
    }
}

#[cfg(test)]
mod exhaustion_tests {
    use super::*;
    struct Scheduler {
        writes: usize,
        current: Task,
    }
    unsafe impl Backend for Scheduler {
        type Park = Ticket;
        type Wake = Ticket;
        type Cancel = Ticket;
        fn identity(&self) -> u64 {
            1
        }
        fn observation(&self) -> Observation {
            Observation {
                cpu: 0,
                task: Some(self.current),
                in_irq: false,
                daif: 0,
            }
        }
        fn validate_task(&self, t: Task) -> bool {
            t.valid()
        }
        fn prepare_park(&mut self, t: Ticket) -> Result<Ticket, Failure> {
            self.writes += 1;
            Ok(t)
        }
        fn commit_park(&mut self, _: Ticket) {
            self.writes += 1
        }
        fn prepare_wake(&mut self, t: Ticket) -> Result<Ticket, Failure> {
            self.writes += 1;
            Ok(t)
        }
        fn commit_wake(&mut self, _: Ticket) {
            self.writes += 1
        }
        fn authorize_cancel(&self, caller: Task, waiter: Task) -> bool {
            caller == waiter || caller.id == 1
        }
        fn prepare_cancel(&mut self, t: Ticket) -> Result<Ticket, Failure> {
            self.writes += 1;
            Ok(t)
        }
        fn commit_cancel(&mut self, _: Ticket) {
            self.writes += 1
        }
    }
    #[test]
    fn ownership_epoch_exhaustion_never_calls_scheduler() {
        let mut mutex = Mutex::<_, 2>::new(Scheduler {
            writes: 0,
            current: Task {
                id: 1,
                generation: 1,
            },
        })
        .ok()
        .unwrap();
        mutex.epoch = u64::MAX;
        let ctx = unsafe {
            Context::new(
                0,
                Some(Task {
                    id: 1,
                    generation: 1,
                }),
                false,
            )
            .unwrap()
        };
        assert!(matches!(mutex.request(&ctx), Err(Error::Overflow)));
        assert_eq!(mutex.owner(), None);
        assert_eq!(mutex.with_backend(|b| b.writes), 0);
    }
    #[test]
    fn stale_wait_generation_never_claims_a_reused_slot() {
        let first = Task {
            id: 1,
            generation: 1,
        };
        let second = Task {
            id: 2,
            generation: 1,
        };
        let mut mutex = Mutex::<_, 2>::new(Scheduler {
            writes: 0,
            current: first,
        })
        .ok()
        .unwrap();
        let c1 = unsafe { Context::new(0, Some(first), false).unwrap() };
        let c2 = unsafe { Context::new(0, Some(second), false).unwrap() };
        let Request::Owned(lease) = mutex.request(&c1).unwrap() else {
            panic!()
        };
        mutex.backend.current = second;
        let Request::Waiting(wait) = mutex.request(&c2).unwrap() else {
            panic!()
        };
        let forged = Wait {
            scope: wait.scope,
            task: wait.task,
            slot: wait.slot,
            generation: wait.generation,
        };
        mutex.backend.current = first;
        assert!(mutex.cancel(&c1, wait).is_ok());
        mutex.backend.current = second;
        let Request::Waiting(next) = mutex.request(&c2).unwrap() else {
            panic!()
        };
        assert!(next.generation() > forged.generation());
        let before = (mutex.owner(), mutex.records(), mutex.backend.writes);
        let (error, _forged) = mutex.claim(&c2, forged).err().unwrap();
        assert_eq!(error, Error::Stale);
        assert_eq!(
            (mutex.owner(), mutex.records(), mutex.backend.writes),
            before
        );
        mutex.backend.current = first;
        assert!(mutex.release(&c1, lease).is_ok());
        mutex.backend.current = second;
        let real = mutex.claim(&c2, next).unwrap();
        assert!(mutex.release(&c2, real).is_ok());
    }
    #[test]
    fn forged_owner_generation_preserves_the_actual_lease_and_backend() {
        let first = Task {
            id: 1,
            generation: 1,
        };
        let mut mutex = Mutex::<_, 2>::new(Scheduler {
            writes: 0,
            current: first,
        })
        .ok()
        .unwrap();
        let ctx = unsafe { Context::new(0, Some(first), false).unwrap() };
        let Request::Owned(real) = mutex.request(&ctx).unwrap() else {
            panic!()
        };
        let fake = Lease {
            scope: real.scope,
            task: real.task,
            generation: real.generation + 1,
        };
        let before = (mutex.owner(), mutex.records(), mutex.backend.writes);
        let (error, _fake) = mutex.release(&ctx, fake).err().unwrap();
        assert_eq!(error, Error::Stale);
        assert_eq!(
            (mutex.owner(), mutex.records(), mutex.backend.writes),
            before
        );
        assert!(mutex.release(&ctx, real).is_ok());
    }
}