use crate::{
    error::{Error, HasSpan},
    interpreter::Interpreter,
    utils::Ptr,
};
use quote::ToTokens;
use std::{collections::HashMap, fmt::Formatter};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum BacktraceEvent<'ast> {
    FnCall(Ptr<'ast, syn::ExprCall>),
    MethodCall(Ptr<'ast, syn::ExprMethodCall>),
    If(Ptr<'ast, syn::ExprIf>),
    ForLoop(Ptr<'ast, syn::ExprForLoop>),
}

impl BacktraceEvent<'_> {
    fn is_call(&self) -> bool {
        !matches!(self, BacktraceEvent::If(_))
    }
}

impl<'ast> From<&'ast syn::ExprCall> for BacktraceEvent<'ast> {
    fn from(arg: &'ast syn::ExprCall) -> Self {
        BacktraceEvent::FnCall(Ptr(arg))
    }
}

impl<'ast> From<&'ast syn::ExprMethodCall> for BacktraceEvent<'ast> {
    fn from(arg: &'ast syn::ExprMethodCall) -> Self {
        BacktraceEvent::MethodCall(Ptr(arg))
    }
}
impl<'ast> From<&'ast syn::ExprIf> for BacktraceEvent<'ast> {
    fn from(arg: &'ast syn::ExprIf) -> Self {
        BacktraceEvent::If(Ptr(arg))
    }
}
impl<'ast> From<&'ast syn::ExprForLoop> for BacktraceEvent<'ast> {
    fn from(arg: &'ast syn::ExprForLoop) -> Self {
        BacktraceEvent::ForLoop(Ptr(arg))
    }
}

impl ToTokens for BacktraceEvent<'_> {
    fn to_tokens(&self, toks: &mut proc_macro2::TokenStream) {
        match self {
            BacktraceEvent::FnCall(expr) => expr.to_tokens(toks),
            BacktraceEvent::MethodCall(expr) => expr.to_tokens(toks),
            BacktraceEvent::If(expr) => expr.to_tokens(toks),
            BacktraceEvent::ForLoop(expr) => expr.to_tokens(toks),
        }
    }
}

#[derive(Debug, Clone, Default)]
pub struct Backtrace<'ast> {
    events: Vec<BacktraceEvent<'ast>>,
    counts: HashMap<BacktraceEvent<'ast>, usize>,
}

impl<'ast> Backtrace<'ast> {
    fn push(&mut self, event: BacktraceEvent<'ast>) -> Result<(), Error> {
        self.events.push(event);
        let count = self.counts.entry(event).or_default();
        *count += 1;
        // event.is_call() is not necessary but leads to better error messages.
        if event.is_call() && *count >= 16 {
            Err(Error::new_spanned(event, "Recursion is not allowed."))
        } else {
            Ok(())
        }
    }
    fn pop(&mut self, event: BacktraceEvent<'ast>) {
        let popped = self.events.pop().expect("empty backtrace");
        assert_eq!(popped, event);
        let count = self.counts.entry(event).or_default();
        *count -= 1;
    }
    fn last_if(&self) -> Option<&'ast syn::ExprIf> {
        self.events.iter().rev().find_map(move |event| {
            if let BacktraceEvent::If(a) = event {
                Some(a.0)
            } else {
                None
            }
        })
    }
    fn iter(&self) -> impl Iterator<Item = &BacktraceEvent<'ast>> {
        self.events.iter().rev()
    }
    pub fn make_conditional_execution_error(
        &self,
        expr: &dyn ToTokens,
        fn_name: &'static str,
    ) -> Error {
        let mut err = Error::new_spanned(
            expr,
            format!("You cannot call {fn_name} in conditional execution."),
        );
        if let Some(expr_if) = self.last_if() {
            err.combine(Error::new_spanned(
                expr_if.if_token,
                format!("You cannot call {fn_name} in conditional execution (if location)."),
            ));
        }
        err
    }
}

impl std::fmt::Display for Backtrace<'_> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        for event in self.events.iter() {
            writeln!(f, "{}", event.to_token_stream())?;
        }
        Ok(())
    }
}

impl<'ast> Interpreter<'ast> {
    pub fn improve_error<T>(
        &self,
        res: Result<T, Error>,
        event: Option<&dyn ToTokens>,
    ) -> Result<T, Error> {
        res.map_err(|mut err| {
            for err_msg in err.iter_mut() {
                if !err_msg.is_span_set() {
                    err_msg
                        .msg
                        .push_str("\n  This error may have occurred in another file.");
                    err_msg
                        .msg
                        .push_str("\n  Error backtrace (up to this file):");
                    let mut event_iter = event
                        .into_iter()
                        .chain(self.backtrace.iter().map(|e| e as &dyn ToTokens));
                    while !err_msg.is_span_set() {
                        let Some(event) = event_iter.next() else {
                            break;
                        };
                        let extra_str = format!("\n  - `{}`", event.to_token_stream());
                        err_msg.msg.push_str(&extra_str);
                        err_msg.set_span(event.span());
                    }
                }
            }
            err
        })
    }
    pub fn with_backtrace<T>(
        &mut self,
        event: impl Into<BacktraceEvent<'ast>>,
        func: impl FnOnce(&mut Self) -> Result<T, Error>,
    ) -> Result<T, Error> {
        let event = event.into();
        self.backtrace.push(event)?;
        let res = func(self);
        let improved_res = self.improve_error(res, None);
        self.backtrace.pop(event);
        improved_res
    }
}
