use crate::{
    parser::{Tree, TreeKind, TreeRef},
    semantics::{Location, Semantics, Type},
    tokens::TokenKind,
};

// TODO: If I were cool this would by actual bytecode.
//       But I'm not cool.
pub enum Instruction {
    Panic,

    BoolNot,
    Discard,
    FloatAdd,
    FloatDivide,
    FloatMultiply,
    FloatSubtract,
    Jump(usize),
    JumpFalse(usize),
    JumpTrue(usize),
    LoadArgument(usize),
    LoadLocal(usize),
    LoadModule(usize),
    PushFalse,
    PushFloat(f64),
    PushNothing,
    PushString(usize),
    PushTrue,
    StoreLocal(usize),
}

pub struct Function {
    instructions: Vec<Instruction>,
    strings: Vec<String>,
}

type CR = Option<()>;
const OK: CR = CR::Some(());

pub fn compile_expression(code: &mut Function, semantics: &Semantics, t: TreeRef) {
    let tree = &semantics.tree()[t];
    let cr = match tree.kind {
        TreeKind::Error => None,
        TreeKind::LiteralExpression => compile_literal(code, semantics, t, tree),
        TreeKind::GroupingExpression => compile_grouping(code, semantics, tree),
        TreeKind::UnaryExpression => compile_unary_operator(code, semantics, tree),
        TreeKind::ConditionalExpression => compile_condition_expression(code, semantics, tree),
        TreeKind::BinaryExpression => compile_binary_expression(code, semantics, tree),
        TreeKind::Identifier => compile_identifier_expression(code, semantics, t, tree),
        TreeKind::CallExpression => todo!(),
        TreeKind::Block => compile_block_expression(code, semantics, tree),
        _ => {
            semantics.internal_compiler_error(Some(t), "tree is not an expression, cannot compile")
        }
    };
    if matches!(cr, None) {
        code.instructions.push(Instruction::Panic);
    }
}

fn compile_literal(code: &mut Function, semantics: &Semantics, t: TreeRef, tr: &Tree) -> CR {
    let tok = tr.nth_token(0)?;
    match semantics.type_of(t) {
        Type::F64 => code
            .instructions
            .push(Instruction::PushFloat(tok.as_str().parse().unwrap())),
        Type::Bool => code.instructions.push(if tok.kind == TokenKind::True {
            Instruction::PushTrue
        } else {
            Instruction::PushFalse
        }),
        Type::String => {
            let index = code.strings.len();

            // TODO: Interpret string here make good!
            let mut result = String::new();
            let mut input = tok.as_str().chars();
            while let Some(ch) = input.next() {
                if ch == '\\' {
                    if let Some(ch) = input.next() {
                        match ch {
                            'n' => result.push('\n'),
                            'r' => result.push('\r'),
                            't' => result.push('\t'),
                            _ => result.push(ch),
                        }
                    } else {
                        result.push(ch)
                    }
                } else {
                    result.push(ch)
                }
            }
            code.strings.push(result);

            code.instructions.push(Instruction::PushString(index))
        }
        Type::Error => code.instructions.push(Instruction::Panic),
        _ => panic!("unsupported literal type: {t:?}"),
    };
    OK
}

fn compile_grouping(code: &mut Function, semantics: &Semantics, t: &Tree) -> CR {
    compile_expression(code, semantics, t.nth_tree(1)?);
    OK
}

fn compile_unary_operator(code: &mut Function, semantics: &Semantics, t: &Tree) -> CR {
    compile_expression(code, semantics, t.nth_tree(1)?);

    let tok = t.nth_token(0)?;
    match tok.kind {
        TokenKind::Minus => {
            code.instructions.push(Instruction::PushFloat(-1.0));
            code.instructions.push(Instruction::FloatMultiply);
        }
        TokenKind::Bang => {
            code.instructions.push(Instruction::BoolNot);
        }
        _ => panic!("unsupported unary operator"),
    }
    OK
}

fn compile_condition_expression(code: &mut Function, semantics: &Semantics, t: &Tree) -> CR {
    let condition = t.nth_tree(1)?;
    compile_expression(code, semantics, condition);

    let jump_else_index = code.instructions.len();
    code.instructions.push(Instruction::JumpFalse(0));

    let then_branch = t.nth_tree(2)?;
    compile_expression(code, semantics, then_branch);

    if let Some(else_branch) = t.nth_tree(4) {
        let jump_end_index = code.instructions.len();
        code.instructions.push(Instruction::Jump(0));

        let else_index = code.instructions.len();
        code.instructions[jump_else_index] = Instruction::JumpFalse(else_index);

        compile_expression(code, semantics, else_branch);

        let end_index = code.instructions.len();
        code.instructions[jump_end_index] = Instruction::Jump(end_index);
    } else {
        let else_index = code.instructions.len();
        code.instructions[jump_else_index] = Instruction::JumpFalse(else_index);
    }
    OK
}

fn compile_binary_expression(code: &mut Function, semantics: &Semantics, t: &Tree) -> CR {
    compile_expression(code, semantics, t.nth_tree(0)?);
    match t.nth_token(1)?.kind {
        TokenKind::Plus => {
            compile_expression(code, semantics, t.nth_tree(2)?);
            code.instructions.push(Instruction::FloatAdd);
        }
        TokenKind::Minus => {
            compile_expression(code, semantics, t.nth_tree(2)?);
            code.instructions.push(Instruction::FloatSubtract);
        }
        TokenKind::Star => {
            compile_expression(code, semantics, t.nth_tree(2)?);
            code.instructions.push(Instruction::FloatMultiply);
        }
        TokenKind::Slash => {
            compile_expression(code, semantics, t.nth_tree(2)?);
            code.instructions.push(Instruction::FloatDivide);
        }
        TokenKind::And => {
            let jump_false_index = code.instructions.len();
            code.instructions.push(Instruction::JumpFalse(0));
            code.instructions.push(Instruction::PushTrue);

            let jump_end_index = code.instructions.len();
            code.instructions.push(Instruction::Jump(0));

            let false_index = code.instructions.len();
            code.instructions[jump_false_index] = Instruction::JumpFalse(false_index);

            compile_expression(code, semantics, t.nth_tree(2)?);

            let end_index = code.instructions.len();
            code.instructions[jump_end_index] = Instruction::Jump(end_index);
        }
        TokenKind::Or => {
            let jump_true_index = code.instructions.len();
            code.instructions.push(Instruction::JumpTrue(0));
            code.instructions.push(Instruction::PushTrue);

            let jump_end_index = code.instructions.len();
            code.instructions.push(Instruction::Jump(0));

            let true_index = code.instructions.len();
            code.instructions[jump_true_index] = Instruction::JumpTrue(true_index);

            compile_expression(code, semantics, t.nth_tree(2)?);

            let end_index = code.instructions.len();
            code.instructions[jump_end_index] = Instruction::Jump(end_index);
        }
        _ => panic!("Unsupported binary expression"),
    }
    OK
}

fn compile_identifier_expression(
    code: &mut Function,
    semantics: &Semantics,
    t: TreeRef,
    tree: &Tree,
) -> Option<()> {
    let ident = tree.nth_token(0)?;
    let environment = semantics.environment_of(t);
    let declaration = environment.bind(ident)?;

    let instruction = match declaration.location {
        Location::Local => Instruction::LoadLocal(declaration.index),
        Location::Argument => Instruction::LoadArgument(declaration.index),
        Location::Module => Instruction::LoadModule(declaration.index),
    };
    code.instructions.push(instruction);

    OK
}

fn compile_block_expression(code: &mut Function, semantics: &Semantics, tree: &Tree) -> Option<()> {
    let last_is_brace = tree.nth_token(tree.children.len() - 1).is_some();
    let last_index = tree.children.len() - if last_is_brace { 2 } else { 1 };

    for i in 1..last_index {
        compile_statement(code, semantics, tree.nth_tree(i)?, false);
    }
    compile_statement(code, semantics, tree.nth_tree(last_index)?, true);
    OK
}

pub fn compile_statement(code: &mut Function, semantics: &Semantics, t: TreeRef, gen_value: bool) {
    let tree = &semantics.tree()[t];
    let cr = match tree.kind {
        TreeKind::FunctionDecl => compile_function_declaration(code, semantics, tree, gen_value),
        TreeKind::LetStatement => compile_let_statement(code, semantics, t, tree, gen_value),
        TreeKind::ExpressionStatement => {
            compile_expression_statement(code, semantics, tree, gen_value)
        }
        TreeKind::IfStatement => compile_if_statement(code, semantics, tree, gen_value),
        _ => panic!("unsupported tree kind {:?}", tree.kind),
    };
    if matches!(cr, None) {
        code.instructions.push(Instruction::Panic);
    }
}

fn compile_if_statement(
    code: &mut Function,
    semantics: &Semantics,
    tree: &Tree,
    gen_value: bool,
) -> CR {
    compile_expression(code, semantics, tree.nth_tree(0)?);
    if !gen_value {
        code.instructions.push(Instruction::Discard);
    }

    OK
}

fn compile_expression_statement(
    code: &mut Function,
    semantics: &Semantics,
    tree: &Tree,
    gen_value: bool,
) -> CR {
    compile_expression(code, semantics, tree.nth_tree(0)?);
    if tree
        .nth_token(1)
        .is_some_and(|t| t.kind == TokenKind::Semicolon)
    {
        code.instructions.push(Instruction::Discard);
        if gen_value {
            code.instructions.push(Instruction::PushNothing);
        }
    } else if !gen_value {
        code.instructions.push(Instruction::Discard);
    }

    OK
}

fn compile_let_statement(
    code: &mut Function,
    semantics: &Semantics,
    t: TreeRef,
    tree: &Tree,
    gen_value: bool,
) -> CR {
    compile_expression(code, semantics, tree.nth_tree(3)?);
    let environment = semantics.environment_of(t);
    let declaration = environment.bind(tree.nth_token(1)?)?;

    // NOTE: Because this is a let statement I assume it's local!
    assert!(matches!(declaration.location, Location::Local));
    code.instructions
        .push(Instruction::StoreLocal(declaration.index));
    if gen_value {
        code.instructions.push(Instruction::PushNothing);
    }

    OK
}

fn compile_function_declaration(
    _code: &mut Function,
    _semantics: &Semantics,
    _tree: &Tree,
    _gen_value: bool,
) -> CR {
    todo!()
}