use std::{
    borrow::Cow,
    collections::{hash_map::Entry, BTreeMap, HashMap, HashSet, VecDeque},
    fmt::Display,
    iter,
};

use chumsky::{container::Container, span::Span as _};
use either::Either::{self, Left, Right};
use itertools::{Itertools, Position};
use maybe_sync::{dyn_maybe_send_sync, Rc};

use super::{
    ir::{Call, CallKind, Expression, Literal, TemplateContent, VariableIndex},
    Error, Template, Type,
};

use crate::{
    ast::{
        Command, Expression as AstExpression, IndentationMarker, Literal as AstLiteral,
        MemberAccess, StringLiteralComponent, Template as AstTemplate,
        TemplateBody as AstTemplateBody, TemplateContent as AstTemplateContent,
        TemplateContentEscapeSequence, TemplateLine, TemplateLineNewline, Type as AstType,
        TypeDefinition as AstTypeDefintion,
    },
    type_checking::ty::MaybeWeak,
    ContextWorker, ContextualizationError, Contextualizer, Escaper, Function, Span,
};

#[derive(Debug, Clone)]
struct Context<'s> {
    functions: HashMap<&'s str, CallKind>,
    type_definitions: HashMap<&'s str, (Either<MaybeWeak<Type>, AstType<'s>>, Span<'s>)>,
}

impl<'s> Context<'s> {
    // true, but it just fits better as a method on context
    #[allow(clippy::only_used_in_recursion)]
    /// `currently_defining`: inserts weak recursive types for these to avoid strong cycles
    fn convert_ast_type<Ctx: Contextualizer<Esc, Wrk>, Esc: Escaper, Wrk: ContextWorker<Esc>>(
        &mut self,
        ast: &AstType<'s>,
        ctxer: &Ctx,
        currently_defining: &HashSet<&'s str>,
    ) -> Type {
        match ast {
            AstType::Ident(base) => match *base {
                "int" => Type::Int,
                "bool" => Type::Bool,
                "float" => Type::Float,
                _ if self.type_definitions.contains_key(*base) => {
                    let ty = match self.type_definitions[*base].0.clone() {
                        Right(ref ast_ty) => {
                            let mut next_currently_defining = currently_defining.clone();
                            next_currently_defining.push(*base);

                            self.type_definitions.get_mut(*base).unwrap().0 =
                                Left(MaybeWeak::new(Type::Int)); // placeholder

                            let ty = self.convert_ast_type(ast_ty, ctxer, &next_currently_defining);

                            let res_ty = self.type_definitions[*base].0.clone().unwrap_left();
                            res_ty.replace(ty.clone());

                            res_ty
                        }
                        Left(ty) => ty.clone(),
                    };
                    Type::Recursive {
                        payload: if currently_defining.contains(base) {
                            ty.weak()
                        } else {
                            ty
                        },
                        name: (*base).to_string(),
                    }
                }
                txt => Type::Text(txt.to_string()),
            },
            AstType::Array(inner) => Type::Array(Some(Box::new(inner.0.as_ref().map_or_else(
                || Type::Text(ctxer.default_text_type()),
                |inner| self.convert_ast_type(inner.as_ref(), ctxer, currently_defining),
            )))),
            AstType::Tuple(inner) => Type::Tuple(
                inner
                    .iter()
                    .map(|(inner, _)| self.convert_ast_type(inner, ctxer, currently_defining))
                    .collect(),
            ),
            AstType::Options(options) => Type::Sum(
                options
                    .iter()
                    .map(|(inner, _)| self.convert_ast_type(inner, ctxer, currently_defining))
                    .collect(),
            ),
            AstType::Struct(inner) => Type::Struct(
                inner
                    .iter()
                    .map(|((name, _), (inner, _))| {
                        (
                            (*name).to_string(),
                            inner.as_ref().map_or_else(
                                || Type::Text(ctxer.default_text_type()),
                                |inner| self.convert_ast_type(inner, ctxer, currently_defining),
                            ),
                        )
                    })
                    .collect(),
            ),
        }
    }
}

#[derive(Debug, Clone)]
struct TemplateContext<'s, 'a> {
    parent: Option<&'a TemplateContext<'s, 'a>>,
    variable_index_offset: usize,
    variables: HashMap<&'s str, (VariableIndex, Span<'s>, Type)>,
}

impl<'s> TemplateContext<'s, '_> {
    fn add_variable(&mut self, name: &'s str, location: Span<'s>, ty: Type) -> VariableIndex {
        let id = self.variables.len() + self.variable_index_offset;
        self.variables.insert(name, (id, location, ty));

        id
    }

    /// get a variable by name in this or outer scopes
    fn get_variable(&self, name: &'s str) -> Option<(VariableIndex, Span<'s>, &Type)> {
        let mut next = Some(self);

        while let Some(current) = next {
            if let Some((index, span, ty)) = current.variables.get(name) {
                return Some((*index, *span, ty));
            }

            next = current.parent;
        }

        None
    }

    /// get a variables id by it's name in the current ctx only to allow for shadowing
    /// or if it's not present adds it to the inner scope
    fn try_assign_var<'c, Ctx: Contextualizer<Esc, Wrk>, Esc: Escaper, Wrk: ContextWorker<Esc>>(
        &mut self,
        name: &'s str,
        location: Span<'s>,
        ty: Type,
        ctxer: &'c Ctx,
    ) -> Result<(usize, Option<Esc>), Error<'s>> {
        let next_var_index = self.variables.len() + self.variable_index_offset;
        match self.variables.entry(name) {
            Entry::Occupied(mut occupied) => {
                let (old_id, original_location, ref mut original_type) = occupied.get_mut();

                let Ok(escer) = original_type.unify(&ty, ctxer) else {
                    return Err(Error::VariableAssignmentIncompatibleType {
                        var_name: name,
                        original_location: *original_location,
                        original_type: original_type.clone(),
                        second_assignment: location,
                        second_type: ty,
                    });
                };

                Ok((*old_id, escer))
            }
            Entry::Vacant(vacant) => {
                vacant.insert((next_var_index, location, ty));

                Ok((next_var_index, None))
            }
        }
    }
}

/// type checks the templates
pub(crate) fn type_check_templates<
    's,
    Ctx: Contextualizer<Esc, Wrk>,
    Esc: Escaper,
    Wrk: ContextWorker<Esc>,
>(
    templates: Vec<AstTemplate<'s>>,
    functions: HashMap<&'s str, Rc<dyn_maybe_send_sync!(Function)>>,
    type_definitions: Vec<(AstTypeDefintion<'s>, Span<'s>)>,
    ctxer: &Ctx,
) -> Result<(Vec<Template<Esc>>, TemplateNameMapping), Error<'s>> {
    let mut ctx = Context {
        functions: HashMap::with_capacity(functions.len() + templates.len()),
        type_definitions: HashMap::with_capacity(type_definitions.len()),
    };

    let mut seen: HashMap<&'s str, Span<'s>> = HashMap::with_capacity(type_definitions.len());

    for type_def in type_definitions {
        match seen.entry(type_def.0.name.0) {
            Entry::Occupied(old) => {
                return Err(Error::DuplicateTypeDefinition {
                    first_span: *old.get(),
                    second_span: type_def.1,
                });
            }
            Entry::Vacant(entry) => {
                entry.insert(type_def.1);

                ctx.type_definitions.insert(
                    type_def.0.name.0,
                    (Either::Right(type_def.0.ty.0), type_def.0.ty.1),
                );
            }
        }
    }

    // notice the order in which functions are added: functions first, templates second so that templates take priority

    ctx.functions.extend(
        functions
            .into_iter()
            .map(|(name, f)| (name, CallKind::Function(f))),
    );

    let tmpl_funcs = templates
        .iter()
        .enumerate()
        .map(|(idx, template)| {
            (
                template.name.0,
                CallKind::Template {
                    idx,
                    out_ty: template
                        .output_type
                        .0
                        .map_or_else(|| ctxer.default_text_type(), str::to_string),
                    parameters: template
                        .parameters
                        .iter()
                        .map(|(name, ty)| {
                            (
                                name.0.to_string(),
                                ty.0.as_ref().map_or_else(
                                    || Type::Text(ctxer.default_text_type()),
                                    |ty| ctx.convert_ast_type(ty, ctxer, &HashSet::new()),
                                ),
                            )
                        })
                        .collect(),
                },
            )
        })
        .collect_vec();

    ctx.functions.extend(tmpl_funcs);

    let name_mapping = templates
        .iter()
        .enumerate()
        .map(|(index, tmpl)| (tmpl.name.0.to_string(), index))
        .collect();

    let type_checked_templates: Vec<Template<Esc>> = templates
        .into_iter()
        .map(|tmpl| type_check_template(tmpl, &mut ctx, ctxer))
        .collect::<Result<_, _>>()?;

    Ok((type_checked_templates, name_mapping))
}

type TemplateNameMapping = HashMap<String, usize>;

fn type_check_template<'s, Ctx: Contextualizer<Esc, Wrk>, Esc: Escaper, Wrk: ContextWorker<Esc>>(
    tmpl: AstTemplate<'s>,
    ctx: &mut Context<'s>,
    ctxer: &Ctx,
) -> Result<Template<Esc>, Error<'s>> {
    let mut parameters = tmpl
        .parameters
        .iter()
        .map(|(name, (ty, location))| {
            (
                name.0,
                (
                    ty.as_ref().map_or_else(
                        || Type::Text(ctxer.default_text_type()),
                        |ty| ctx.convert_ast_type(ty, ctxer, &HashSet::new()),
                    ),
                    location,
                ),
            )
        })
        .collect_vec();

    let mut tmpl_ctx = TemplateContext {
        parent: None,
        variable_index_offset: 0,
        variables: HashMap::with_capacity(parameters.len()),
    };

    parameters.sort_unstable_by(|(name, (_, _)), (name2, (_, _))| name.cmp(name2));

    for (name, ty) in &parameters {
        tmpl_ctx.add_variable(name, *ty.1, ty.0.clone());
    }

    let (body, variable_count) = type_check_template_body(
        tmpl.body,
        (
            &tmpl
                .output_type
                .0
                .map_or_else(|| Cow::Owned(ctxer.default_text_type()), Cow::Borrowed),
            tmpl.output_type.1,
        ),
        ctx,
        tmpl_ctx,
        ctxer,
    )?;

    let type_checked_template = Template {
        body,
        params: Type::Struct(
            parameters
                .into_iter()
                .map(|(name, ty)| (name.to_string(), ty.0))
                .collect(),
        ),
        variable_count,
    };

    Ok(type_checked_template)
}

fn push_static<Esc: Escaper, Wrk: ContextWorker<Esc>>(
    content: &mut Vec<TemplateContent<Esc>>,
    wrker: &mut Wrk,
    s: &str,
) -> Result<(), ContextualizationError> {
    // reuse last static content if there is one
    match content.last_mut() {
        Some(TemplateContent::Static(str)) => str.push_str(s),
        _ => content.push(TemplateContent::Static(s.to_string())),
    }

    wrker.push_static(s)
}

#[allow(clippy::too_many_lines)]
#[allow(clippy::too_many_arguments)]
fn push_template_content<
    'b,
    's,
    Ctx: Contextualizer<Esc, Wrk>,
    Esc: Escaper,
    Wrk: ContextWorker<Esc>,
>(
    output_ty: &Cow<'s, str>,
    content: &mut Vec<TemplateContent<Esc>>,
    content_pos_and_span: (AstTemplateContent<'s>, Position, Span<'s>),
    line_indentation_and_newline: (&'b str, &TemplateLineNewline<'s>),
    ctx: &mut Context<'s>,
    tmpl_ctx: &mut TemplateContext<'s, 'b>,
    mut max_additional_variable_count: usize,
    ctxer: &Ctx,
    wrker: &mut Wrk,
) -> Result<usize, Error<'s>> {
    match content_pos_and_span.0 {
        AstTemplateContent::Static(str) => {
            push_static(content, wrker, str)
                .map_err(|err| Error::Contextualization(err, content_pos_and_span.2))?;
        }

        AstTemplateContent::EscapeSequence(TemplateContentEscapeSequence::Normal(seq)) => {
            push_static(content, wrker, seq.str())
                .map_err(|err| Error::Contextualization(err, content_pos_and_span.2))?;
        }

        AstTemplateContent::EscapeSequence(TemplateContentEscapeSequence::Curly(
            _kind,
            _count,
            braces_str,
            _span,
        )) => push_static(content, wrker, braces_str)
            .map_err(|err| Error::Contextualization(err, content_pos_and_span.2))?,

        AstTemplateContent::Command(cmd) => match *cmd {
            Command::Expression(expr, span) => {
                let ((expression, ty), additional_var_count): ((Expression<Esc>, Type), usize) =
                    type_check_expression(
                        ctx,
                        tmpl_ctx,
                        expr,
                        span,
                        Some(&if matches!(content_pos_and_span.1, Position::Only) {
                            Type::Array(Some(Box::new(Type::Text(output_ty.clone().into_owned()))))
                        } else {
                            Type::Text(output_ty.clone().into_owned())
                        }),
                        ctxer,
                    )?;

                max_additional_variable_count =
                    max_additional_variable_count.max(additional_var_count);

                let Some(text_ty) = ty.to_printable() else {
                    return Err(Error::BodyExpressionInvalidType(ty, span));
                };

                let escer = wrker
                    .dynamic(&text_ty.unwrap_or_else(|| ctxer.default_text_type()))
                    .map_err(|err| Error::Contextualization(err, span))?;

                content.push(TemplateContent::BodyExpression {
                    indentation: line_indentation_and_newline.0.to_string(),
                    expression: Expression::Escaped(Box::new(expression), escer),
                    newline: match line_indentation_and_newline.1 {
                        TemplateLineNewline::Normal(inner) => (*inner).to_string(),
                        TemplateLineNewline::Escaped => String::new(),
                    },
                });
            }
            Command::Assignment {
                variable,
                expression,
            } => {
                let ((expr, ty), additional_vars) = type_check_expression(
                    ctx,
                    tmpl_ctx,
                    expression.0,
                    expression.1,
                    tmpl_ctx
                        .variables
                        .get(variable.0)
                        .map(|(_, _, ty)| ty.clone())
                        .as_ref(),
                    ctxer,
                )?;

                max_additional_variable_count = max_additional_variable_count.max(additional_vars);

                let (var_index, escer) =
                    tmpl_ctx.try_assign_var(variable.0, expression.1, ty, ctxer)?;

                content.push(TemplateContent::Assignment(
                    var_index,
                    escer.map_or(expr.clone(), |escer| {
                        Expression::Escaped(Box::new(expr), escer)
                    }),
                ));
            }
        },
    }

    Ok(max_additional_variable_count)
}

fn body_base_indentation<'s>(body: &AstTemplateBody<'s>) -> Option<(&'s str, Span<'s>)> {
    body.lines
        .iter()
        // we need to filter out completely blank lines since editors typically strip out trailing whitespace and we don't want empty lines to dictate minimum indentation
        .filter(|(line, _)| {
            !(line.indentation.0.is_empty()
                && line.content.is_empty()
                && matches!(line.newline.0, TemplateLineNewline::Normal(_)))
        })
        .map(|(line, _)| line.indentation)
        .min_by_key(|indentation| indentation.0.len())
}

fn body_lines_with_indentation_marker<'s>(
    lines: Vec<(TemplateLine<'s>, Span<'s>)>,
    indentation_marker: &IndentationMarker<'s>,
) -> impl Iterator<Item = (TemplateLine<'s>, Span<'s>)> + Clone {
    if let IndentationMarker::Escaped {
        indentation_and_span,
        marker_span,
        newline: newline_and_span,
    } = indentation_marker
    {
        Left(
            iter::once({
                (
                    // in case of an escaped indentation marker create an artifical template line to account for it and prepend it
                    TemplateLine {
                        had_comments: false,
                        indentation: *indentation_and_span,
                        content: vec![(AstTemplateContent::Static(">"), *marker_span)],
                        newline: (
                            TemplateLineNewline::Normal(newline_and_span.0),
                            newline_and_span.1,
                        ),
                    },
                    indentation_and_span.1.union(*marker_span),
                )
            })
            .chain(lines),
        )
    } else {
        Right(lines.into_iter())
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// An amount of indentation. Typically of a template / template literal line.
pub struct Indentation {
    spaces: usize,
    tabs: usize,
}

impl From<&str> for Indentation {
    fn from(value: &str) -> Self {
        let chars = value.chars();
        Self {
            spaces: chars.clone().filter(|c| *c == ' ').count(),
            tabs: chars.filter(|c| *c == '\t').count(),
        }
    }
}

impl Indentation {
    fn remaining<'a>(&self, s: &'a str) -> Option<&'a str> {
        let mut self_remaining = *self;
        for (i, c) in s.char_indices() {
            match c {
                ' ' if self_remaining.spaces > 0 => {
                    self_remaining.spaces -= 1;
                }
                '\t' if self_remaining.tabs > 0 => {
                    self_remaining.tabs -= 1;
                }
                _ => return None,
            }

            if self_remaining.spaces == 0 && self_remaining.tabs == 0 {
                let start = i + c.len_utf8();
                return Some(if s.len() >= start {
                    &s[start..s.len()]
                } else {
                    ""
                });
            }
        }

        None
    }
}

#[cfg(test)]
#[allow(clippy::items_after_test_module)]
mod test_indentation {
    use super::Indentation;

    #[test]
    fn test_from_str() {
        assert_eq!(Indentation { spaces: 0, tabs: 0 }, "".into());
        assert_eq!(Indentation { spaces: 2, tabs: 2 }, "  \t\t".into());
        assert_eq!(Indentation { spaces: 3, tabs: 3 }, " \t  \t\t".into());
    }

    #[test]
    fn test_remaining() {
        assert_eq!(None, Indentation { spaces: 2, tabs: 0 }.remaining(" "));
        assert_eq!(
            None,
            Indentation { spaces: 2, tabs: 2 }.remaining(" \t  \t")
        );
        assert_eq!(
            Some(""),
            Indentation { spaces: 4, tabs: 0 }.remaining("    ")
        );
        assert_eq!(
            Some("\t "),
            Indentation { spaces: 2, tabs: 3 }.remaining(" \t \t\t\t ")
        );
    }
}

impl Display for Indentation {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self.spaces {
            0 => Ok(()),
            1 => f.write_str("1 space"),
            _ => write!(f, "{} spaces", self.spaces),
        }?;

        if self.spaces != 0 && self.tabs != 0 {
            f.write_str(" and ")?;
        }

        match self.tabs {
            0 => Ok(()),
            1 => f.write_str("1 tab"),
            _ => write!(f, "{} tabs", self.tabs),
        }
    }
}

fn body_line_remaining_indentation<'s>(
    indentation_marker: &IndentationMarker<'s>,
    base_indentation: (&'s str, Span<'s>),
    line: &TemplateLine<'s>,
) -> Result<&'s str, Error<'s>> {
    // treat completely empty lines as if they had the necessary indentation; editors
    // usually strip indentation in empty lines
    if line.indentation.0.is_empty()
        && line.content.is_empty()
        && matches!(line.newline.0, TemplateLineNewline::Normal(_))
    {
        return Ok("");
    }
    match indentation_marker {
        IndentationMarker::None | IndentationMarker::Escaped { .. } => {
            let ind: Indentation = base_indentation.0.into();
            ind.remaining(line.indentation.0)
                .ok_or(Error::InconsistentMinimumIndentation(
                    (ind, base_indentation.1),
                    (line.indentation.0.into(), line.indentation.1),
                ))
        }
        IndentationMarker::Present {
            indentation_and_span,
        } => {
            let ind: Indentation = indentation_and_span.0.into();
            ind.remaining(line.indentation.0)
                .ok_or(Error::NotEnoughIndentation {
                    indentation_marker_indentation_and_span: (ind, line.indentation.1),
                    insufficent_indentation_and_span: (
                        line.indentation.0.into(),
                        line.indentation.1,
                    ),
                })
        }
    }
}

fn type_check_template_body<
    's,
    'a,
    Ctx: Contextualizer<Esc, Wrk>,
    Esc: Escaper,
    Wrk: ContextWorker<Esc>,
>(
    body: AstTemplateBody<'s>,
    ty_and_span: (&Cow<'s, str>, Span<'s>),
    ctx: &'a mut Context<'s>,
    mut tmpl_ctx: TemplateContext<'s, 'a>,
    ctxer: &Ctx,
) -> Result<(Vec<TemplateContent<Esc>>, usize), Error<'s>> {
    let mut max_additional_variable_count = 0;
    // todo: possible optimization: we could compute the neccessary content count here and use as initial capacity; at least an upper bound
    let mut content = vec![];

    let mut wrker = ctxer
        .contextualize(ty_and_span.0)
        .map_err(|err| Error::Contextualization(err, ty_and_span.1))?;

    let Some(base_indentation) = body_base_indentation(&body) else {
        return Ok((content, tmpl_ctx.variables.len()));
    };

    let lines = body_lines_with_indentation_marker(body.lines, &body.indentation_marker);

    for (pos, (line, line_span)) in lines
        // ignore comment only lines
        .filter(|(line, _span)| !(line.content.is_empty() && line.had_comments))
        .with_position()
    {
        let remaining_indentation =
            body_line_remaining_indentation(&body.indentation_marker, base_indentation, &line)?;

        if !remaining_indentation.is_empty() {
            push_static(&mut content, &mut wrker, remaining_indentation)
                .map_err(|err| Error::Contextualization(err, line_span))?;
        }

        let sole_assignment = line.content.len() == 1
            && matches!(
                &line.content[0].0,
                AstTemplateContent::Command(cmd) if matches!(**cmd, Command::Assignment { .. })
            );

        for (content_pos, (next, content_span)) in line.content.into_iter().with_position() {
            max_additional_variable_count =
                max_additional_variable_count.max(push_template_content(
                    ty_and_span.0,
                    &mut content,
                    (next, content_pos, content_span),
                    (remaining_indentation, &line.newline.0),
                    ctx,
                    &mut tmpl_ctx,
                    max_additional_variable_count,
                    ctxer,
                    &mut wrker,
                )?);
        }

        // do not add the newline if the line soly consists of a comment or assignment; also not if
        // this is the last line in the body
        if sole_assignment || matches!(pos, Position::Only | Position::Last) {
            continue;
        }

        let newline = match line.newline.0 {
            TemplateLineNewline::Normal(newline) => newline,
            TemplateLineNewline::Escaped => "",
        };

        content.push(TemplateContent::OriginalTemplateNewline(
            newline.to_string(),
        ));

        wrker.push_static(newline).map_err(|err| {
            Error::Contextualization(
                err,
                Span {
                    start: line_span.end,
                    ..line_span
                },
            )
        })?;
    }

    Ok((
        content,
        tmpl_ctx.variables.len() + max_additional_variable_count,
    ))
}

/// type checks this expression and tries to match the target type. Inability to do so here is not
/// considered an error.
#[allow(clippy::too_many_lines)]
fn type_check_expression<
    's,
    Esc: Escaper,
    Wrk: ContextWorker<Esc>,
    Ctx: Contextualizer<Esc, Wrk>,
>(
    ctx: &mut Context<'s>,
    tmpl_ctx: &mut TemplateContext<'s, '_>,
    expr: AstExpression<'s>,
    span: Span<'s>,
    target_ty: Option<&Type>,
    ctxer: &Ctx,
) -> Result<((Expression<Esc>, Type), usize), Error<'s>> {
    Ok(match expr {
        AstExpression::Literal(literal) => type_check_literal(ctx, tmpl_ctx, literal, ctxer)?,

        AstExpression::VariableOrCall(name) => {
            type_check_variable_or_call(ctx, tmpl_ctx, name, span)?
        }

        AstExpression::Call { func, arguments } => {
            type_check_call(ctx, tmpl_ctx, &func, arguments, span, ctxer)?
        }

        AstExpression::If {
            condition,
            then_expression,
            else_expression,
        } => type_check_if(
            ctx,
            tmpl_ctx,
            condition,
            then_expression,
            else_expression,
            span,
            ctxer,
        )?,

        AstExpression::Each {
            variable,
            array_expression,
            body_expression,
        } => type_check_each(
            ctx,
            tmpl_ctx,
            variable,
            array_expression,
            body_expression,
            target_ty,
            ctxer,
        )?,

        AstExpression::Select { expr, arms } => {
            type_check_select(ctx, tmpl_ctx, expr, arms, span, target_ty, ctxer)?
        }

        AstExpression::MemberAccess { expression, access } => {
            let ((expr, mut ty), extra_vars) =
                type_check_expression(ctx, tmpl_ctx, *expression.0, expression.1, None, ctxer)?;

            ty.substitute_recursive_once();

            match (ty.clone(), access.0.clone()) {
                (Type::Tuple(ref inner_types), MemberAccess::Numeric(index, _access_span)) => {
                    let index = index as usize;
                    let Some(ty) = inner_types.get(index) else {
                        return Err(Error::InvalidTupleAccess {
                            access_span: access.1,
                            tuple_ty: ty.clone(),
                            index,
                        });
                    };

                    (
                        (
                            Expression::MemberAccess {
                                expression: Box::new(expr),
                                access: index,
                            },
                            ty.clone(),
                        ),
                        extra_vars,
                    )
                }
                (Type::Struct(inner_names_and_types), MemberAccess::Named(name, access_span)) => {
                    let Some((ty, index)) = inner_names_and_types.into_iter().enumerate().find_map(
                        |(index, (inner_key, ty))| (inner_key == name).then_some((ty, index)),
                    ) else {
                        return Err(Error::InvalidStructAccess {
                            access_span,
                            struct_ty: ty,
                            name,
                        });
                    };

                    (
                        (
                            Expression::MemberAccess {
                                expression: Box::new(expr),
                                access: index,
                            },
                            ty,
                        ),
                        extra_vars,
                    )
                }
                _ => {
                    return Err(Error::InvalidMemberAccess {
                        expr_span: span,
                        expr_ty: ty,
                        access: access.0,
                    })
                }
            }
        }
    })
}

fn type_check_literal<'s, Ctx: Contextualizer<Esc, Wrk>, Esc: Escaper, Wrk: ContextWorker<Esc>>(
    ctx: &mut Context<'s>,
    tmpl_ctx: &mut TemplateContext<'s, '_>,
    literal: AstLiteral<'s>,
    ctxer: &Ctx,
) -> Result<((Expression<Esc>, Type), usize), Error<'s>> {
    Ok(match literal {
        AstLiteral::Int(value) => ((Expression::Literal(Literal::Int(value)), Type::Int), 0),
        AstLiteral::Float(value) => ((Expression::Literal(Literal::Float(value)), Type::Float), 0),
        AstLiteral::Bool(value) => ((Expression::Literal(Literal::Bool(value)), Type::Bool), 0),
        AstLiteral::String(string_literal_components) => {
            let mut value = String::new();

            for component in string_literal_components {
                match component.0 {
                    StringLiteralComponent::EscapedQuote => value.push('"'),
                    StringLiteralComponent::EscapeSequence(sequence) => {
                        value.push_str(sequence.str());
                    }
                    StringLiteralComponent::Literal(str) => value.push_str(str),
                }
            }

            (
                (
                    Expression::Literal(Literal::String(value)),
                    Type::Text(ctxer.default_text_type()),
                ),
                0,
            )
        }
        AstLiteral::Template(output_type, body) => {
            let inner_tmpl_ctx = TemplateContext {
                parent: Some(tmpl_ctx),
                variable_index_offset: tmpl_ctx.variables.len() + tmpl_ctx.variable_index_offset,
                variables: HashMap::new(),
            };

            let out_ty = output_type
                .0
                .map_or_else(|| Cow::Owned(ctxer.default_text_type()), Cow::Borrowed);

            let (lines, var_count) = type_check_template_body(
                *body,
                (&out_ty.clone(), output_type.1),
                ctx,
                inner_tmpl_ctx,
                ctxer,
            )?;

            (
                (
                    Expression::Literal(Literal::Template(lines)),
                    Type::Array(Some(Box::new(Type::Text(out_ty.into_owned())))),
                ),
                var_count,
            )
        }
        AstLiteral::Tuple(expressions) => {
            let mut max_additional_vars = 0;

            let mut exprs = Vec::with_capacity(expressions.len());
            let mut types = Vec::with_capacity(expressions.len());

            for expression in expressions {
                let ((expr, ty), additional_vars) =
                    type_check_expression(ctx, tmpl_ctx, expression.0, expression.1, None, ctxer)?;

                exprs.push(expr);
                types.push(ty);

                max_additional_vars = max_additional_vars.max(additional_vars);
            }

            (
                (
                    Expression::Literal(Literal::ArrayOrTuple(exprs)),
                    Type::Tuple(types),
                ),
                max_additional_vars,
            )
        }
        AstLiteral::Array(expressions) => {
            type_check_array_literal(ctx, tmpl_ctx, expressions, ctxer)?
        }
        AstLiteral::Struct(expressions) => {
            let mut max_additional_vars = 0;

            let (names_and_expressions, names_and_types): (BTreeMap<_, _>, BTreeMap<_, _>) =
                expressions
                    .into_iter()
                    .map(|(name, expression)| {
                        let ((expr, ty), additional_vars) = type_check_expression(
                            ctx,
                            tmpl_ctx,
                            expression.0,
                            expression.1,
                            None,
                            ctxer,
                        )?;

                        max_additional_vars = max_additional_vars.max(additional_vars);

                        Ok(((name.0.to_string(), expr), (name.0.to_string(), ty)))
                    })
                    .process_results(|iter| iter.unzip())?;

            (
                (
                    Expression::Literal(Literal::Struct(names_and_expressions)),
                    Type::Struct(names_and_types),
                ),
                max_additional_vars,
            )
        }
    })
}

fn type_check_array_literal<
    's,
    Ctx: Contextualizer<Esc, Wrk>,
    Esc: Escaper,
    Wrk: ContextWorker<Esc>,
>(
    ctx: &mut Context<'s>,
    tmpl_ctx: &mut TemplateContext<'s, '_>,
    expressions: Vec<(AstExpression<'s>, Span<'s>)>,
    ctxer: &Ctx,
) -> Result<((Expression<Esc>, Type), usize), Error<'s>> {
    let mut expressions = VecDeque::from(expressions);
    let Some(first) = expressions.pop_front() else {
        return Ok((
            (
                Expression::Literal(Literal::ArrayOrTuple(vec![])),
                Type::Array(None),
            ),
            0,
        ));
    };

    let ((first_expr, first_ty), first_additional_vars) =
        type_check_expression(ctx, tmpl_ctx, first.0, first.1, None, ctxer)?;

    let mut exprs = Vec::with_capacity(expressions.len());
    exprs.push(first_expr);

    let mut array_ty = first_ty.clone();

    let mut max_additional_vars = first_additional_vars;
    let mut last_expression_span = first.1;

    for expression in expressions {
        let ((expr, ty), additional_vars) = type_check_expression(
            ctx,
            tmpl_ctx,
            expression.0,
            expression.1,
            Some(&first_ty),
            ctxer,
        )?;

        let Ok(escer) = array_ty.unify(&ty, ctxer) else {
            return Err(Error::ArrayLiteralIncompatibleTypes {
                first_span: last_expression_span,
                first_ty: array_ty,
                second_span: expression.1,
                second_ty: ty,
            });
        };

        exprs.push(escer.map_or(expr.clone(), |escer| {
            Expression::Escaped(Box::new(expr), escer)
        }));

        max_additional_vars = max_additional_vars.max(additional_vars);
        last_expression_span = expression.1;
    }

    Ok((
        (
            Expression::Literal(Literal::ArrayOrTuple(exprs)),
            Type::Array(Some(Box::new(array_ty))),
        ),
        max_additional_vars,
    ))
}

fn type_check_variable_or_call<'s, Esc: Escaper>(
    ctx: &mut Context<'s>,
    tmpl_ctx: &mut TemplateContext<'s, '_>,
    name: &'s str,
    span: Span<'s>,
) -> Result<((Expression<Esc>, Type), usize), Error<'s>> {
    Ok(
        match (tmpl_ctx.get_variable(name), ctx.functions.get(name)) {
            // prioritize variables even if a function with the same name is available to enable shadowing and backwards compatibility
            (Some((var_index, _span, ty)), _) => ((Expression::Variable(var_index), ty.clone()), 0),
            (None, Some(call_kind)) => match call_kind {
                CallKind::Function(f) => {
                    let out_ty = f.accepts(Vec::new()).ok_or(Error::CallWrongArguments {
                        call_span: span,
                        name: Cow::Borrowed(name),
                    })?;

                    (
                        (
                            Expression::Call(Call::Function(f.clone(), Vec::new())),
                            out_ty,
                        ),
                        0,
                    )
                }
                CallKind::Template {
                    idx,
                    out_ty,
                    parameters,
                } => {
                    if !parameters.is_empty() {
                        return Err(Error::CallWrongArguments {
                            call_span: span,
                            name: Cow::Borrowed(name),
                        });
                    }

                    (
                        (
                            Expression::Call(Call::Template(*idx, Vec::new())),
                            Type::Array(Some(Box::new(Type::Text(out_ty.clone())))),
                        ),
                        0,
                    )
                }
            },
            _ => return Err(Error::UnknownVariableOrFunction(span)),
        },
    )
}

fn type_check_call<'s, Ctx: Contextualizer<Esc, Wrk>, Esc: Escaper, Wrk: ContextWorker<Esc>>(
    ctx: &mut Context<'s>,
    tmpl_ctx: &mut TemplateContext<'s, '_>,
    func: &(Box<AstExpression<'s>>, Span<'s>),
    arguments: Vec<(AstExpression<'s>, Span<'s>)>,
    span: Span<'s>,
    ctxer: &Ctx,
) -> Result<((Expression<Esc>, Type), usize), Error<'s>> {
    let AstExpression::VariableOrCall(name) = *func.0 else {
        return Err(Error::CallInvalidFunction(func.1));
    };

    let Some(call_kind) = ctx.functions.get(name).cloned() else {
        return Err(Error::UnknownFunction(func.1));
    };

    Ok(match call_kind {
        CallKind::Function(f) => {
            let mut max_var_count = 0;

            let (expressions, tys): (Vec<_>, Vec<_>) = arguments
                .iter()
                .cloned()
                .map(|expression| {
                    let ((expression, mut got_ty), additional_var_count) =
                        type_check_expression(ctx, tmpl_ctx, expression.0, span, None, ctxer)?;

                    got_ty.substitute_recursive_once();

                    max_var_count = max_var_count.max(additional_var_count);

                    Ok::<_, Error<'s>>((expression, got_ty))
                })
                .try_fold(
                    (Vec::new(), Vec::new()),
                    |(mut expressions, mut tys), res| {
                        let (expression, ty) = res?;
                        expressions.push(expression);
                        tys.push(ty);

                        Ok((expressions, tys))
                    },
                )?;

            let out_ty = f.accepts(tys).ok_or(Error::CallWrongArguments {
                call_span: span,
                name: Cow::Borrowed(name),
            })?;

            (
                (Expression::Call(Call::Function(f, expressions)), out_ty),
                max_var_count,
            )
        }
        CallKind::Template {
            parameters,
            out_ty,
            idx,
        } => {
            if arguments.len() != parameters.len() {
                return Err(Error::CallWrongArguments {
                    name: Cow::Borrowed(name),
                    call_span: span,
                });
            }

            let mut max_var_count = 0;

            let argument_expressions = arguments
                .into_iter()
                .zip(parameters)
                .map(|(expression, (name, mut expected_ty))| {
                    let ((mut expression, got_ty), additional_var_count) = type_check_expression(
                        ctx,
                        tmpl_ctx,
                        expression.0,
                        expression.1,
                        Some(&expected_ty),
                        ctxer,
                    )?;

                    if let Some(escer) = expected_ty.unify(&got_ty, ctxer).map_err(|()| {
                        Error::CallWrongArguments {
                            name: Cow::Owned(name.clone()),
                            call_span: span,
                        }
                    })? {
                        expression = Expression::Escaped(Box::new(expression), escer);
                    }

                    max_var_count = max_var_count.max(additional_var_count);

                    Ok((name.clone(), expression))
                })
                .collect::<Result<_, _>>()?;
            (
                (
                    Expression::Call(Call::Template(idx, argument_expressions)),
                    Type::Array(Some(Box::new(Type::Text(out_ty.clone())))),
                ),
                max_var_count,
            )
        }
    })
}

fn type_check_if<'s, Ctx: Contextualizer<Esc, Wrk>, Esc: Escaper, Wrk: ContextWorker<Esc>>(
    ctx: &mut Context<'s>,
    tmpl_ctx: &mut TemplateContext<'s, '_>,
    condition: (Box<AstExpression<'s>>, Span<'s>),
    then_expression: (Box<AstExpression<'s>>, Span<'s>),
    else_expression: Option<(Box<AstExpression<'s>>, Span<'s>)>,
    span: Span<'s>,
    ctxer: &Ctx,
) -> Result<((Expression<Esc>, Type), usize), Error<'s>> {
    let ((condition_expr, condition_ty), condition_extra_vars) = type_check_expression(
        ctx,
        tmpl_ctx,
        *condition.0,
        condition.1,
        Some(&Type::Bool),
        ctxer,
    )?;

    if condition_ty != Type::Bool {
        return Err(Error::IfExpressionConditionNotBoolType {
            condition_span: condition.1,
            condition_ty,
        });
    }

    let ((then_expr, then_ty), then_extra_vars) = type_check_expression(
        ctx,
        tmpl_ctx,
        *then_expression.0,
        then_expression.1,
        None,
        ctxer,
    )?;

    Ok(match (else_expression, then_ty) {
        (None, ty @ Type::Text(_)) => (
            (
                Expression::If {
                    condition: Box::new(condition_expr),
                    then: Box::new(then_expr),
                    otherwise: Box::new(Expression::Literal(Literal::String(String::new()))),
                },
                ty,
            ),
            condition_extra_vars.max(then_extra_vars),
        ),
        (None, ty) => return Err(Error::IfExpressionMissingElseExpression { if_span: span, ty }),
        (Some(else_expression), then_ty) => {
            let ((else_expr, mut else_ty), else_extra_vars) = type_check_expression(
                ctx,
                tmpl_ctx,
                *else_expression.0,
                else_expression.1,
                Some(&then_ty),
                ctxer,
            )?;

            let Ok(escer) = else_ty.unify(&then_ty, ctxer) else {
                return Err(Error::IfExpressionThenAndElseIncompatibleType {
                    then_span: then_expression.1,
                    then_ty,
                    else_span: else_expression.1,
                    else_ty,
                });
            };

            (
                (
                    Expression::If {
                        condition: Box::new(condition_expr),
                        then: Box::new(then_expr),
                        otherwise: Box::new(escer.map_or(else_expr.clone(), |escer| {
                            Expression::Escaped(Box::new(else_expr), escer)
                        })),
                    },
                    else_ty,
                ),
                condition_extra_vars
                    .max(then_extra_vars)
                    .max(else_extra_vars),
            )
        }
    })
}

fn type_check_each<'s, Ctx: Contextualizer<Esc, Wrk>, Esc: Escaper, Wrk: ContextWorker<Esc>>(
    ctx: &mut Context<'s>,
    tmpl_ctx: &mut TemplateContext<'s, '_>,
    variable: (&'s str, Span<'s>),
    array_expression: (Box<AstExpression<'s>>, Span<'s>),
    body_expression: (Box<AstExpression<'s>>, Span<'s>),
    target_ty: Option<&Type>,
    ctxer: &Ctx,
) -> Result<((Expression<Esc>, Type), usize), Error<'s>> {
    let ((array_expr, array_ty), array_extra_vars) = type_check_expression(
        ctx,
        tmpl_ctx,
        *array_expression.0,
        array_expression.1,
        None,
        ctxer,
    )?;

    let var_ty = match array_ty {
        // int is a placeholder type; it's concrete value is irrelevant since it will never be used because an array literal has to be empty for this to occur
        Type::Array(inner) => inner.map_or(Type::Int, |inner| *inner),
        _ => {
            return Err(Error::EachExpressionNotArray {
                expression_span: array_expression.1,
                expression_ty: array_ty,
            })
        }
    };

    // new template context to hold the iteration variable
    let var_index = tmpl_ctx.variables.len() + tmpl_ctx.variable_index_offset;

    let mut inner_tmpl_ctx = TemplateContext {
        parent: Some(tmpl_ctx),
        variable_index_offset: var_index,
        variables: HashMap::from_iter([(variable.0, (var_index, variable.1, var_ty))]),
    };

    let body_ty = if let Some(Type::Array(Some(ref inner))) = target_ty {
        Some(*inner.clone())
    } else {
        None
    };

    let ((body_expr, body_ty), body_extra_vars) = type_check_expression(
        ctx,
        &mut inner_tmpl_ctx,
        *body_expression.0,
        body_expression.1,
        body_ty.as_ref(),
        ctxer,
    )?;

    Ok((
        (
            Expression::Each {
                assignment_variable: var_index,
                array_expression: Box::new(array_expr),
                body: Box::new(body_expr),
            },
            Type::Array(Some(Box::new(body_ty))),
        ),
        array_extra_vars.max(1 + body_extra_vars),
    ))
}

#[allow(clippy::type_complexity)]
fn type_check_select<'s, Ctx: Contextualizer<Esc, Wrk>, Esc: Escaper, Wrk: ContextWorker<Esc>>(
    ctx: &mut Context<'s>,
    tmpl_ctx: &mut TemplateContext<'s, '_>,
    expr: (Box<AstExpression<'s>>, Span<'s>),
    arms: Vec<(
        (&'s str, Span<'s>),
        (AstType<'s>, Span<'s>),
        (AstExpression<'s>, Span<'s>),
    )>,
    total_span: Span<'s>,
    target_ty: Option<&Type>,
    ctxer: &Ctx,
) -> Result<((Expression<Esc>, Type), usize), Error<'s>> {
    let ((select_expr, mut select_ty), select_expr_extra_vars) =
        type_check_expression(ctx, tmpl_ctx, *expr.0, expr.1, None, ctxer)?;

    select_ty.substitute_recursive_once();

    let Type::Sum(ref options) = select_ty else {
        return Err(Error::NonSumTypeInSelect {
            source_expr_span: expr.1,
            source_expr_ty: select_ty,
        });
    };

    #[allow(clippy::mutable_key_type)]
    let mut options: HashSet<_> = options.iter().collect();

    let mut total_arm_ty: Option<(Type, Span<'s>)> = None;

    let var_index = tmpl_ctx.variables.len() + tmpl_ctx.variable_index_offset;

    let mut total_arms_extra_vars = 0;

    let arms: Vec<(Type, Expression<Esc>)> = arms
        .into_iter()
        .map(|(ident, ty, expr)| {
            let ty = (ctx.convert_ast_type(&ty.0, ctxer, &HashSet::new()), ty.1);

            if !options.remove(&ty.0) {
                return Err(Error::SelectArmDuplicateOrInvalidType {
                    arm_ty: ty.0,
                    arm_ty_span: ty.1,
                    source_expr_span: expr.1,
                    source_expr_ty: select_ty.clone(),
                });
            }

            let mut inner_tmpl_ctx = TemplateContext {
                parent: Some(tmpl_ctx),
                variable_index_offset: var_index,
                variables: HashMap::from_iter([(ident.0, (var_index, ident.1, ty.0.clone()))]),
            };

            let ((mut arm_expr, arm_ty), arm_extra_vars) =
                type_check_expression(ctx, &mut inner_tmpl_ctx, expr.0, expr.1, target_ty, ctxer)?;

            total_arms_extra_vars = total_arms_extra_vars.max(arm_extra_vars);

            let (prev_ty, prev_span) = total_arm_ty.get_or_insert((arm_ty.clone(), expr.1));

            if let Some(escer) =
                prev_ty
                    .unify(&arm_ty, ctxer)
                    .map_err(|()| Error::SelectArmsIncompatibleTypes {
                        first_expr_span: *prev_span,
                        first_expr_ty: prev_ty.clone(),
                        second_expr_span: expr.1,
                        second_expr_ty: arm_ty,
                    })?
            {
                arm_expr = Expression::Escaped(Box::new(arm_expr), escer);
            }

            Ok((ty.0, arm_expr))
        })
        .collect::<Result<Vec<_>, _>>()?;

    let Some((res_ty, _)) = total_arm_ty else {
        return Err(Error::SelectNoArms {
            select_span: total_span,
        });
    };

    let is_text = matches!(&res_ty, Type::Text(_));

    if !is_text && !options.is_empty() {
        return Err(Error::SelectNotAllOptionsCovered {
            source_expr_span: expr.1,
            source_expr_ty: select_ty.clone(),

            remaining: options.into_iter().cloned().collect(),
        });
    }

    Ok((
        (
            Expression::Select {
                select_expression: Box::new(select_expr),
                empty_text_fallback: is_text,
                variable: var_index,
                arms,
            },
            res_ty,
        ),
        select_expr_extra_vars.max(1 + total_arms_extra_vars),
    ))
}
