use ecow::{EcoVec, eco_vec};
use typst_library::diag::{At, SourceResult, bail, error, warning};
use typst_library::engine::Engine;
use typst_library::foundations::{
    Array, Capturer, Closure, ClosureNode, Content, ContextElem, Dict, Func,
    NativeElement, Selector, Str, Value, ops,
};
use typst_library::introspection::{Counter, State};
use typst_syntax::Span;
use typst_syntax::ast::{self, AstNode};
use typst_utils::singleton;

use crate::{CapturesVisitor, Eval, FlowEvent, Vm};

impl Eval for ast::Code<'_> {
    type Output = Value;

    fn eval(self, vm: &mut Vm) -> SourceResult<Self::Output> {
        eval_code(vm, &mut self.exprs())
    }
}

/// Evaluate a stream of expressions.
fn eval_code<'a>(
    vm: &mut Vm,
    exprs: &mut impl Iterator<Item = ast::Expr<'a>>,
) -> SourceResult<Value> {
    let flow = vm.flow.take();
    let mut output = Value::None;

    while let Some(expr) = exprs.next() {
        let span = expr.span();
        let value = match expr {
            ast::Expr::SetRule(set) => {
                let styles = set.eval(vm)?;
                if vm.flow.is_some() {
                    break;
                }

                let tail = eval_code(vm, exprs)?.display();
                Value::Content(tail.styled_with_map(styles))
            }
            ast::Expr::ShowRule(show) => {
                let recipe = show.eval(vm)?;
                if vm.flow.is_some() {
                    break;
                }

                let tail = eval_code(vm, exprs)?.display();
                Value::Content(tail.styled_with_recipe(
                    &mut vm.engine,
                    vm.context,
                    recipe,
                )?)
            }
            _ => expr.eval(vm)?,
        };

        output = ops::join(output, value).at(span)?;

        if let Some(event) = &vm.flow {
            warn_for_discarded_content(&mut vm.engine, event, &output);
            break;
        }
    }

    if flow.is_some() {
        vm.flow = flow;
    }

    Ok(output)
}

impl Eval for ast::Expr<'_> {
    type Output = Value;

    fn eval(self, vm: &mut Vm) -> SourceResult<Self::Output> {
        let span = self.span();
        let forbidden = |name| {
            error!(span, "{} is only allowed directly in code and content blocks", name)
        };

        let value = match self {
            Self::Text(v) => v.eval(vm).map(Value::Content),
            Self::Space(v) => v.eval(vm).map(Value::Content),
            Self::Linebreak(v) => v.eval(vm).map(Value::Content),
            Self::Parbreak(v) => v.eval(vm).map(Value::Content),
            Self::Escape(v) => v.eval(vm),
            Self::Shorthand(v) => v.eval(vm),
            Self::SmartQuote(v) => v.eval(vm).map(Value::Content),
            Self::Strong(v) => v.eval(vm).map(Value::Content),
            Self::Emph(v) => v.eval(vm).map(Value::Content),
            Self::Raw(v) => v.eval(vm).map(Value::Content),
            Self::Link(v) => v.eval(vm).map(Value::Content),
            Self::Label(v) => v.eval(vm),
            Self::Ref(v) => v.eval(vm).map(Value::Content),
            Self::Heading(v) => v.eval(vm).map(Value::Content),
            Self::ListItem(v) => v.eval(vm).map(Value::Content),
            Self::EnumItem(v) => v.eval(vm).map(Value::Content),
            Self::TermItem(v) => v.eval(vm).map(Value::Content),
            Self::Equation(v) => v.eval(vm).map(Value::Content),
            Self::Math(v) => v.eval(vm).map(Value::Content),
            Self::MathText(v) => v.eval(vm).map(Value::Content),
            Self::MathIdent(v) => v.eval(vm),
            Self::MathFieldAccess(v) => v.eval(vm),
            Self::MathShorthand(v) => v.eval(vm),
            Self::MathAlignPoint(v) => v.eval(vm).map(Value::Content),
            Self::MathCall(v) => v.eval(vm),
            Self::MathDelimited(v) => v.eval(vm).map(Value::Content),
            Self::MathAttach(v) => v.eval(vm).map(Value::Content),
            Self::MathPrimes(v) => v.eval(vm).map(Value::Content),
            Self::MathFrac(v) => v.eval(vm).map(Value::Content),
            Self::MathRoot(v) => v.eval(vm).map(Value::Content),
            Self::Ident(v) => v.eval(vm),
            Self::None(v) => v.eval(vm),
            Self::Auto(v) => v.eval(vm),
            Self::Bool(v) => v.eval(vm),
            Self::Int(v) => v.eval(vm),
            Self::Float(v) => v.eval(vm),
            Self::Numeric(v) => v.eval(vm),
            Self::Str(v) => v.eval(vm),
            Self::CodeBlock(v) => v.eval(vm),
            Self::ContentBlock(v) => v.eval(vm).map(Value::Content),
            Self::Array(v) => v.eval(vm).map(Value::Array),
            Self::Dict(v) => v.eval(vm).map(Value::Dict),
            Self::Parenthesized(v) => v.eval(vm),
            Self::FieldAccess(v) => v.eval(vm),
            Self::FuncCall(v) => v.eval(vm),
            Self::Closure(v) => v.eval(vm),
            Self::Unary(v) => v.eval(vm),
            Self::Binary(v) => v.eval(vm),
            Self::LetBinding(v) => v.eval(vm),
            Self::DestructAssignment(v) => v.eval(vm),
            Self::SetRule(_) => bail!(forbidden("set")),
            Self::ShowRule(_) => bail!(forbidden("show")),
            Self::Contextual(v) => v.eval(vm).map(Value::Content),
            Self::Conditional(v) => v.eval(vm),
            Self::WhileLoop(v) => v.eval(vm),
            Self::ForLoop(v) => v.eval(vm),
            Self::ModuleImport(v) => v.eval(vm),
            Self::ModuleInclude(v) => v.eval(vm).map(Value::Content),
            Self::LoopBreak(v) => v.eval(vm),
            Self::LoopContinue(v) => v.eval(vm),
            Self::FuncReturn(v) => v.eval(vm),
        }?
        .spanned(span);

        // This satisfies the obligation to call `Vm::trace` for almost all
        // value-producing expressions!
        vm.trace_at(span, &value);

        Ok(value)
    }
}

impl Eval for ast::Ident<'_> {
    type Output = Value;

    fn eval(self, vm: &mut Vm) -> SourceResult<Self::Output> {
        let span = self.span();
        Ok(vm
            .scopes
            .get(&self)
            .at(span)?
            .read_checked((&mut vm.engine, span))
            .clone())
    }
}

impl Eval for ast::None<'_> {
    type Output = Value;

    fn eval(self, _: &mut Vm) -> SourceResult<Self::Output> {
        Ok(Value::None)
    }
}

impl Eval for ast::Auto<'_> {
    type Output = Value;

    fn eval(self, _: &mut Vm) -> SourceResult<Self::Output> {
        Ok(Value::Auto)
    }
}

impl Eval for ast::Bool<'_> {
    type Output = Value;

    fn eval(self, _: &mut Vm) -> SourceResult<Self::Output> {
        Ok(Value::Bool(self.get()))
    }
}

impl Eval for ast::Int<'_> {
    type Output = Value;

    fn eval(self, _: &mut Vm) -> SourceResult<Self::Output> {
        Ok(Value::Int(self.get()))
    }
}

impl Eval for ast::Float<'_> {
    type Output = Value;

    fn eval(self, _: &mut Vm) -> SourceResult<Self::Output> {
        Ok(Value::Float(self.get()))
    }
}

impl Eval for ast::Numeric<'_> {
    type Output = Value;

    fn eval(self, _: &mut Vm) -> SourceResult<Self::Output> {
        Ok(Value::numeric(self.get()))
    }
}

impl Eval for ast::Str<'_> {
    type Output = Value;

    fn eval(self, _: &mut Vm) -> SourceResult<Self::Output> {
        Ok(Value::Str(self.get().into()))
    }
}

impl Eval for ast::Array<'_> {
    type Output = Array;

    fn eval(self, vm: &mut Vm) -> SourceResult<Self::Output> {
        let mut items = self.items();

        let mut vec = EcoVec::with_capacity(items.size_hint().0);

        // We raise an error when one of the array items is the spread of a
        // dictionary. If _all_ of the array items are spreads of dictionaries,
        // the user probably wanted to write `(: ..dict_a, ..dict_b)` instead
        // to create a dictionary, not an array.
        let mut all_dict_spreads = true;

        while let Some(item) = items.next() {
            match item {
                ast::ArrayItem::Pos(expr) => {
                    all_dict_spreads = false;
                    vec.push(expr.eval(vm)?)
                }
                ast::ArrayItem::Spread(spread) => match spread.expr().eval(vm)? {
                    Value::None => {}
                    Value::Array(array) => {
                        all_dict_spreads = false;
                        vec.extend(array);
                    }
                    v @ Value::Dict(_)
                        if all_dict_spreads
                        // Lookahead to see whether remaining items are spreads
                        // of dicts
                        && items.all(|item| matches!(
                            item,
                            ast::ArrayItem::Spread(spread) if matches!(
                                spread.expr().eval(vm),
                                Ok(Value::Dict(_)),
                            ),
                        )) =>
                    {
                        let fixed =
                            self.to_untyped().clone().into_text().replacen("(", "(: ", 1);
                        bail!(
                            spread.span(), "cannot spread {} into array", v.ty();
                            hint: "add a colon to create a dictionary instead: `{fixed}`";
                        )
                    }
                    v => bail!(spread.span(), "cannot spread {} into array", v.ty()),
                },
            }
        }

        Ok(vec.into())
    }
}

impl Eval for ast::Dict<'_> {
    type Output = Dict;

    fn eval(self, vm: &mut Vm) -> SourceResult<Self::Output> {
        let mut map = indexmap::IndexMap::default();
        let mut invalid_keys = eco_vec![];

        for item in self.items() {
            match item {
                ast::DictItem::Named(named) => {
                    map.insert(named.name().get().clone().into(), named.expr().eval(vm)?);
                }
                ast::DictItem::Keyed(keyed) => {
                    let raw_key = keyed.key();
                    let key = raw_key.eval(vm)?;
                    let key =
                        key.cast::<Str>().at(raw_key.span()).unwrap_or_else(|errors| {
                            invalid_keys.extend(errors);
                            Str::default()
                        });
                    map.insert(key, keyed.expr().eval(vm)?);
                }
                ast::DictItem::Spread(spread) => match spread.expr().eval(vm)? {
                    Value::None => {}
                    Value::Dict(dict) => map.extend(dict),
                    v => bail!(spread.span(), "cannot spread {} into dictionary", v.ty()),
                },
            }
        }

        if !invalid_keys.is_empty() {
            return Err(invalid_keys);
        }

        Ok(map.into())
    }
}

impl Eval for ast::CodeBlock<'_> {
    type Output = Value;

    fn eval(self, vm: &mut Vm) -> SourceResult<Self::Output> {
        vm.scopes.enter();
        let output = self.body().eval(vm)?;
        vm.scopes.exit();
        Ok(output)
    }
}

impl Eval for ast::ContentBlock<'_> {
    type Output = Content;

    fn eval(self, vm: &mut Vm) -> SourceResult<Self::Output> {
        vm.scopes.enter();
        let content = self.body().eval(vm)?;
        vm.scopes.exit();
        Ok(content)
    }
}

impl Eval for ast::Parenthesized<'_> {
    type Output = Value;

    fn eval(self, vm: &mut Vm) -> SourceResult<Self::Output> {
        self.expr().eval(vm)
    }
}

impl Eval for ast::FieldAccess<'_> {
    type Output = Value;

    fn eval(self, vm: &mut Vm) -> SourceResult<Self::Output> {
        let target = self.target().eval(vm)?;
        let field = self.field();
        access_field(vm, target, field.as_str(), field.span())
    }
}

/// Access a field on a target value.
pub(crate) fn access_field(
    vm: &mut Vm,
    target: Value,
    field: &str,
    field_span: Span,
) -> SourceResult<Value> {
    let err = match target.field(field, (&mut vm.engine, field_span)).at(field_span) {
        Ok(value) => return Ok(value),
        Err(err) => err,
    };

    // Check whether this is a get rule field access.
    if let Value::Func(func) = &target
        && let Some(element) = func.to_element()
        && let Some(id) = element.field_id(field)
        && let styles = vm.context.styles().at(field_span)
        && let Ok(value) =
            element.field_from_styles(id, styles.as_ref().map(|&s| s).unwrap_or_default())
    {
        // Only validate the contextual styles once we know that this is indeed
        // a field from the style chain.
        let _ = styles?;
        return Ok(value);
    }

    Err(err)
}

impl Eval for ast::Contextual<'_> {
    type Output = Content;

    fn eval(self, vm: &mut Vm) -> SourceResult<Self::Output> {
        let body = self.body();

        // Collect captured variables.
        let captured = {
            let mut visitor = CapturesVisitor::new(Some(&vm.scopes), Capturer::Context);
            visitor.visit(body.to_untyped());
            visitor.finish()
        };

        // Define the closure.
        let closure = Closure {
            node: ClosureNode::Context(self.body().to_untyped().clone()),
            defaults: vec![],
            captured,
            num_pos_params: 0,
        };

        let func = Func::from(closure).spanned(body.span());
        Ok(ContextElem::new(func).pack().spanned(body.span()))
    }
}

/// Emits a warning when we discard content while returning unconditionally.
fn warn_for_discarded_content(engine: &mut Engine, event: &FlowEvent, joined: &Value) {
    let FlowEvent::Return(span, Some(_), false) = event else { return };
    let Value::Content(tree) = &joined else { return };

    let selector = singleton!(
        Selector,
        Selector::Or(eco_vec![State::select_any(), Counter::select_any()])
    );

    let mut warning = warning!(
        *span,
        "this return unconditionally discards the content before it";
        hint: "try omitting the `return` to automatically join all values";
    );

    if tree.query_first_naive(selector).is_some() {
        warning.hint("state/counter updates are content that must end up in the document to have an effect");
    }

    engine.sink.warn(warning);
}
