use core::{array, marker::PhantomData};
use turkos_os_threads::threads::ThreadId;

pub const MAX_WAITS: usize = 8;
pub const MAX_EVENTS: usize = 16;

#[cfg(test)]
#[allow(dead_code)]
mod private_tests;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    Configuration,
    Context,
    Cpu,
    Full,
    Foreign,
    Stale,
    Pending,
    State,
    Exhausted,
    Source,
    Backend,
    Quarantined,
    Duplicate,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Operation {
    Arm,
    Depart,
    Land,
    Service,
    Cancel,
    Retire,
    Produce,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Execution {
    pub cpu: u32,
    pub irq_masked: bool,
    pub in_irq: bool,
    pub current: Option<ThreadId>,
}
/// Architecture ingress calls this before forming a raw mutable loan. APIs
/// also repeat their actual immutable resource observation.
pub fn context(cpu: u32, execution: Execution, operation: Operation) -> Result<(), Error> {
    if execution.cpu != cpu {
        return Err(Error::Cpu);
    }
    if !execution.irq_masked {
        return Err(Error::Context);
    }
    if execution.in_irq && operation != Operation::Produce {
        return Err(Error::Context);
    }
    Ok(())
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Source {
    pub channel: u64,
    pub epoch: u64,
    pub owner: u64,
}
impl Source {
    fn valid(self) -> bool {
        self.channel != 0 && self.epoch != 0 && self.owner != 0
    }
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Observation {
    pub source: Source,
    pub sequence: u64,
    pub ready: bool,
    pub now: u64,
}
/// # Safety
/// Owns a genuine immutable producer registration/source pin. Observations
/// read actual CPU/DAIF and producer/clock progression; no caller argument is
/// substituted. IRQ observation never borrows allocator, runtime or WaitSet.
pub unsafe trait Producer {
    fn execution(&self) -> Execution;
    fn observe(&self) -> Result<Observation, Error>;
}

#[derive(Clone, Copy)]
struct Event {
    observation: Observation,
}
/// A separate bounded owner; an IRQ can mutate it while task I/O or a thread
/// continuation holds unrelated resources. Events have no public constructor.
pub struct Ingress<P: Producer, const N: usize> {
    producer: P,
    identity: u64,
    cpu: u32,
    source: Source,
    events: [Option<Event>; N],
    first: usize,
    count: usize,
    sequence: u64,
    clock: u64,
    lost: u64,
    last_lost_sequence: Option<u64>,
    non_send: PhantomData<*mut ()>,
}
impl<P: Producer, const N: usize> Ingress<P, N> {
    pub fn new(producer: P) -> Result<Self, (Error, P)> {
        if N == 0 || N > MAX_EVENTS {
            return Err((Error::Configuration, producer));
        }
        let execution = producer.execution();
        if let Err(error) = context(execution.cpu, execution, Operation::Arm) {
            return Err((error, producer));
        }
        let observation = match producer.observe() {
            Ok(o) if o.source.valid() => o,
            Ok(_) => return Err((Error::Source, producer)),
            Err(e) => return Err((e, producer)),
        };
        let identity = match identity() {
            Ok(i) => i,
            Err(e) => return Err((e, producer)),
        };
        Ok(Self {
            producer,
            identity,
            cpu: execution.cpu,
            source: observation.source,
            events: [None; N],
            first: 0,
            count: 0,
            sequence: observation.sequence,
            clock: observation.now,
            lost: 0,
            last_lost_sequence: None,
            non_send: PhantomData,
        })
    }
    fn observe(&self) -> Result<Observation, Error> {
        let o = self.producer.observe()?;
        if o.source != self.source
            || o.sequence < self.sequence
            || self.last_lost_sequence.is_some_and(|v| o.sequence < v)
            || o.now < self.clock
        {
            return Err(Error::Source);
        }
        Ok(o)
    }
    pub fn poll(&mut self) -> Result<bool, Error> {
        context(self.cpu, self.producer.execution(), Operation::Produce)?;
        let o = self.observe()?;
        if o.sequence == self.sequence {
            self.clock = o.now;
            return Ok(false);
        }
        if self.count == N {
            if self.last_lost_sequence != Some(o.sequence) {
                self.lost = self.lost.checked_add(1).ok_or(Error::Exhausted)?;
                self.last_lost_sequence = Some(o.sequence);
            }
            self.clock = o.now;
            return Err(Error::Full);
        }
        self.events[(self.first + self.count) % N] = Some(Event { observation: o });
        self.count += 1;
        self.sequence = o.sequence;
        self.clock = o.now;
        Ok(true)
    }
    pub fn pending(&self) -> usize {
        self.count
    }
    pub fn lost(&self) -> u64 {
        self.lost
    }
    pub fn source(&self) -> Source {
        self.source
    }
    pub fn into_producer(self) -> Result<P, Self> {
        if context(self.cpu, self.producer.execution(), Operation::Retire).is_ok()
            && self.count == 0
        {
            Ok(self.producer)
        } else {
            Err(self)
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Phase {
    Empty,
    Armed,
    Staged,
    Departed,
    Sleeping,
    Terminal,
    Quarantined,
    Retired,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Outcome {
    Signal,
    Timeout,
    Cancelled,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FailurePhase {
    Before,
    After,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Failure {
    pub error: Error,
    pub phase: FailurePhase,
}
/// # Safety
/// Pins and Prepared are actual nonclone custody with inert destructors. The
/// backend owns a frozen Threads Table and AsyncBank exclusively. `stage`
/// calls actual AsyncBank.sleep_to while Table remains Running; `depart`
/// changes the real Table and retains its exact Wait in Prepared. `landed`
/// validates actual target CPU/root/SP/assembly episode. No mutable borrow
/// of this backend or WaitSet spans assembly execution. A stage error returns
/// before acquiring a Prepared or changing the thread/context tables. Before failures change
/// nothing; every uncertain change is After. `resolve` checks the actual live
/// sleeping thread before wake/cancel, never calls a dead callback, and consumes
/// its exact Wait once. Release errors return every real owner unchanged.
pub unsafe trait Backend {
    type Pins;
    type Prepared;
    type Capture;
    fn execution(&self) -> Execution;
    fn observe(&self, source: Source) -> Result<Observation, Error>;
    fn retain(
        &mut self,
        thread: ThreadId,
        target: ThreadId,
        source: Source,
    ) -> Result<Self::Pins, Error>;
    fn validate(
        &self,
        pins: &Self::Pins,
        thread: ThreadId,
        target: ThreadId,
        source: Source,
    ) -> Result<(), Error>;
    fn stage(
        &mut self,
        pins: &mut Self::Pins,
        capture: Self::Capture,
    ) -> Result<Self::Prepared, Error>;
    fn abort(&mut self, prepared: &mut Self::Prepared) -> Result<(), Failure>;
    fn depart(&mut self, prepared: &mut Self::Prepared) -> Result<(), Failure>;
    fn landed(&mut self, prepared: &mut Self::Prepared) -> Result<(), Failure>;
    fn resolve(
        &mut self,
        pins: &mut Self::Pins,
        prepared: &mut Self::Prepared,
        outcome: Outcome,
    ) -> Result<(), Failure>;
    fn release(
        &mut self,
        owned: Retained<Self::Pins, Self::Prepared>,
    ) -> Result<(), ReleaseFailure<Self::Pins, Self::Prepared>>;
}
pub struct Retained<P, T> {
    pub pins: P,
    pub prepared: Option<T>,
}
pub type ReleaseFailure<P, T> = (Error, Retained<P, T>);
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Record {
    pub phase: Phase,
    pub generation: u64,
    pub thread: Option<ThreadId>,
    pub target: Option<ThreadId>,
    pub source: Option<Source>,
    pub ingress: u64,
    pub sequence: u64,
    pub clock: u64,
    pub deadline: Option<u64>,
    pub signal_link: bool,
    pub deadline_link: bool,
    pub winner: Option<Outcome>,
}
impl Record {
    fn empty() -> Self {
        Self {
            phase: Phase::Empty,
            generation: 0,
            thread: None,
            target: None,
            source: None,
            ingress: 0,
            sequence: 0,
            clock: 0,
            deadline: None,
            signal_link: false,
            deadline_link: false,
            winner: None,
        }
    }
}
struct Slot<P, T> {
    record: Record,
    owned: Option<Retained<P, T>>,
}
/// Nonclone authority for one wait generation. Forgetting it leaves custody in
/// the set; it never releases resources implicitly.
pub struct WaitToken {
    set: u64,
    slot: usize,
    generation: u64,
}
impl WaitToken {
    pub fn generation(&self) -> u64 {
        self.generation
    }
    pub fn slot(&self) -> usize {
        self.slot
    }
}
pub struct WaitTransfer {
    token: WaitToken,
}
pub enum Departure {
    Ready(WaitToken),
    Switch(WaitTransfer),
}
pub struct ArmFailure {
    pub error: Error,
    pub retained: Option<WaitToken>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Completion {
    pub thread: ThreadId,
    pub source: Source,
    pub generation: u64,
    pub outcome: Outcome,
}

pub struct WaitSet<B: Backend, const N: usize> {
    backend: B,
    identity: u64,
    cpu: u32,
    slots: [Slot<B::Pins, B::Prepared>; N],
    winners: u64,
    non_send: PhantomData<*mut ()>,
}
#[cfg(not(target_os = "none"))]
fn identity() -> Result<u64, Error> {
    use core::sync::atomic::{AtomicU64, Ordering};
    static NEXT: AtomicU64 = AtomicU64::new(1);
    let mut value = NEXT.load(Ordering::Relaxed);
    loop {
        let next = value.checked_add(1).ok_or(Error::Exhausted)?;
        match NEXT.compare_exchange_weak(value, next, Ordering::Relaxed, Ordering::Relaxed) {
            Ok(v) => return Ok(v),
            Err(v) => value = v,
        }
    }
}
#[cfg(target_os = "none")]
fn identity() -> Result<u64, Error> {
    unsafe extern "C" {
        fn __os_k1080_wait_identity() -> u64;
    }
    let value = unsafe { __os_k1080_wait_identity() };
    if value == 0 {
        Err(Error::Exhausted)
    } else {
        Ok(value)
    }
}

impl<B: Backend, const N: usize> WaitSet<B, N> {
    pub fn new(backend: B) -> Result<Self, (Error, B)> {
        if N == 0 || N > MAX_WAITS {
            return Err((Error::Configuration, backend));
        }
        let execution = backend.execution();
        if let Err(e) = context(execution.cpu, execution, Operation::Arm) {
            return Err((e, backend));
        }
        let identity = match identity() {
            Ok(i) => i,
            Err(e) => return Err((e, backend)),
        };
        Ok(Self {
            backend,
            identity,
            cpu: execution.cpu,
            slots: array::from_fn(|_| Slot {
                record: Record::empty(),
                owned: None,
            }),
            winners: 0,
            non_send: PhantomData,
        })
    }
    fn gate(&self, operation: Operation) -> Result<Execution, Error> {
        let e = self.backend.execution();
        context(self.cpu, e, operation)?;
        Ok(e)
    }
    fn checked(&self, token: &WaitToken) -> Result<usize, Error> {
        if token.set != self.identity {
            return Err(Error::Foreign);
        }
        let slot = self.slots.get(token.slot).ok_or(Error::Stale)?;
        if slot.record.generation != token.generation || slot.owned.is_none() {
            return Err(Error::Stale);
        }
        Ok(token.slot)
    }
    fn observed(&self, i: usize) -> Result<Observation, Error> {
        let s = &self.slots[i];
        let r = s.record;
        let owned = s.owned.as_ref().ok_or(Error::State)?;
        self.backend.validate(
            &owned.pins,
            r.thread.ok_or(Error::State)?,
            r.target.ok_or(Error::State)?,
            r.source.ok_or(Error::State)?,
        )?;
        let o = self.backend.observe(r.source.ok_or(Error::State)?)?;
        if Some(o.source) != r.source || o.sequence < r.sequence || o.now < r.clock {
            return Err(Error::Source);
        }
        Ok(o)
    }
    fn winner(&mut self, i: usize, outcome: Outcome, clock: u64) -> Result<(), Error> {
        let next = self.winners.checked_add(1).ok_or(Error::Exhausted)?;
        let r = self.slots[i].record;
        if r.winner.is_some() {
            return Err(Error::Duplicate);
        }
        let o = self.slots[i].owned.as_mut().ok_or(Error::State)?;
        let result = match r.phase {
            Phase::Armed => Ok(()),
            Phase::Staged => self.backend.abort(o.prepared.as_mut().ok_or(Error::State)?),
            Phase::Sleeping => self.backend.resolve(
                &mut o.pins,
                o.prepared.as_mut().ok_or(Error::State)?,
                outcome,
            ),
            Phase::Departed => return Err(Error::Pending),
            Phase::Quarantined => return Err(Error::Quarantined),
            _ => return Err(Error::State),
        };
        if let Err(f) = result {
            if f.phase == FailurePhase::After {
                self.slots[i].record.phase = Phase::Quarantined;
            }
            return Err(f.error);
        }
        let r = &mut self.slots[i].record;
        r.phase = Phase::Terminal;
        r.signal_link = false;
        r.deadline_link = false;
        r.winner = Some(outcome);
        r.clock = clock;
        self.winners = next;
        Ok(())
    }
    fn rollback_arm(&mut self, i: usize, error: Error, token: WaitToken) -> ArmFailure {
        let generation = self.slots[i].record.generation;
        let owned = self.slots[i].owned.take().unwrap();
        match self.backend.release(owned) {
            Ok(()) => {
                self.slots[i].record = Record {
                    generation,
                    ..Record::empty()
                };
                ArmFailure {
                    error,
                    retained: None,
                }
            }
            Err((_, owned)) => {
                self.slots[i].owned = Some(owned);
                self.slots[i].record.phase = Phase::Quarantined;
                ArmFailure {
                    error,
                    retained: Some(token),
                }
            }
        }
    }
    pub fn arm<P: Producer, const E: usize>(
        &mut self,
        ingress: &Ingress<P, E>,
        thread: ThreadId,
        target: ThreadId,
        capture: B::Capture,
        deadline: Option<u64>,
    ) -> Result<WaitToken, ArmFailure> {
        let fail = |error| ArmFailure {
            error,
            retained: None,
        };
        let execution = self.gate(Operation::Arm).map_err(fail)?;
        if execution.current != Some(thread) || thread == target {
            return Err(fail(Error::State));
        }
        context(self.cpu, ingress.producer.execution(), Operation::Arm).map_err(fail)?;
        let producer = ingress.observe().map_err(fail)?;
        let initial = self.backend.observe(producer.source).map_err(fail)?;
        if initial.source != producer.source
            || initial.sequence < producer.sequence
            || initial.now < producer.now
        {
            return Err(fail(Error::Source));
        }
        if self
            .slots
            .iter()
            .any(|s| s.owned.is_some() && s.record.thread == Some(thread))
        {
            return Err(fail(Error::Duplicate));
        }
        let i = self
            .slots
            .iter()
            .position(|s| s.record.phase == Phase::Empty)
            .ok_or_else(|| fail(Error::Full))?;
        let generation = match self.slots[i].record.generation.checked_add(1) {
            Some(g) => g,
            None => {
                self.slots[i].record.phase = Phase::Retired;
                return Err(fail(Error::Exhausted));
            }
        };
        let pins = self
            .backend
            .retain(thread, target, initial.source)
            .map_err(fail)?;
        self.slots[i] = Slot {
            record: Record {
                phase: Phase::Armed,
                generation,
                thread: Some(thread),
                target: Some(target),
                source: Some(initial.source),
                ingress: ingress.identity,
                sequence: initial.sequence,
                clock: initial.now,
                deadline,
                signal_link: true,
                deadline_link: deadline.is_some(),
                winner: None,
            },
            owned: Some(Retained {
                pins,
                prepared: None,
            }),
        };
        let token = WaitToken {
            set: self.identity,
            slot: i,
            generation,
        };
        if let Err(e) = self.backend.validate(
            &self.slots[i].owned.as_ref().unwrap().pins,
            thread,
            target,
            initial.source,
        ) {
            return Err(self.rollback_arm(i, e, token));
        }
        let immediate = if initial.ready {
            Some(Outcome::Signal)
        } else if deadline.is_some_and(|d| initial.now >= d) {
            Some(Outcome::Timeout)
        } else {
            None
        };
        if let Some(outcome) = immediate {
            if let Err(e) = self.winner(i, outcome, initial.now) {
                return Err(ArmFailure {
                    error: e,
                    retained: Some(token),
                });
            }
            return Ok(token);
        }
        let prepared = match self
            .backend
            .stage(&mut self.slots[i].owned.as_mut().unwrap().pins, capture)
        {
            Ok(p) => p,
            Err(e) => return Err(self.rollback_arm(i, e, token)),
        };
        self.slots[i].owned.as_mut().unwrap().prepared = Some(prepared);
        self.slots[i].record.phase = Phase::Staged;
        Ok(token)
    }
    pub fn depart(&mut self, token: WaitToken) -> Result<Departure, (Error, WaitToken)> {
        if let Err(e) = self.gate(Operation::Depart) {
            return Err((e, token));
        }
        let i = match self.checked(&token) {
            Ok(i) => i,
            Err(e) => return Err((e, token)),
        };
        let r = self.slots[i].record;
        if r.phase == Phase::Terminal {
            return Ok(Departure::Ready(token));
        }
        if self.backend.execution().current != r.thread || r.phase != Phase::Staged {
            return Err((Error::State, token));
        }
        let o = match self.observed(i) {
            Ok(o) => o,
            Err(e) => return Err((e, token)),
        };
        let early = if o.ready || o.sequence > r.sequence {
            Some(Outcome::Signal)
        } else if r.deadline.is_some_and(|d| o.now >= d) {
            Some(Outcome::Timeout)
        } else {
            None
        };
        if let Some(outcome) = early {
            return match self.winner(i, outcome, o.now) {
                Ok(()) => Ok(Departure::Ready(token)),
                Err(e) => Err((e, token)),
            };
        }
        let prepared = self.slots[i]
            .owned
            .as_mut()
            .unwrap()
            .prepared
            .as_mut()
            .unwrap();
        if let Err(f) = self.backend.depart(prepared) {
            if f.phase == FailurePhase::After {
                self.slots[i].record.phase = Phase::Quarantined;
            }
            return Err((f.error, token));
        }
        self.slots[i].record.phase = Phase::Departed;
        self.slots[i].record.clock = o.now;
        Ok(Departure::Switch(WaitTransfer { token }))
    }
    pub fn landed(&mut self, transfer: WaitTransfer) -> Result<WaitToken, (Error, WaitTransfer)> {
        if let Err(e) = self.gate(Operation::Land) {
            return Err((e, transfer));
        }
        let i = match self.checked(&transfer.token) {
            Ok(i) => i,
            Err(e) => return Err((e, transfer)),
        };
        if self.slots[i].record.phase != Phase::Departed
            || self.backend.execution().current != self.slots[i].record.target
        {
            return Err((Error::State, transfer));
        }
        let observation = match self.observed(i) {
            Ok(o) => o,
            Err(e) => return Err((e, transfer)),
        };
        let prepared = self.slots[i]
            .owned
            .as_mut()
            .unwrap()
            .prepared
            .as_mut()
            .unwrap();
        if let Err(f) = self.backend.landed(prepared) {
            if f.phase == FailurePhase::After {
                self.slots[i].record.phase = Phase::Quarantined;
            }
            return Err((f.error, transfer));
        }
        self.slots[i].record.phase = Phase::Sleeping;
        self.slots[i].record.clock = observation.now;
        Ok(transfer.token)
    }
    /// Each winner commits independently. On a later error, earlier Terminal
    /// records and the winners counter remain queryable; their tokens retire
    /// separately. Err never rolls back or redelivers a committed winner.
    pub fn service<P: Producer, const E: usize>(
        &mut self,
        ingress: &mut Ingress<P, E>,
    ) -> Result<usize, Error> {
        self.gate(Operation::Service)?;
        context(self.cpu, ingress.producer.execution(), Operation::Service)?;
        let now = ingress.observe()?;
        if self.slots.iter().any(|s| {
            s.record.ingress == ingress.identity
                && s.record.signal_link
                && s.record.phase == Phase::Departed
        }) {
            return Err(Error::Pending);
        }
        let event = if ingress.count == 0 {
            None
        } else {
            ingress.events[ingress.first]
        };
        let mut completed = 0;
        for i in 0..N {
            let r = self.slots[i].record;
            if r.ingress != ingress.identity || !r.signal_link {
                continue;
            }
            if r.phase == Phase::Departed {
                return Err(Error::Pending);
            }
            let observed = self.observed(i)?;
            if observed.now < now.now {
                return Err(Error::Source);
            }
            let signal = observed.ready
                || observed.sequence > r.sequence
                || event.is_some_and(|e| {
                    e.observation.source == observed.source && e.observation.sequence > r.sequence
                });
            let outcome = if signal {
                Some(Outcome::Signal)
            } else if r.deadline.is_some_and(|d| observed.now >= d) {
                Some(Outcome::Timeout)
            } else {
                None
            };
            if let Some(outcome) = outcome {
                self.winner(i, outcome, observed.now)?;
                completed += 1;
            } else {
                self.slots[i].record.clock = observed.now;
            }
        }
        if event.is_some() {
            ingress.events[ingress.first] = None;
            ingress.first = (ingress.first + 1) % E;
            ingress.count -= 1;
        }
        ingress.clock = now.now;
        Ok(completed)
    }
    pub fn cancel(&mut self, token: &WaitToken) -> Result<(), Error> {
        self.gate(Operation::Cancel)?;
        let i = self.checked(token)?;
        let observation = self.observed(i)?;
        self.winner(i, Outcome::Cancelled, observation.now)
    }
    pub fn retire(&mut self, token: WaitToken) -> Result<Completion, (Error, WaitToken)> {
        if let Err(e) = self.gate(Operation::Retire) {
            return Err((e, token));
        }
        let i = match self.checked(&token) {
            Ok(i) => i,
            Err(e) => return Err((e, token)),
        };
        let r = self.slots[i].record;
        if r.phase != Phase::Terminal || r.signal_link || r.deadline_link {
            return Err((Error::State, token));
        }
        let owned = self.slots[i].owned.take().unwrap();
        if let Err((e, owned)) = self.backend.release(owned) {
            self.slots[i].owned = Some(owned);
            self.slots[i].record.phase = Phase::Quarantined;
            return Err((e, token));
        }
        self.slots[i].record = Record {
            generation: r.generation,
            ..Record::empty()
        };
        Ok(Completion {
            thread: r.thread.unwrap(),
            source: r.source.unwrap(),
            generation: r.generation,
            outcome: r.winner.unwrap(),
        })
    }
    pub fn records(&self) -> [Record; N] {
        array::from_fn(|i| self.slots[i].record)
    }
    pub fn winners(&self) -> u64 {
        self.winners
    }
    pub fn with_backend_observation<R>(&self, f: impl FnOnce(&B) -> R) -> R {
        f(&self.backend)
    }
    pub fn into_backend(self) -> Result<B, Self> {
        if self.gate(Operation::Retire).is_ok() && self.slots.iter().all(|s| s.owned.is_none()) {
            Ok(self.backend)
        } else {
            Err(self)
        }
    }
}