use std::collections::BTreeMap;

use crate::ast::TemplateAstExpr;
use crate::input::NomoInput;
use crate::parser::TemplateToken;
use crate::parser::TokenOperator;
use crate::value::NomoValue;

pub struct EmitMachine {
    current_index: usize,
    labels: BTreeMap<LabelSlot, usize>,
}

impl EmitMachine {
    fn reserve_slot(&mut self) -> VariableSlot {
        VariableSlot {
            index: {
                let val = self.current_index;
                self.current_index += 1;
                val
            },
        }
    }

    fn reserve_label(&mut self) -> LabelSlot {
        LabelSlot {
            index: {
                let val = self.current_index;
                self.current_index += 1;
                val
            },
        }
    }

    fn assign_label(&mut self, slot: LabelSlot, idx: usize) {
        let no_prev = self.labels.insert(slot, idx).is_none();

        assert!(no_prev, "A label slot was already assigned")
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct VariableSlot {
    index: usize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct LabelSlot {
    index: usize,
}

#[derive(Debug, Clone)]
pub enum Instruction {
    AppendContent {
        content: NomoInput,
    },
    LoadFromContextToSlot {
        name: NomoInput,
        slot: VariableSlot,
    },
    EmitFromSlot {
        slot: VariableSlot,
    },
    PushScope {
        inherit_parent: bool,
    },
    Abort,
    JumpIfNotTrue {
        emit_slot: VariableSlot,
        jump: LabelSlot,
    },
    Jump {
        jump: LabelSlot,
    },
    NoOp,
    CreateIteratorFromSlotToSlot {
        iterator_slot: VariableSlot,
        iterator_source_slot: VariableSlot,
    },
    AdvanceIteratorOrJump {
        iterator_slot: VariableSlot,
        value_slot: VariableSlot,
        jump: LabelSlot,
    },
    GetIteratorEmptyOrJump {
        iterator_slot: VariableSlot,
        jump: LabelSlot,
    },
    PopScope,
    LoadFromSlotToContext {
        value_ident: NomoInput,
        value_slot: VariableSlot,
    },
    LoadLiteralToSlot {
        source: TemplateToken,
        value: NomoValue,
        slot: VariableSlot,
    },
    MathOperate {
        op: TokenOperator,
        left_slot: VariableSlot,
        right_slot: VariableSlot,
        result_slot: VariableSlot,
    },
    FunctionCall {
        name: NomoInput,
        args: Vec<VariableSlot>,
        slot: VariableSlot,
    },
}

#[derive(Debug, Clone)]
pub struct VMInstructions {
    pub labels: BTreeMap<LabelSlot, usize>,
    pub instructions: Vec<Instruction>,
}

pub fn emit_machine(input: crate::ast::TemplateAst<'_>) -> VMInstructions {
    let mut eval = vec![];

    let mut machine = EmitMachine {
        current_index: 0,
        labels: BTreeMap::new(),
    };

    for ast in input.root() {
        emit_ast_expr(&mut machine, &mut eval, ast);
    }

    VMInstructions {
        labels: machine.labels,
        instructions: eval,
    }
}

fn emit_ast_expr(
    machine: &mut EmitMachine,
    eval: &mut Vec<Instruction>,
    ast: &TemplateAstExpr<'_>,
) {
    match ast {
        TemplateAstExpr::StaticContent(template_token) => {
            eval.push(Instruction::AppendContent {
                content: template_token.source().clone(),
            });
        }
        TemplateAstExpr::Interpolation {
            prev_whitespace_content,
            expression,
            post_whitespace_content,
        } => {
            if let Some(ws) = prev_whitespace_content {
                eval.push(Instruction::AppendContent {
                    content: ws.source().clone(),
                });
            }

            let emit_slot = machine.reserve_slot();
            emit_expr_load(machine, eval, emit_slot, expression);
            eval.push(Instruction::EmitFromSlot { slot: emit_slot });

            if let Some(ws) = post_whitespace_content {
                eval.push(Instruction::AppendContent {
                    content: ws.source().clone(),
                });
            }
        }
        TemplateAstExpr::ConditionalChain { chain } => {
            let mut chain = chain.iter();

            let end_label = machine.reserve_label();
            let mut end_indices = vec![];

            let mut previous_post_whitespace_content: &Option<crate::parser::TemplateToken> = &None;
            let mut previous_jump: Option<LabelSlot> = None;

            loop {
                let next = chain.next().unwrap();
                if let Some(ws) = previous_post_whitespace_content {
                    eval.push(Instruction::AppendContent {
                        content: ws.source().clone(),
                    });
                }

                if let TemplateAstExpr::ConditionalContent { content } = &next {
                    for ast in content {
                        emit_ast_expr(machine, eval, ast);
                    }

                    end_indices.push(eval.len());
                    eval.push(Instruction::Jump { jump: end_label });
                } else if let TemplateAstExpr::Block {
                    prev_whitespace_content,
                    post_whitespace_content,
                    expression,
                } = &next
                {
                    previous_post_whitespace_content = post_whitespace_content;
                    if let Some(ws) = prev_whitespace_content {
                        let idx = end_indices.last().copied();
                        eval.insert(
                            idx.unwrap_or(eval.len()),
                            Instruction::AppendContent {
                                content: ws.source().clone(),
                            },
                        );
                        if let Some(idx) = end_indices.last_mut() {
                            *idx += 1;
                        }
                    }

                    if let TemplateAstExpr::IfConditional { expression } = &**expression {
                        let emit_slot = machine.reserve_slot();
                        emit_expr_load(machine, eval, emit_slot, expression);

                        let jmp_label = machine.reserve_label();
                        previous_jump = Some(jmp_label);
                        eval.push(Instruction::JumpIfNotTrue {
                            emit_slot,
                            jump: jmp_label,
                        });
                    } else if let TemplateAstExpr::ElseConditional { expression } = &**expression {
                        if let Some(previous_jump) = previous_jump.take() {
                            machine.assign_label(previous_jump, eval.len());
                        } else {
                            panic!("Got an else without a previous if?");
                        }

                        if let Some(expression) = expression {
                            let emit_slot = machine.reserve_slot();
                            emit_expr_load(machine, eval, emit_slot, expression);

                            let jmp_label = machine.reserve_label();
                            previous_jump = Some(jmp_label);
                            eval.push(Instruction::JumpIfNotTrue {
                                emit_slot,
                                jump: jmp_label,
                            });
                        } else {
                            // We don't have to do anything in the else case
                        }
                    } else if let TemplateAstExpr::EndBlock = &**expression {
                        break;
                    }
                }
            }

            if let Some(previous_jump) = previous_jump.take() {
                machine.assign_label(previous_jump, eval.len());
            }
            machine.assign_label(end_label, eval.len());

            if let Some(ws) = previous_post_whitespace_content {
                eval.push(Instruction::AppendContent {
                    content: ws.source().clone(),
                });
            } else {
                eval.push(Instruction::NoOp);
            }
        }
        TemplateAstExpr::ForChain {
            for_block,
            content,
            else_block,
            else_content,
            end_block,
        } => {
            let post_for_whitespace_content;
            let label_to_else_or_empty_index = machine.reserve_label();
            let label_to_end_index = machine.reserve_label();
            let label_start_loop = machine.reserve_label();
            if let TemplateAstExpr::Block {
                prev_whitespace_content,
                expression,
                post_whitespace_content,
            } = &**for_block
                && let TemplateAstExpr::For {
                    value_ident,
                    value_expression,
                } = &**expression
            {
                if let Some(ws) = prev_whitespace_content {
                    eval.push(Instruction::AppendContent {
                        content: ws.source().clone(),
                    });
                }
                post_for_whitespace_content = post_whitespace_content;

                eval.push(Instruction::PushScope {
                    inherit_parent: true,
                });

                let value_slot = machine.reserve_slot();
                let iterator_source_slot = machine.reserve_slot();
                let iterator_slot = machine.reserve_slot();

                emit_expr_load(machine, eval, iterator_source_slot, value_expression);
                eval.push(Instruction::CreateIteratorFromSlotToSlot {
                    iterator_source_slot,
                    iterator_slot,
                });

                eval.push(Instruction::GetIteratorEmptyOrJump {
                    iterator_slot,
                    jump: label_to_else_or_empty_index,
                });

                machine.assign_label(label_start_loop, eval.len());
                eval.push(Instruction::AdvanceIteratorOrJump {
                    iterator_slot,
                    value_slot,
                    jump: label_to_end_index,
                });

                eval.push(Instruction::LoadFromSlotToContext {
                    value_slot,
                    value_ident: value_ident.source(),
                });
            } else {
                panic!("For block should be a for block");
            };

            if let Some(ws) = post_for_whitespace_content {
                eval.push(Instruction::AppendContent {
                    content: ws.source().clone(),
                });
            }

            for content in content {
                emit_ast_expr(machine, eval, content);
            }

            let end_of_content_jump = eval.len();
            eval.push(Instruction::Jump {
                jump: label_start_loop,
            });

            let has_else = else_block.is_some();

            if let Some(TemplateAstExpr::Block {
                prev_whitespace_content,
                expression,
                post_whitespace_content,
            }) = else_block.as_deref()
                && let TemplateAstExpr::ForElse = &**expression
            {
                if let Some(ws) = prev_whitespace_content {
                    eval.insert(
                        end_of_content_jump.saturating_sub(1),
                        Instruction::AppendContent {
                            content: ws.source().clone(),
                        },
                    );
                }

                machine.assign_label(label_to_else_or_empty_index, eval.len());

                if let Some(ws) = post_whitespace_content {
                    eval.push(Instruction::AppendContent {
                        content: ws.source().clone(),
                    });
                }

                for content in else_content
                    .as_ref()
                    .expect("If there is a for block, there should be for content (even if empty)")
                {
                    emit_ast_expr(machine, eval, content);
                }
            }

            let post_end_whitespace_content;

            if let TemplateAstExpr::Block {
                prev_whitespace_content,
                expression,
                post_whitespace_content,
            } = &**end_block
                && let TemplateAstExpr::EndBlock = &**expression
            {
                post_end_whitespace_content = post_whitespace_content;

                if let Some(ws) = prev_whitespace_content {
                    if has_else {
                        eval.push(Instruction::AppendContent {
                            content: ws.source().clone(),
                        });
                    } else {
                        eval.insert(
                            end_of_content_jump.saturating_sub(1),
                            Instruction::AppendContent {
                                content: ws.source().clone(),
                            },
                        );
                    }
                }

                if !has_else {
                    machine.assign_label(label_to_else_or_empty_index, eval.len());
                }

                machine.assign_label(label_to_end_index, eval.len());

                eval.push(Instruction::PopScope);

                if let Some(ws) = post_end_whitespace_content {
                    eval.push(Instruction::AppendContent {
                        content: ws.source().clone(),
                    });
                }
            } else {
                panic!("End block should be an endblock");
            }
        }

        TemplateAstExpr::Block { .. }
        | TemplateAstExpr::EndBlock
        | TemplateAstExpr::IfConditional { .. }
        | TemplateAstExpr::ConditionalContent { .. }
        | TemplateAstExpr::ElseConditional { .. }
        | TemplateAstExpr::For { .. }
        | TemplateAstExpr::ForElse
        | TemplateAstExpr::Invalid { .. }
        | TemplateAstExpr::Literal { .. }
        | TemplateAstExpr::FunctionCall { .. }
        | TemplateAstExpr::Operation { .. }
        | TemplateAstExpr::VariableAccess { .. } => eval.push(Instruction::Abort),
    }
}

fn emit_expr_load(
    machine: &mut EmitMachine,
    eval: &mut Vec<Instruction>,
    emit_slot: VariableSlot,
    expression: &TemplateAstExpr<'_>,
) {
    match expression {
        TemplateAstExpr::VariableAccess(template_token) => {
            eval.push(Instruction::LoadFromContextToSlot {
                name: template_token.source().clone(),
                slot: emit_slot,
            });
        }
        TemplateAstExpr::Literal { source, value } => {
            eval.push(Instruction::LoadLiteralToSlot {
                source: source.clone(),
                value: value.clone(),
                slot: emit_slot,
            });
        }
        TemplateAstExpr::Operation { op, lhs, rhs } => {
            let left_slot = machine.reserve_slot();
            emit_expr_load(machine, eval, left_slot, lhs);
            let right_slot = machine.reserve_slot();
            emit_expr_load(machine, eval, right_slot, rhs);

            eval.push(Instruction::MathOperate {
                op: *op,
                left_slot,
                right_slot,
                result_slot: emit_slot,
            });
        }
        TemplateAstExpr::FunctionCall { name, args } => {
            let mut arg_slots = vec![];
            for arg in args {
                let slot = machine.reserve_slot();
                emit_expr_load(machine, eval, slot, arg);
                arg_slots.push(slot);
            }

            eval.push(Instruction::FunctionCall {
                name: name.source(),
                args: arg_slots,
                slot: emit_slot,
            });
        }
        TemplateAstExpr::Invalid { .. } => eval.push(Instruction::Abort),
        TemplateAstExpr::StaticContent { .. } | TemplateAstExpr::Interpolation { .. } => {
            unreachable!("Invalid AST here")
        }
        TemplateAstExpr::ConditionalChain { .. } => todo!(),
        TemplateAstExpr::ElseConditional { .. } => todo!(),
        TemplateAstExpr::EndBlock => todo!(),
        TemplateAstExpr::Block { .. } => todo!(),
        TemplateAstExpr::ForChain { .. } => todo!(),
        TemplateAstExpr::For { .. } => todo!(),
        TemplateAstExpr::ForElse => todo!(),

        TemplateAstExpr::IfConditional { .. } => todo!(),
        TemplateAstExpr::ConditionalContent { .. } => todo!(),
    }
}

#[cfg(test)]
mod tests {
    use crate::emit::emit_machine;

    #[test]
    fn check_simple_variable_interpolation() {
        let input = "Hello {{= world }}";

        let parsed = crate::parser::parse(input.into()).unwrap();

        let ast = crate::ast::parse(parsed.tokens()).unwrap();

        let emit = emit_machine(ast);

        insta::assert_debug_snapshot!(emit, @r#"
        VMInstructions {
            labels: {},
            instructions: [
                AppendContent {
                    content: "Hello" (0..5),
                },
                AppendContent {
                    content: " " (5..6),
                },
                LoadFromContextToSlot {
                    name: "world" (10..15),
                    slot: VariableSlot {
                        index: 0,
                    },
                },
                EmitFromSlot {
                    slot: VariableSlot {
                        index: 0,
                    },
                },
            ],
        }
        "#);
    }

    #[test]
    fn check_if_else_if() {
        let input = "{{ if foo }} foo {{ else if bar }} bar {{ else }} foobar {{ end }}";

        let parsed = crate::parser::parse(input.into()).unwrap();

        let ast = crate::ast::parse(parsed.tokens()).unwrap();

        let emit = emit_machine(ast);

        insta::assert_debug_snapshot!(emit);
    }

    #[test]
    fn check_function_call() {
        let input = "{{ if foo(23) }} bar {{ else }} foobar {{ end }}";

        let parsed = crate::parser::parse(input.into()).unwrap();

        let ast = crate::ast::parse(parsed.tokens()).unwrap();

        let emit = emit_machine(ast);

        insta::assert_debug_snapshot!(emit);
    }
}