/*  Mycelium Scheme
 *  Copyright (C) 2025 Ava Affine
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */

use core::fmt::{self, Debug, Formatter};
use core::ops::Index;
use alloc::rc::Rc;

struct StackInner<T: Sized> {
    pub next: Stack<T>,
    pub data: T
}

struct Stack<T: Sized> (Rc<Option<StackInner<T>>>);

struct StackStackInner<T: Sized> {
    next: StackStack<T>,
    count: usize,
    stack: Stack<T>,
}

pub struct StackStack<T: Sized> (Rc<Option<StackStackInner<T>>>);

impl<T> From<T> for StackInner<T> {
    fn from(t: T) -> StackInner<T> {
        StackInner {
            next: Stack(Rc::from(None)),
            data: t,
        }
    }
}

impl<T> From<StackInner<T>> for Stack<T> {
    fn from(t: StackInner<T>) -> Stack<T> {
        Stack(Rc::from(Some(t)))
    }
}

impl<T> Index<usize> for StackStack<T> {
    type Output = T;
    fn index(&self, index: usize) -> &T {
        if let Some(ref inner) = *self.0 {
            // pass on to next
            if inner.count <= index {
                &inner.next[index - inner.count]

            // fetch from our stack
            } else {
                let mut idx = index;
                let mut cursor = &inner.stack;
                while let Some(ref node) = *cursor.0 {
                    if idx == 0 {
                        return &node.data
                    }
                    idx -= 1;
                    cursor = &node.next;
                }
                // should never hit this case
                panic!("encountered inconsistent lengths in stackstack")
            }

        // guaranteed out of bounds
        } else {
            panic!("index out of bounds on stackstack access")
        }
    }
}

impl<T: Debug> Debug for StackStack<T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        let mut ss_idx = 1;
        let mut ss_cur = &*self.0;
        while let Some(inner) = ss_cur {
            write!(f, "Frame {ss_idx}:")?;
            let mut s_cur = &*inner.stack.0;
            while let Some(node) = s_cur {
                write!(f, " {:#?}", node.data)?;
                s_cur = &*node.next.0;
            }
            write!(f, "\n")?;
            ss_cur = &*inner.next.0;
            ss_idx += 1;
        }

        write!(f, "\n")
    }
}

impl<T> Stack<T> {
    fn push(&mut self, item: T) {
        self.0 = Rc::from(Some(StackInner{
            data: item,
            next: Stack(self.0.clone()),
        }))
    }

    fn pop(&mut self) -> T {
        // clone self.0 and then drop first ref, decreasing strong count back to 1
        let d = self.0.clone();
        self.0 = Rc::new(None);

        // deconstruct the rc that formerly held self.0
        let b = Rc::into_inner(d).unwrap();
        if let Some(inner) = b {
            let data = inner.data;
            self.0 = inner.next.0;
            data
        } else {
            panic!("pop from 0 length stack")
        }
    }
}

impl<T> StackStack<T> {
    pub fn push_current_stack(&mut self, item: T) {
        if let Some(inner) = Rc::get_mut(&mut self.0).unwrap() {
            inner.stack.push(item);
            inner.count += 1;
        } else {
            panic!("push to uninitialized stackstack")
        }
    }

    pub fn pop_current_stack(&mut self) -> T {
        if let Some(inner) = Rc::get_mut(&mut self.0).unwrap() {
            inner.count -= 1;
            inner.stack.pop()
        } else {
            panic!("pop from uninitialized stackstack")
        }
    }

    pub fn add_stack(&mut self) {
        self.0 = Rc::from(Some(StackStackInner{
            next: StackStack(self.0.clone()),
            count: 0,
            stack: Stack(Rc::from(None)),
        }))
    }

    pub fn destroy_top_stack(&mut self) {
        let s = Rc::get_mut(&mut self.0).unwrap();
        if let Some(inner) = s {
            self.0 = inner.next.0.clone()
        } else {
            panic!("del from empty stackstack")
        }
    }

    pub fn new() -> StackStack<T> {
        StackStack(Rc::from(Some(StackStackInner{
            count: 0,
            next: StackStack(Rc::from(None)),
            stack: Stack(Rc::from(None)),
        })))
    }

    pub fn len(&self) -> usize {
        if let Some(ref inner) = *self.0 {
            if let Some(_) = *inner.next.0 {
                inner.next.len() + inner.count
            } else {
                inner.count
            }
        } else {
            0
        }
    }
}

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

    #[test]
    fn test_alloc_new_stack_and_push_many() {
        let mut g = StackStack::<i8>::new();
        g.add_stack();
        g.push_current_stack(0);
        g.push_current_stack(1);
        g.push_current_stack(2);
        assert_eq!(g.len(), 3);
        g.add_stack();
        g.push_current_stack(3);
        g.push_current_stack(4);
        assert_eq!(g.len(), 5);

        assert_eq!(g.pop_current_stack(), 4);
        assert_eq!(g.pop_current_stack(), 3);
        g.destroy_top_stack();
        assert_eq!(g.pop_current_stack(), 2);
        assert_eq!(g.pop_current_stack(), 1);
        assert_eq!(g.pop_current_stack(), 0);
    }

    #[test]
    fn test_stack_index_bounds() {
        let mut g = StackStack::<i8>::new();
        g.add_stack();
        g.push_current_stack(0);
        g.push_current_stack(1);
        g.push_current_stack(2);
        assert_eq!(g.len(), 3);
        g.add_stack();
        g.push_current_stack(3);
        g.push_current_stack(4);
        assert_eq!(g.len(), 5);

        assert_eq!(g[0], 4);
        assert_eq!(g[1], 3);
        assert_eq!(g[2], 2);
        assert_eq!(g[3], 1);
        assert_eq!(g[4], 0);

        g.destroy_top_stack();
        assert_eq!(g.len(), 3);
        assert_eq!(g[0], 2);
        assert_eq!(g[1], 1);
        assert_eq!(g[2], 0);
    }
}