use std::{cell::OnceCell, collections::HashMap, fmt::Write};

use crate::{
    data::Data,
    error::Error,
    parser::{Node, Template},
    render::provider::SourceProvider,
};

/// Methods for getting templates
pub mod provider;

/// You can use the render manager to cache any generated templates across
/// multiple render calls.
pub struct RenderManager<'c, P: SourceProvider> {
    /// End of the source cache to be appended to.
    cache_head: CacheReference<'c>,
    templates: HashMap<P::Key, (usize, Template<'c>)>,
    /// Used exclusively for error handling.
    pub(crate) sources: Vec<(P::Key, &'c str)>,
    pub(crate) provider: P,
}

impl<'c, P: SourceProvider> RenderManager<'c, P> {
    /// Create a new render manager using a [`SourceProvider`] and [`SourceCache`].
    /// Note that re-using a source cache for multiple render managers will *not*
    /// improve performance.
    ///
    /// # Example
    ///
    /// ```rust
    /// # use mustache2::render::{SourceCache, RenderManager, provider::FsProvider};
    /// let cache = SourceCache::default();
    /// let mut renderer = RenderManager::new(FsProvider, &cache);
    /// ```
    pub fn new(provider: P, cache: &'c SourceCache) -> Self {
        Self {
            cache_head: CacheReference::Empty(cache),
            templates: HashMap::new(),
            sources: Vec::new(),
            provider,
        }
    }

    fn build_template(&mut self, key: &P::Key, strict: bool) -> Result<bool, Error> {
        let src = match self.provider.get_src(key) {
            Ok(src) => src,
            Err(e) => {
                if strict {
                    return Err(Error::new_general(e));
                } else {
                    return Ok(false);
                }
            }
        };
        let src = self.cache_head.append(src);
        let src_id = self.sources.len();
        let template = Template::parse(src).map_err(|e| e.with_source_id(src_id))?;
        // dbg!(&template);
        self.templates.insert(key.clone(), (src_id, template));
        self.sources.push((key.clone(), src));
        Ok(true)
    }

    /// Parse and render a template, returning the rendered file.
    pub fn render(&mut self, key: P::Key, data: Data) -> Result<String, Error> {
        let mut out = String::new();
        self.render_to(&mut out, key, data)?;
        Ok(out)
    }

    #[cfg(feature = "serde")]
    /// Parse and render a template using data from a serde value.
    pub fn render_serde<T: serde::Serialize>(
        &mut self,
        key: P::Key,
        value: T,
    ) -> Result<String, Error> {
        let mut out = String::new();
        self.render_serde_to(&mut out, key, value)?;
        Ok(out)
    }

    #[cfg(feature = "serde")]
    /// Parse and render a template using data from a serde value. See [`render_to`](Self::render_to) for more information.
    pub fn render_serde_to<T: serde::Serialize>(
        &mut self,
        out: &mut impl Write,
        key: P::Key,
        value: T,
    ) -> Result<(), Error> {
        self.render_to(out, key, crate::to_data(value)?)
    }

    /// Parse and render a template, appending the rendered content to `out`.
    pub fn render_to(
        &mut self,
        out: &mut impl Write,
        key: P::Key,
        data: Data,
    ) -> Result<(), Error> {
        self.render_datastack(out, key, &mut DataStack::from(vec![data]))
    }

    fn render_datastack(
        &mut self,
        out: &mut impl Write,
        key: P::Key,
        stack: &mut DataStack,
    ) -> Result<(), Error> {
        self.build_template(&key, true)?;
        let mut indents = vec![];
        #[cfg(feature = "ext-inheritance")]
        let mut overrides = vec![];
        let mut history = vec![HistoryFrame {
            template_key: key,
            stack_attached: true,
            indent_attached: false,
            #[cfg(feature = "ext-inheritance")]
            override_attached: false,
            state: RenderState::Root { pos: 0 },
        }];

        'step: while let Some(HistoryFrame {
            template_key,
            stack_attached,
            indent_attached,
            #[cfg(feature = "ext-inheritance")]
            override_attached,
            state,
        }) = history.last_mut()
        {
            let (src_id, template) = self.templates.get(&template_key).unwrap();
            let src_id = *src_id;

            let (nodes, pos) = match state {
                RenderState::Root { pos } => (&template.root, pos),
                RenderState::Group { id, pos } => (&template.groups[*id], pos),
                RenderState::Loop {
                    data,
                    index,
                    group_id,
                } => {
                    if *index == data.len() {
                        stack.pop();
                        history.pop();
                    } else {
                        stack.push(data[*index].clone());
                        *index += 1;

                        let template_key = template_key.clone();
                        let id = *group_id;
                        history.push(HistoryFrame {
                            template_key,
                            stack_attached: true,
                            indent_attached: false,
                            #[cfg(feature = "ext-inheritance")]
                            override_attached: false,
                            state: RenderState::Group { id, pos: 0 },
                        });
                    }
                    continue 'step;
                }
            };

            while let Some(node) = nodes.get(*pos) {
                *pos += 1;
                match node {
                    Node::Content(str) => out.write_str(str),
                    Node::Newline => 'bl: {
                        match out.write_char('\n') {
                            Ok(()) => (),
                            Err(e) => break 'bl Err(e),
                        }
                        if *pos < nodes.len() {
                            for indent in &indents {
                                match out.write_str(*indent) {
                                    Ok(()) => (),
                                    Err(e) => break 'bl Err(e),
                                }
                            }
                        }
                        if let Some(indent) = nodes.get(*pos).and_then(Node::get_indent) {
                            match out.write_str(indent) {
                                Ok(()) => (),
                                Err(e) => break 'bl Err(e),
                            }
                        }
                        Ok(())
                    }
                    Node::Variable { key, raw } => 'bl: {
                        if let Some(datum) = stack.get_key(key) {
                            match datum {
                                Data::String(str) if !raw => {
                                    for c in str.chars() {
                                        let r = if c == '&' {
                                            out.write_str("&amp;")
                                        } else if c == '"' {
                                            out.write_str("&quot;")
                                        } else if c == '<' {
                                            out.write_str("&lt;")
                                        } else if c == '>' {
                                            out.write_str("&gt;")
                                        } else {
                                            out.write_char(c)
                                        };
                                        match r {
                                            Ok(()) => (),
                                            Err(e) => break 'bl Err(e),
                                        }
                                    }
                                    Ok(())
                                }
                                #[cfg(feature = "ext-lambdas")]
                                Data::Provider(func) => {
                                    use provider::LambdaSourceProvider;

                                    let Ok(mut func) = func.0.lock() else {
                                        return Err(Error::new_general("provider poisoned"));
                                    };
                                    let src = func();
                                    drop(func);

                                    let cache = SourceCache::default();
                                    let mut renderer =
                                        RenderManager::new(LambdaSourceProvider(Some(src)), &cache);
                                    renderer.render_datastack(out, (), stack)?;
                                    Ok(())
                                }
                                _ => write!(out, "{datum}"),
                            }
                        } else {
                            Ok(())
                        }
                    }
                    #[allow(unused_labels)]
                    Node::Partial {
                        name,
                        source,
                        indent,
                        #[cfg(feature = "ext-dynamic-names")]
                        dynamic,
                    } => 'bl: {
                        #[allow(unused_mut)]
                        let mut name = *name;
                        let source = *source;

                        #[cfg(feature = "ext-dynamic-names")]
                        if *dynamic {
                            if let Some(datum) = stack.get_key(name) {
                                match datum {
                                    Data::String(str) => name = str,
                                    _ => {
                                        break 'bl Ok(());
                                    }
                                }
                            } else {
                                break 'bl Ok(());
                            }
                        }

                        let indent = *indent;
                        let key = self.provider.resolve_partial(template_key, name);
                        let success = self
                            .build_template(&key, cfg!(feature = "strict-resolver"))
                            .map_err(|mut e| {
                                e.source = e.source.or(Some((src_id, source)));
                                e
                            })?;
                        if !success {
                            // borrowing reasons
                            continue 'step;
                        }

                        indents.push(indent);
                        history.push(HistoryFrame {
                            template_key: key,
                            stack_attached: false,
                            indent_attached: true,
                            #[cfg(feature = "ext-inheritance")]
                            override_attached: false,
                            state: RenderState::Root { pos: 0 },
                        });
                        continue 'step;
                    }

                    Node::Section {
                        key,
                        group_id,
                        #[cfg(feature = "ext-lambdas")]
                        span_inside,
                    } => {
                        if let Some(datum) = stack.get_key(key) {
                            let k = template_key.clone();
                            let group_id = *group_id;
                            match datum {
                                Data::List(data) => {
                                    history.push(HistoryFrame {
                                        template_key: k,
                                        stack_attached: false,
                                        indent_attached: false,
                                        #[cfg(feature = "ext-inheritance")]
                                        override_attached: false,
                                        state: RenderState::Loop {
                                            data: data.clone(),
                                            index: 0,
                                            group_id,
                                        },
                                    });
                                    continue 'step;
                                }
                                #[cfg(feature = "ext-lambdas")]
                                Data::Transformer(func) => {
                                    use provider::LambdaSourceProvider;

                                    let Ok(mut func) = func.0.lock() else {
                                        return Err(Error::new_general("transformer poisoned"));
                                    };
                                    let src = func(
                                        &self.sources[src_id].1[span_inside.start..span_inside.end],
                                    );
                                    drop(func);

                                    let cache = SourceCache::default();
                                    let mut renderer =
                                        RenderManager::new(LambdaSourceProvider(Some(src)), &cache);
                                    renderer.render_datastack(out, (), stack)?;
                                    Ok(())
                                }
                                _ => {
                                    if datum.is_truthy() {
                                        stack.push(datum.clone());
                                        history.push(HistoryFrame {
                                            template_key: k,
                                            stack_attached: true,
                                            indent_attached: false,
                                            #[cfg(feature = "ext-inheritance")]
                                            override_attached: false,
                                            state: RenderState::Group {
                                                id: group_id,
                                                pos: 0,
                                            },
                                        });
                                        continue 'step;
                                    }
                                    Ok(())
                                }
                            }
                        } else {
                            Ok(())
                        }
                    }
                    Node::Inverted { key, group_id } => {
                        if stack.get_key(key).is_none_or(|d| !d.is_truthy()) {
                            let template_key = template_key.clone();
                            let group_id = *group_id;
                            history.push(HistoryFrame {
                                template_key,
                                stack_attached: false,
                                indent_attached: false,
                                #[cfg(feature = "ext-inheritance")]
                                override_attached: false,
                                state: RenderState::Group {
                                    id: group_id,
                                    pos: 0,
                                },
                            });
                            continue 'step;
                        } else {
                            Ok(())
                        }
                    }

                    #[cfg(feature = "ext-inheritance")]
                    #[allow(unused_labels)]
                    Node::Parent {
                        name,
                        parameters,
                        indent,
                        source,
                        #[cfg(feature = "ext-dynamic-names")]
                        dynamic,
                    } => 'bl: {
                        #[allow(unused_mut)]
                        let mut name = *name;
                        let source = *source;

                        #[cfg(feature = "ext-dynamic-names")]
                        if *dynamic {
                            if let Some(datum) = stack.get_key(name) {
                                match datum {
                                    Data::String(str) => name = str,
                                    _ => {
                                        break 'bl Ok(());
                                    }
                                }
                            } else {
                                break 'bl Ok(());
                            }
                        }

                        let parameters = parameters.clone();
                        let indent = *indent;
                        let parent_key = self.provider.resolve_partial(template_key, name);
                        let success = self
                            .build_template(&parent_key, cfg!(feature = "strict-resolver"))
                            .map_err(|mut e| {
                                e.source = e.source.or(Some((src_id, source)));
                                e
                            })?;
                        if !success {
                            // borrowing reasons
                            continue 'step;
                        }

                        overrides.push((template_key.clone(), parameters));
                        indents.push(indent);
                        history.push(HistoryFrame {
                            template_key: parent_key,
                            stack_attached: false,
                            indent_attached: true,
                            override_attached: true,
                            state: RenderState::Root { pos: 0 },
                        });

                        continue 'step;
                    }
                    #[cfg(feature = "ext-inheritance")]
                    Node::Block {
                        name,
                        indent,
                        group_id,
                    } => {
                        for (key, blocks) in &overrides {
                            if let Some(group_id) = blocks.get(name) {
                                let template_key = key.clone();
                                let group_id = *group_id;
                                indents.push(indent);
                                history.push(HistoryFrame {
                                    template_key,
                                    stack_attached: false,
                                    indent_attached: true,
                                    #[cfg(feature = "ext-inheritance")]
                                    override_attached: false,
                                    state: RenderState::Group {
                                        id: group_id,
                                        pos: 0,
                                    },
                                });

                                continue 'step;
                            }
                        }

                        let template_key = template_key.clone();
                        let group_id = *group_id;
                        history.push(HistoryFrame {
                            template_key,
                            stack_attached: false,
                            indent_attached: false,
                            #[cfg(feature = "ext-inheritance")]
                            override_attached: false,
                            state: RenderState::Group {
                                id: group_id,
                                pos: 0,
                            },
                        });
                        continue 'step;
                    }
                }
                .map_err(|_| Error::new_general("failed to write output"))?;
            }

            if *stack_attached {
                stack.pop();
            }
            if *indent_attached {
                indents.pop();
            }
            #[cfg(feature = "ext-inheritance")]
            if *override_attached {
                overrides.pop();
            }
            history.pop();
        }

        Ok(())
    }
}

#[derive(Debug)]
struct DataStack {
    vec: Vec<Data>,
}

impl DataStack {
    fn push(&mut self, value: Data) {
        self.vec.push(value);
    }
    fn pop(&mut self) {
        self.vec.pop();
    }

    fn get_key<'a>(&'a self, key: &str) -> Option<&'a Data> {
        if key == "." {
            self.vec.last()
        } else {
            let mut segments = key.split(".");
            let key = segments.next()?;
            let mut value = self.vec.iter().rev().find_map(|d| match d {
                Data::Map(map) => map.get(key),
                _ => None,
            })?;

            for key in segments {
                value = match value {
                    Data::Map(map) => map.get(key),
                    _ => None,
                }?;
            }

            Some(value)
        }
    }
}

impl From<Vec<Data>> for DataStack {
    fn from(vec: Vec<Data>) -> Self {
        Self { vec }
    }
}

struct HistoryFrame<K> {
    template_key: K,
    stack_attached: bool,
    indent_attached: bool,
    #[cfg(feature = "ext-inheritance")]
    override_attached: bool,
    state: RenderState,
}

enum RenderState {
    Root {
        pos: usize,
    },
    Group {
        id: usize,
        pos: usize,
    },
    Loop {
        data: Vec<Data>,
        index: usize,
        group_id: usize,
    },
}

/// Linked list responsible for storing source code for the [`RenderManager`].
///
/// # Performance
///
/// We are using a linked list because we need to append items to it with an
/// immutable reference. It's okay performance wise because we only ever append
/// to the list; we never search or index it.
#[derive(Default)]
pub struct SourceCache(OnceCell<CacheItem>);

struct CacheItem {
    src: String,
    next: OnceCell<Box<CacheItem>>,
}

#[derive(Clone, Copy)]
enum CacheReference<'a> {
    Empty(&'a SourceCache),
    Tail(&'a CacheItem),
}

impl<'a> CacheReference<'a> {
    fn append(&mut self, src: String) -> &'a str {
        let item = CacheItem {
            src,
            next: OnceCell::new(),
        };
        let cell = match self {
            CacheReference::Empty(c) => c.0.get_or_init(|| item),
            CacheReference::Tail(c) => &*c.next.get_or_init(|| Box::new(item)),
        };
        *self = Self::Tail(cell);
        &cell.src
    }
}
