//! Single-worker, in-memory actor runtime. No durability or preemption guarantees.

use std::collections::{BTreeMap, VecDeque};
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::task::{Context, Wake, Waker};

use runa_actor::{ActorId, ActorSnapshot, Event, MailboxId, Message, Status};
use runa_code::CompiledArtifact;
use runa_core::ImplementationId;
use runa_rune::{ExperimentResult, Invocation, Slice};

type WireTransition = (i64, Vec<(u64, i64, i64)>);
pub type RuntimeResult<T> = ExperimentResult<T>;

struct WakeFlag(AtomicBool);
impl Wake for WakeFlag {
    fn wake(self: Arc<Self>) {
        self.0.store(true, Ordering::SeqCst);
    }

    fn wake_by_ref(self: &Arc<Self>) {
        self.0.store(true, Ordering::SeqCst);
    }
}

struct Actor {
    snapshot: ActorSnapshot,
    mailbox: VecDeque<Message>,
    capacity: NonZeroUsize,
    active: Option<Invocation<WireTransition>>,
    wake: Arc<WakeFlag>,
}

#[derive(Default)]
pub struct Runtime {
    actors: BTreeMap<ActorId, Actor>,
    implementations: BTreeMap<ImplementationId, Arc<CompiledArtifact>>,
    next_actor: u64,
    events: Vec<Event>,
}

impl Runtime {
    pub fn publish(
        &mut self,
        artifact: Arc<CompiledArtifact>,
    ) -> ExperimentResult<ImplementationId> {
        let id = artifact.implementation_id;
        if let Some(existing) = self.implementations.get(&id) {
            if existing.source_id != artifact.source_id {
                return Err("implementation identity collision".into());
            }
            return Ok(id);
        }
        self.implementations.insert(id, artifact);
        Ok(id)
    }

    pub fn spawn(
        &mut self,
        implementation: ImplementationId,
        state: i64,
        capacity: NonZeroUsize,
    ) -> ExperimentResult<ActorId> {
        if !self.implementations.contains_key(&implementation) {
            return Err("unknown implementation".into());
        }
        self.next_actor = self
            .next_actor
            .checked_add(1)
            .ok_or("actor IDs exhausted")?;
        let id = ActorId(self.next_actor);
        self.actors.insert(
            id,
            Actor {
                snapshot: ActorSnapshot {
                    id,
                    mailbox: MailboxId(id.0),
                    implementation,
                    incarnation: 1,
                    state,
                    status: Status::Idle,
                    queued: Vec::new(),
                    reserved: None,
                },
                mailbox: VecDeque::new(),
                capacity,
                active: None,
                wake: Arc::new(WakeFlag(AtomicBool::new(true))),
            },
        );
        self.events.push(Event::Spawned(id));
        Ok(id)
    }

    pub fn send(&mut self, id: ActorId, message: Message) -> ExperimentResult<()> {
        let actor = self.actors.get_mut(&id).ok_or("unknown actor")?;
        if matches!(actor.snapshot.status, Status::Stopped | Status::Failed(_)) {
            return Err("actor is not accepting messages".into());
        }
        if actor.mailbox.len() + usize::from(actor.snapshot.reserved.is_some())
            >= actor.capacity.get()
        {
            return Err("mailbox full".into());
        }
        actor.mailbox.push_back(message);
        Ok(())
    }

    pub fn inspect(&self, id: ActorId) -> ExperimentResult<ActorSnapshot> {
        let actor = self.actors.get(&id).ok_or("unknown actor")?;
        let mut snapshot = actor.snapshot.clone();
        snapshot.queued = actor.mailbox.iter().copied().collect();
        Ok(snapshot)
    }

    pub fn actors(&self) -> Vec<ActorSnapshot> {
        self.actors
            .values()
            .map(|actor| {
                let mut snapshot = actor.snapshot.clone();
                snapshot.queued = actor.mailbox.iter().copied().collect();
                snapshot
            })
            .collect()
    }

    pub fn stop(&mut self, id: ActorId) -> ExperimentResult<()> {
        let actor = self.actors.get_mut(&id).ok_or("unknown actor")?;
        actor.active = None;
        actor.snapshot.status = Status::Stopped;
        self.events.push(Event::Stopped(id));
        Ok(())
    }

    pub fn restart(&mut self, id: ActorId) -> ExperimentResult<()> {
        let actor = self.actors.get_mut(&id).ok_or("unknown actor")?;
        if !matches!(actor.snapshot.status, Status::Stopped | Status::Failed(_)) {
            return Err("stop the actor before restarting".into());
        }
        actor.snapshot.incarnation = actor
            .snapshot
            .incarnation
            .checked_add(1)
            .ok_or("incarnation exhausted")?;
        if let Some(message) = actor.snapshot.reserved.take() {
            actor.mailbox.push_front(message);
        }
        actor.snapshot.status = Status::Idle;
        self.events.push(Event::Restarted(id));
        Ok(())
    }

    pub fn drain_events(&mut self) -> Vec<Event> {
        std::mem::take(&mut self.events)
    }

    fn fail(&mut self, id: ActorId, error: String) {
        let actor = self.actors.get_mut(&id).expect("scheduled actor exists");
        actor.active = None;
        actor.snapshot.status = Status::Failed(error.clone());
        self.events.push(Event::Failed(id, error));
    }

    fn commit(&mut self, id: ActorId, transition: WireTransition) -> ExperimentResult<()> {
        let (state, effects) = transition;
        let mut counts = BTreeMap::<ActorId, usize>::new();
        for &(target, _, _) in &effects {
            let target = ActorId(target);
            let actor = self
                .actors
                .get(&target)
                .ok_or("effect targets unknown actor")?;
            if matches!(actor.snapshot.status, Status::Stopped | Status::Failed(_)) {
                return Err("effect targets unavailable actor".into());
            }
            let count = counts.entry(target).or_default();
            *count += 1;
            let retained = usize::from(target != id && actor.snapshot.reserved.is_some());
            if actor.mailbox.len() + retained + *count > actor.capacity.get() {
                return Err("effect exceeds mailbox capacity".into());
            }
        }
        let actor = self.actors.get_mut(&id).expect("scheduled actor exists");
        actor.snapshot.state = state;
        actor.snapshot.reserved = None;
        actor.snapshot.status = Status::Idle;
        actor.active = None;
        for (target, kind, value) in effects {
            self.actors
                .get_mut(&ActorId(target))
                .expect("validated target")
                .mailbox
                .push_back(Message { kind, value });
        }
        self.events.push(Event::Committed(id));
        Ok(())
    }

    /// Give each runnable actor one bounded slice. Awaiting actors require a wake.
    pub fn tick(&mut self, budget: NonZeroUsize) {
        let ids: Vec<_> = self.actors.keys().copied().collect();
        for id in ids {
            let actor = self.actors.get_mut(&id).expect("scheduled actor exists");
            if matches!(actor.snapshot.status, Status::Stopped | Status::Failed(_)) {
                continue;
            }
            if actor.active.is_none() {
                let Some(message) = actor.mailbox.pop_front() else {
                    continue;
                };
                actor.snapshot.reserved = Some(message);
                let program = self.implementations[&actor.snapshot.implementation].program();
                match Invocation::start(
                    program.vm(),
                    "handle",
                    (actor.snapshot.state, message.kind, message.value),
                ) {
                    Ok(invocation) => actor.active = Some(invocation),
                    Err(error) => {
                        self.fail(id, error.to_string());
                        continue;
                    }
                }
                actor.snapshot.status = Status::Running;
            }
            if actor.snapshot.status == Status::Awaiting && !actor.wake.0.load(Ordering::SeqCst) {
                continue;
            }
            actor.wake.0.store(false, Ordering::SeqCst);
            let waker = Waker::from(Arc::clone(&actor.wake));
            let mut cx = Context::from_waker(&waker);
            let outcome = actor
                .active
                .as_mut()
                .expect("active invocation")
                .poll_slice(&mut cx, budget);
            match outcome {
                Slice::Completed(transition) => {
                    if let Err(error) = self.commit(id, transition) {
                        self.fail(id, error.to_string());
                    }
                }
                Slice::Awaiting => actor.snapshot.status = Status::Awaiting,
                Slice::BudgetExhausted => actor.snapshot.status = Status::Running,
                Slice::Failed(error) => self.fail(id, error.to_string()),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use runa_code::CodeStore;

    fn capacity(n: usize) -> NonZeroUsize {
        NonZeroUsize::new(n).unwrap()
    }

    fn publish(runtime: &mut Runtime, source: &str) -> ExperimentResult<ImplementationId> {
        runtime.publish(CodeStore::default().compile(source)?)
    }

    #[test]
    fn two_actors_isolation_and_effects() -> ExperimentResult<()> {
        let mut runtime = Runtime::default();
        let code = publish(
            &mut runtime,
            "pub fn handle(state, kind, value) { if kind == 1 { (state + 1, [(value, 0, 5)]) } else { (state + value, []) } }",
        )?;
        let first = runtime.spawn(code, 0, capacity(4))?;
        let second = runtime.spawn(code, 10, capacity(4))?;
        runtime.send(
            first,
            Message {
                kind: 1,
                value: i64::try_from(second.0)?,
            },
        )?;
        for _ in 0..100 {
            runtime.tick(capacity(10));
        }
        assert_eq!(runtime.inspect(first)?.state, 1);
        assert_eq!(runtime.inspect(second)?.state, 15);
        assert!(runtime.inspect(first)?.reserved.is_none());
        assert!(
            runtime
                .drain_events()
                .iter()
                .any(|e| matches!(e, Event::Committed(id) if *id == second))
        );
        Ok(())
    }

    #[test]
    fn runaway_does_not_starve_healthy_actor() -> ExperimentResult<()> {
        let mut runtime = Runtime::default();
        let runaway = publish(
            &mut runtime,
            "pub fn handle(state, kind, value) { loop {} }",
        )?;
        let healthy = publish(
            &mut runtime,
            "pub fn handle(state, kind, value) { (state + value, []) }",
        )?;
        let first = runtime.spawn(runaway, 0, capacity(1))?;
        let second = runtime.spawn(healthy, 0, capacity(2))?;
        runtime.send(first, Message { kind: 0, value: 1 })?;
        runtime.send(second, Message { kind: 0, value: 42 })?;
        for _ in 0..100 {
            runtime.tick(capacity(5));
        }
        assert_eq!(runtime.inspect(first)?.status, Status::Running);
        assert_eq!(runtime.inspect(first)?.state, 0);
        assert_eq!(runtime.inspect(second)?.state, 42);
        assert!(runtime.send(first, Message { kind: 0, value: 2 }).is_err());
        runtime.stop(first)?;
        assert_eq!(runtime.inspect(first)?.status, Status::Stopped);
        runtime.restart(first)?;
        assert_eq!(runtime.inspect(first)?.incarnation, 2);
        assert_eq!(runtime.inspect(first)?.queued.len(), 1);
        Ok(())
    }

    #[test]
    fn invalid_effect_does_not_commit_state_or_partial_sends() -> ExperimentResult<()> {
        let mut runtime = Runtime::default();
        let code = publish(
            &mut runtime,
            "pub fn handle(state, kind, value) { (99, [(value, 0, 1), (999, 0, 2)]) }",
        )?;
        let actor = runtime.spawn(code, 5, capacity(4))?;
        runtime.send(
            actor,
            Message {
                kind: 0,
                value: i64::try_from(actor.0)?,
            },
        )?;
        for _ in 0..100 {
            runtime.tick(capacity(10));
        }
        let snapshot = runtime.inspect(actor)?;
        assert_eq!(snapshot.state, 5);
        assert!(matches!(snapshot.status, Status::Failed(_)));
        assert!(snapshot.queued.is_empty());
        assert!(snapshot.reserved.is_some());
        Ok(())
    }

    #[test]
    fn awaiting_retains_message_and_state_until_woken() -> ExperimentResult<()> {
        use std::future::Future;
        use std::pin::Pin;
        use std::sync::Mutex;
        use std::task::Poll;

        #[derive(Default)]
        struct Gate {
            ready: AtomicBool,
            dropped: AtomicBool,
            polls: std::sync::atomic::AtomicUsize,
            waker: Mutex<Option<Waker>>,
        }
        struct Wait(Arc<Gate>);
        impl Future for Wait {
            type Output = i64;
            fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<i64> {
                self.0.polls.fetch_add(1, Ordering::SeqCst);
                if self.0.ready.load(Ordering::SeqCst) {
                    Poll::Ready(7)
                } else {
                    *self.0.waker.lock().unwrap() = Some(cx.waker().clone());
                    Poll::Pending
                }
            }
        }
        impl Drop for Wait {
            fn drop(&mut self) {
                self.0.dropped.store(true, Ordering::SeqCst);
            }
        }

        let gate = Arc::new(Gate::default());
        let native_gate = Arc::clone(&gate);
        let mut module = rune::Module::new();
        module
            .function("wait", move || Wait(Arc::clone(&native_gate)))
            .build()?;
        let mut runtime = Runtime::default();
        let code = runtime.publish(
            CodeStore::default().compile_with_module(
                "pub async fn handle(state, kind, value) { if kind == 1 { let x = wait().await; (state + x, []) } else { (state + value, []) } }",
                "test-gate:v1",
                module,
            )?,
        )?;
        let actor = runtime.spawn(code, 10, capacity(2))?;
        runtime.send(actor, Message { kind: 1, value: 0 })?;
        runtime.send(actor, Message { kind: 0, value: 3 })?;
        for _ in 0..100 {
            runtime.tick(capacity(5));
        }
        let pending = runtime.inspect(actor)?;
        assert_eq!(pending.status, Status::Awaiting);
        assert_eq!(pending.state, 10);
        assert_eq!(pending.queued.len(), 1);
        assert_eq!(pending.reserved, Some(Message { kind: 1, value: 0 }));
        let polls = gate.polls.load(Ordering::SeqCst);
        for _ in 0..10 {
            runtime.tick(capacity(5));
        }
        assert_eq!(gate.polls.load(Ordering::SeqCst), polls);
        gate.ready.store(true, Ordering::SeqCst);
        gate.waker
            .lock()
            .unwrap()
            .take()
            .expect("registered wake")
            .wake();
        for _ in 0..100 {
            runtime.tick(capacity(5));
        }
        let completed = runtime.inspect(actor)?;
        assert_eq!(completed.state, 20);
        assert_eq!(completed.status, Status::Idle);
        assert!(completed.reserved.is_none());
        assert!(completed.queued.is_empty());

        gate.ready.store(false, Ordering::SeqCst);
        gate.dropped.store(false, Ordering::SeqCst);
        runtime.send(actor, Message { kind: 1, value: 0 })?;
        for _ in 0..100 {
            runtime.tick(capacity(5));
        }
        assert_eq!(runtime.inspect(actor)?.status, Status::Awaiting);
        runtime.stop(actor)?;
        assert!(gate.dropped.load(Ordering::SeqCst));
        runtime.restart(actor)?;
        let restarted = runtime.inspect(actor)?;
        assert_eq!(restarted.mailbox, pending.mailbox);
        assert_eq!(restarted.state, 20);
        assert_eq!(restarted.incarnation, 2);
        assert_eq!(restarted.queued.len(), 1);
        Ok(())
    }

    #[test]
    fn mailbox_capacity_and_runtime_failures_are_explicit() -> ExperimentResult<()> {
        let mut runtime = Runtime::default();
        let code = publish(
            &mut runtime,
            "pub fn handle(state, kind, value) { (state / value, []) }",
        )?;
        let actor = runtime.spawn(code, 5, capacity(1))?;
        runtime.send(actor, Message { kind: 0, value: 0 })?;
        assert!(runtime.send(actor, Message { kind: 0, value: 1 }).is_err());
        assert!(runtime.inspect(ActorId(999)).is_err());
        for _ in 0..100 {
            runtime.tick(capacity(5));
        }
        assert!(matches!(runtime.inspect(actor)?.status, Status::Failed(_)));
        assert_eq!(runtime.inspect(actor)?.state, 5);
        assert!(
            runtime
                .drain_events()
                .iter()
                .any(|e| matches!(e, Event::Failed(id, _) if *id == actor))
        );
        Ok(())
    }

    #[test]
    fn aggregate_self_sends_respect_capacity_without_partial_commit() -> ExperimentResult<()> {
        let mut runtime = Runtime::default();
        let code = publish(
            &mut runtime,
            "pub fn handle(state, kind, value) { (99, [(value, 0, 1), (value, 0, 2)]) }",
        )?;
        let actor = runtime.spawn(code, 5, capacity(1))?;
        runtime.send(
            actor,
            Message {
                kind: 0,
                value: i64::try_from(actor.0)?,
            },
        )?;
        for _ in 0..100 {
            runtime.tick(capacity(5));
        }
        let snapshot = runtime.inspect(actor)?;
        assert!(matches!(snapshot.status, Status::Failed(_)));
        assert_eq!(snapshot.state, 5);
        assert!(snapshot.queued.is_empty());
        assert!(snapshot.reserved.is_some());
        Ok(())
    }

    #[test]
    fn invalid_transition_shape_fails_without_committing() -> ExperimentResult<()> {
        let mut runtime = Runtime::default();
        let code = publish(&mut runtime, "pub fn handle(state, kind, value) { 99 }")?;
        let actor = runtime.spawn(code, 5, capacity(1))?;
        runtime.send(actor, Message { kind: 0, value: 1 })?;
        for _ in 0..100 {
            runtime.tick(capacity(5));
        }
        let snapshot = runtime.inspect(actor)?;
        assert!(matches!(snapshot.status, Status::Failed(_)));
        assert_eq!(snapshot.state, 5);
        assert!(snapshot.reserved.is_some());
        Ok(())
    }

    #[test]
    fn restart_retries_failed_reservation_from_last_committed_state() -> ExperimentResult<()> {
        use std::sync::atomic::AtomicBool;

        let fail_once = Arc::new(AtomicBool::new(true));
        let native_fail_once = Arc::clone(&fail_once);
        let mut module = rune::Module::new();
        module
            .function("fail_once", move || {
                native_fail_once.swap(false, Ordering::SeqCst)
            })
            .build()?;
        let mut runtime = Runtime::default();
        let code = runtime.publish(
            CodeStore::default().compile_with_module(
                "fn check() { if fail_once() { 1 / 0 } else { 0 } } pub fn handle(state, kind, value) { let _ = check(); (state + value, []) }",
                "test-fail-once:v1",
                module,
            )?,
        )?;
        let actor = runtime.spawn(code, 10, capacity(3))?;
        let first = Message { kind: 1, value: 2 };
        let second = Message { kind: 1, value: 3 };
        runtime.send(actor, first)?;
        runtime.send(actor, second)?;
        for _ in 0..100 {
            runtime.tick(capacity(5));
        }

        let failed = runtime.inspect(actor)?;
        assert!(matches!(failed.status, Status::Failed(_)));
        assert_eq!(failed.state, 10);
        assert_eq!(failed.reserved, Some(first));
        assert_eq!(failed.queued, vec![second]);
        assert!(runtime.send(actor, Message { kind: 1, value: 4 }).is_err());
        assert!(
            runtime
                .drain_events()
                .iter()
                .any(|event| matches!(event, Event::Failed(id, _) if *id == actor))
        );

        runtime.restart(actor)?;
        let restarted = runtime.inspect(actor)?;
        assert_eq!(restarted.id, failed.id);
        assert_eq!(restarted.mailbox, failed.mailbox);
        assert_eq!(restarted.implementation, failed.implementation);
        assert_eq!(restarted.incarnation, failed.incarnation + 1);
        assert_eq!(restarted.state, 10);
        assert_eq!(restarted.queued, vec![first, second]);
        assert_eq!(restarted.reserved, None);

        for _ in 0..100 {
            runtime.tick(capacity(5));
        }
        let completed = runtime.inspect(actor)?;
        assert_eq!(completed.state, 15, "{completed:?}");
        assert_eq!(completed.status, Status::Idle);
        assert_eq!(completed.queued, Vec::<Message>::new());
        assert_eq!(completed.reserved, None);
        Ok(())
    }

    #[test]
    fn stop_and_restart_preserve_reserved_and_queued_fifo_order() -> ExperimentResult<()> {
        let mut runtime = Runtime::default();
        let code = publish(
            &mut runtime,
            "pub fn handle(state, kind, value) { (state * 10 + value, []) }",
        )?;
        let actor = runtime.spawn(code, 0, capacity(3))?;
        let first = Message { kind: 0, value: 1 };
        let second = Message { kind: 0, value: 2 };
        let third = Message { kind: 0, value: 3 };
        runtime.send(actor, first)?;
        runtime.send(actor, second)?;
        runtime.send(actor, third)?;

        runtime.tick(capacity(1));
        let running = runtime.inspect(actor)?;
        assert_eq!(running.status, Status::Running);
        assert_eq!(running.state, 0);
        assert_eq!(running.reserved, Some(first));
        assert_eq!(running.queued, vec![second, third]);

        runtime.stop(actor)?;
        let stopped = runtime.inspect(actor)?;
        assert_eq!(stopped.status, Status::Stopped);
        assert_eq!(stopped.state, 0);
        assert_eq!(stopped.reserved, Some(first));
        assert_eq!(stopped.queued, vec![second, third]);
        assert!(runtime.send(actor, Message { kind: 0, value: 4 }).is_err());

        runtime.restart(actor)?;
        let restarted = runtime.inspect(actor)?;
        assert_eq!(restarted.state, 0);
        assert_eq!(restarted.reserved, None);
        assert_eq!(restarted.queued, vec![first, second, third]);
        for _ in 0..100 {
            runtime.tick(capacity(5));
        }
        let completed = runtime.inspect(actor)?;
        assert_eq!(completed.state, 123);
        assert_eq!(completed.status, Status::Idle);
        assert_eq!(completed.queued, Vec::<Message>::new());
        assert!(completed.reserved.is_none());
        Ok(())
    }
}