mathic 0.2.0

use std::collections::HashMap;

use crate::{
    diagnostics::LoweringError,
    lowering::ir::{
        basic_block::Terminator,
        function::{FunctionBuilder, LocalKind},
        instruction::{InitInstruct, LValInstruct, RValInstruct, RValueKind},
        symbols::TypeIndex,
        types::{FloatTy, MathicType, NumericTy, SintTy, UintTy, lower_inner_ast_type},
        value::{ConstExpr, NumericConst, Value, ValueModifier},
    },
    parser::{
        Span,
        ast::expression::{BinaryOp, ExprStmt, ExprStmtKind, LogicalOp, PrimaryExpr, UnaryOp},
    },
};

pub fn lower_expr(
    func: &mut FunctionBuilder,
    expr: &ExprStmt,
    ty_hint: Option<TypeIndex>,
) -> Result<(RValInstruct, TypeIndex), LoweringError> {
    let rvalue = match &expr.kind {
        ExprStmtKind::Primary(val) => lower_primary_value(func, val, expr.span, ty_hint)?,
        ExprStmtKind::Binary { lhs, op, rhs } => lower_binary_op(func, lhs, *op, rhs, expr.span)?,
        ExprStmtKind::Unary { op, rhs } => lower_unary_op(func, *op, rhs, expr.span, ty_hint)?,
        ExprStmtKind::Group(expr) => {
            return lower_expr(func, expr, ty_hint);
        }
        ExprStmtKind::Call { callee, args } => {
            if callee == "eval" {
                return lower_eval_builtin(func, args, expr.span);
            }
            lower_call(func, callee.clone(), args, expr.span)?
        }
        ExprStmtKind::Assign {
            name,
            expr: assign_expr,
        } => lower_assignment(func, name, assign_expr, expr.span)?,
        ExprStmtKind::Logical { lhs, op, rhs } => lower_logical_op(func, lhs, *op, rhs, expr.span)?,
        ExprStmtKind::StructInit { name, fields } => lower_adt_init(func, name, fields, expr.span)?,
        ExprStmtKind::Index { .. } => todo!(),
        ExprStmtKind::StructGet {
            expr: struct_expr,
            field_name,
        } => lower_struct_get(func, struct_expr, field_name, expr.span, ty_hint)?,
        ExprStmtKind::StructSet {
            lhs,
            field_name,
            rhs,
        } => lower_struct_set(func, lhs, field_name, rhs, expr.span)?,
    };

    Ok((
        rvalue,
        lower_expression_type(func, &expr.kind, ty_hint, expr.span)?,
    ))
}

fn lower_assignment(
    func: &mut FunctionBuilder,
    name: &str,
    expr: &ExprStmt,
    span: Span,
) -> Result<RValInstruct, LoweringError> {
    let local = func.sym_table.get_local_from_name(name, span)?;
    let (value, expr_ty_idx) = lower_expr(func, expr, Some(local.ty))?;
    // The new value should be of the same type as the local's.
    if local.ty != expr_ty_idx {
        return Err(LoweringError::MismatchedType {
            expected: func.get_type(local.ty, span)?,
            found: func.get_type(expr_ty_idx, span)?,
            span,
        });
    }

    func.get_basic_block_mut(func.last_block_idx())
        .instructions
        .push(LValInstruct::Assign {
            local_idx: local.local_idx,
            value,
            modifier: vec![],
            span: Some(span),
        });

    Ok(RValInstruct {
        kind: RValueKind::Use {
            value: Value::Const(ConstExpr::Void),
            span: None,
        },
        ty: func.get_or_insert_global_type_idx(MathicType::Void),
    })
}

fn lower_call(
    func: &mut FunctionBuilder,
    callee: String,
    func_args: &[ExprStmt],
    span: Span,
) -> Result<RValInstruct, LoweringError> {
    let mut arg_values: Vec<RValInstruct> = Vec::new();
    let func_prototype = func.get_function_decl(&callee, span)?;

    if func_prototype.params.len() != func_args.len() {
        return Err(LoweringError::WrongArgumentCount {
            name: callee.to_string(),
            expected: func_prototype.params.len(),
            got: func_args.len(),
            span,
        });
    }

    for (arg, param) in func_args.iter().zip(func_prototype.params.iter()) {
        let param_ty_idx = lower_inner_ast_type(func, &param.ty, param.span)?;
        let (arg_val, arg_ty_idx) = lower_expr(func, arg, Some(param_ty_idx))?;

        if arg_ty_idx != param_ty_idx {
            return Err(LoweringError::MismatchedType {
                expected: func.get_type(param_ty_idx, param.span)?,
                found: func.get_type(arg_ty_idx, param.span)?,
                span: arg.span,
            });
        }

        arg_values.push(arg_val);
    }

    // Since we represent function calls as expressions, we need to return a
    // value. Due to the fact that function calls are actually terminators and
    // not RValue instructions, we need to create a temporary local to store
    // the return value and then create the RValue instruction pointing to that
    // new local.
    let return_ty_idx = match func_prototype.return_ty {
        Some(ty) => lower_inner_ast_type(func, &ty, span)?,
        None => func.get_or_insert_global_type_idx(MathicType::Void),
    };
    let local_idx = func
        .sym_table
        .add_local(None, return_ty_idx, None, LocalKind::Temp)?;

    let dest_block_idx = func.last_block_idx() + 1;

    func.get_basic_block_mut(func.last_block_idx()).terminator = Terminator::Call {
        callee,
        args: arg_values,
        span: Some(span),
        return_dest: Value::InMemory {
            local_idx,
            modifier: vec![],
        },
        return_ty: return_ty_idx,
        dest_block: dest_block_idx,
    };

    func.add_block(Terminator::Return(None, None), None);

    Ok(RValInstruct {
        kind: RValueKind::Use {
            value: Value::InMemory {
                local_idx,
                modifier: vec![],
            },
            span: None,
        },
        ty: return_ty_idx,
    })
}

fn lower_eval_builtin(
    func: &mut FunctionBuilder,
    func_args: &[ExprStmt],
    span: Span,
) -> Result<(RValInstruct, TypeIndex), LoweringError> {
    if func_args.len() != 3 {
        return Err(LoweringError::WrongArgumentCount {
            name: "eval".into(),
            expected: 3,
            got: func_args.len(),
            span,
        });
    }

    let (expr_rv, expr_ty_idx) = lower_expr(func, &func_args[0], None)?;
    let expr_mty = func.get_type(expr_ty_idx, func_args[0].span)?;
    let inner_ty = match expr_mty {
        MathicType::SymbolicExpr(num_ty) => num_ty,
        _ => {
            return Err(LoweringError::MismatchedType {
                expected: expr_mty,
                found: expr_mty,
                span: func_args[0].span,
            });
        }
    };

    let sym_name = match &func_args[1].kind {
        ExprStmtKind::Primary(PrimaryExpr::Ident(name)) => {
            let local = func
                .sym_table
                .get_local_from_name(name, func_args[1].span)?;
            let local_ty = func.get_type(local.ty, func_args[1].span)?;
            if !local_ty.is_symbolic() {
                return Err(LoweringError::MismatchedType {
                    expected: expr_mty,
                    found: local_ty,
                    span: func_args[1].span,
                });
            }
            name.clone()
        }
        _ => {
            return Err(LoweringError::MismatchedType {
                expected: expr_mty,
                found: expr_mty,
                span: func_args[1].span,
            });
        }
    };

    let inner_ty_idx = func.get_or_insert_global_type_idx(MathicType::Numeric(inner_ty));
    let (value_rv, val_ty_idx) = lower_expr(func, &func_args[2], Some(inner_ty_idx))?;
    if val_ty_idx != inner_ty_idx {
        return Err(LoweringError::MismatchedType {
            expected: func.get_type(inner_ty_idx, func_args[2].span)?,
            found: func.get_type(val_ty_idx, func_args[2].span)?,
            span: func_args[2].span,
        });
    }

    let local_idx = func
        .sym_table
        .add_local(None, inner_ty_idx, None, LocalKind::Temp)?;
    let dest_block_idx = func.last_block_idx() + 1;

    func.get_basic_block_mut(func.last_block_idx()).terminator = Terminator::Eval {
        expr: expr_rv,
        sym_name,
        value: value_rv,
        return_dest: Value::InMemory {
            local_idx,
            modifier: vec![],
        },
        return_ty_idx: inner_ty_idx,
        dest_block: dest_block_idx,
        span: Some(span),
    };
    func.add_block(Terminator::Return(None, None), None);

    Ok((
        RValInstruct {
            kind: RValueKind::Use {
                value: Value::InMemory {
                    local_idx,
                    modifier: vec![],
                },
                span: None,
            },
            ty: inner_ty_idx,
        },
        inner_ty_idx,
    ))
}

fn lower_binary_op(
    func: &mut FunctionBuilder,
    lhs: &ExprStmt,
    op: BinaryOp,
    rhs: &ExprStmt,
    span: Span,
) -> Result<RValInstruct, LoweringError> {
    let (lhs, lhs_ty_idx) = lower_expr(func, lhs, None)?;
    let lhs_ty = func.get_type(lhs_ty_idx, span)?;

    let (rhs, rhs_ty_idx) = lower_expr(
        func,
        rhs,
        // If lhs_ty is symbolic, we don't want to add a ty_hint.
        if !lhs_ty.is_symbolic() {
            Some(lhs_ty_idx)
        } else {
            None
        },
    )?;

    let rhs_ty = func.get_type(rhs_ty_idx, span)?;

    let is_symbolic = lhs_ty.is_symbolic() || rhs_ty.is_symbolic();

    Ok(match op {
        BinaryOp::Arithmetic(arith) if is_symbolic => {
            if lhs_ty.is_symbolic() && rhs_ty.is_symbolic() && lhs_ty_idx != rhs_ty_idx {
                return Err(LoweringError::MismatchedType {
                    expected: lhs_ty,
                    found: rhs_ty,
                    span,
                });
            }

            let kind = RValueKind::SymbolicBinary {
                op: arith,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
                span,
            };

            RValInstruct {
                kind,
                ty: if lhs_ty.is_symbolic() {
                    lhs_ty_idx
                } else {
                    rhs_ty_idx
                },
            }
        }
        _ => {
            let inst_ty_idx = match op {
                BinaryOp::Compare(_) => func.get_or_insert_global_type_idx(MathicType::Bool),
                BinaryOp::Arithmetic(_) => lhs_ty_idx,
            };

            // Operands' types must match.
            if lhs_ty_idx != rhs_ty_idx {
                return Err(LoweringError::MismatchedType {
                    expected: func.get_type(lhs_ty_idx, span)?,
                    found: func.get_type(rhs_ty_idx, span)?,
                    span,
                });
            }

            let kind = RValueKind::Binary {
                op,
                lhs: Box::new(lhs),
                rhs: Box::new(rhs),
                span,
            };

            RValInstruct {
                kind,
                ty: inst_ty_idx,
            }
        }
    })
}

fn lower_logical_op(
    func: &mut FunctionBuilder,
    lhs: &ExprStmt,
    op: LogicalOp,
    rhs: &ExprStmt,
    span: Span,
) -> Result<RValInstruct, LoweringError> {
    let (lhs, lhs_ty_idx) = lower_expr(func, lhs, None)?;
    let (rhs, rhs_ty_idx) = lower_expr(func, rhs, Some(lhs_ty_idx))?;

    let lhs_ty = func.get_type(lhs_ty_idx, span)?;
    let rhs_ty = func.get_type(rhs_ty_idx, span)?;

    // Operands' types must be boolean.
    if !lhs_ty.is_bool() {
        return Err(LoweringError::MismatchedType {
            expected: MathicType::Bool,
            found: lhs_ty,
            span,
        });
    }
    if !rhs_ty.is_bool() {
        return Err(LoweringError::MismatchedType {
            expected: MathicType::Bool,
            found: rhs_ty,
            span,
        });
    }

    Ok(RValInstruct {
        kind: RValueKind::Logical {
            op,
            lhs: Box::new(lhs),
            rhs: Box::new(rhs),
            span,
        },
        ty: func.get_or_insert_global_type_idx(MathicType::Bool),
    })
}

fn lower_unary_op(
    func: &mut FunctionBuilder,
    op: UnaryOp,
    rhs: &ExprStmt,
    span: Span,
    ty_hint: Option<TypeIndex>,
) -> Result<RValInstruct, LoweringError> {
    let (rhs, rhs_ty) = lower_expr(func, rhs, ty_hint)?;

    Ok(RValInstruct {
        kind: RValueKind::Unary {
            op,
            rhs: Box::new(rhs),
            span,
        },
        ty: rhs_ty,
    })
}

fn lower_adt_init(
    func: &mut FunctionBuilder,
    name: &str,
    fields: &HashMap<String, ExprStmt>,
    span: Span,
) -> Result<RValInstruct, LoweringError> {
    let adt_ty = func.get_user_def_type(name, span)?;
    let adt_body = func.get_adt(adt_ty, span)?.clone();
    let mut init_fields = vec![None; fields.len()];

    if fields.len() != adt_body.fields_len() {
        let adt_fields_names = adt_body.get_field_names();
        let missing = adt_fields_names
            .into_iter()
            .filter(|n| fields.contains_key(n))
            .collect::<Vec<_>>()
            .join(", ");

        return Err(LoweringError::MissingStructFields { missing, span });
    }
    for (name, expr) in fields {
        let (rvalue, rvalue_ty_idx) = lower_expr(func, expr, adt_body.get_field_ty(name))?;
        let field_ty_idx =
            adt_body
                .get_field_ty(name)
                .ok_or(LoweringError::UndeclaredStructField {
                    found: name.to_string(),
                    span,
                })?;

        if field_ty_idx != rvalue_ty_idx {
            return Err(LoweringError::MismatchedType {
                expected: func.get_type(field_ty_idx, span)?,
                found: func.get_type(rvalue_ty_idx, span)?,
                span,
            });
        }

        let field_idx =
            adt_body
                .get_field_index(name)
                .ok_or(LoweringError::UndeclaredStructField {
                    found: name.to_string(),
                    span,
                })?;

        init_fields[field_idx] = Some(rvalue);
    }

    Ok(RValInstruct {
        kind: RValueKind::Init {
            init_inst: InitInstruct::StructInit {
                fields: init_fields.into_iter().map(Option::unwrap).collect(),
            },
            span,
        },
        ty: adt_ty,
    })
}

fn lower_struct_get(
    func: &mut FunctionBuilder,
    expr: &ExprStmt,
    field_name: &str,
    span: Span,
    ty_hint: Option<TypeIndex>,
) -> Result<RValInstruct, LoweringError> {
    let (struct_expr, struct_ty) = lower_expr(func, expr, ty_hint)?;

    let RValueKind::Use { value, .. } = struct_expr.kind else {
        unreachable!()
    };
    let Value::InMemory {
        local_idx,
        mut modifier,
    } = value
    else {
        unreachable!()
    };

    let struct_adt = func.get_adt(struct_ty, expr.span)?;
    let field_index =
        struct_adt
            .get_field_index(field_name)
            .ok_or(LoweringError::UndeclaredStructField {
                found: field_name.to_string(),
                span: expr.span,
            })?;
    let field_ty =
        struct_adt
            .get_field_ty(field_name)
            .ok_or(LoweringError::UndeclaredStructField {
                found: field_name.to_string(),
                span: expr.span,
            })?;

    modifier.push(ValueModifier::Field(field_index));

    Ok(RValInstruct {
        kind: RValueKind::Use {
            value: Value::InMemory {
                local_idx,
                modifier,
            },
            span: Some(span),
        },
        ty: field_ty,
    })
}

fn lower_struct_set(
    func: &mut FunctionBuilder,
    lhs: &ExprStmt,
    field_name: &str,
    rhs: &ExprStmt,
    span: Span,
) -> Result<RValInstruct, LoweringError> {
    let struct_field_value = lower_struct_get(func, lhs, field_name, span, None)?;
    let (local_idx, modifier) = {
        let RValueKind::Use { value, .. } = struct_field_value.kind else {
            unreachable!()
        };
        let Value::InMemory {
            local_idx,
            modifier,
        } = value
        else {
            unreachable!()
        };

        (local_idx, modifier)
    };
    let (value, value_ty_idx) = lower_expr(func, rhs, Some(struct_field_value.ty))?;

    if value_ty_idx != struct_field_value.ty {
        return Err(LoweringError::MismatchedType {
            expected: func.get_type(struct_field_value.ty, span)?,
            found: func.get_type(value_ty_idx, span)?,
            span,
        });
    }

    func.get_basic_block_mut(func.last_block_idx())
        .instructions
        .push(LValInstruct::Assign {
            local_idx,
            value,
            modifier,
            span: Some(span),
        });

    Ok(RValInstruct {
        kind: RValueKind::Use {
            value: Value::Const(ConstExpr::Void),
            span: None,
        },
        ty: func.get_or_insert_global_type_idx(MathicType::Void),
    })
}

fn lower_primary_value(
    func: &mut FunctionBuilder,
    expr: &PrimaryExpr,
    span: Span,
    ty_hint: Option<TypeIndex>,
) -> Result<RValInstruct, LoweringError> {
    let (value, ty) = match expr {
        PrimaryExpr::Ident(name) => {
            let local = func.sym_table.get_local_from_name(name, span)?;
            let local_ty = func.get_type(local.ty, span)?;
            // Use Symbol variant for symbolic expressions (SSA, no memory).
            let value = if local_ty.is_symbolic() {
                Value::Symbol {
                    local_idx: local.local_idx,
                }
            } else {
                Value::InMemory {
                    local_idx: local.local_idx,
                    modifier: vec![],
                }
            };
            (value, local.ty)
        }
        PrimaryExpr::Num(n) => match ty_hint {
            Some(ty) => (
                Value::Const(match func.get_type(ty, span)? {
                    MathicType::Numeric(NumericTy::Uint(uint_ty)) => match uint_ty {
                        UintTy::Usize => {
                            ConstExpr::Numeric(NumericConst::Usize(n.parse::<usize>().unwrap()))
                        }
                        UintTy::U8 => {
                            ConstExpr::Numeric(NumericConst::U8(n.parse::<u8>().unwrap()))
                        }
                        UintTy::U16 => {
                            ConstExpr::Numeric(NumericConst::U16(n.parse::<u16>().unwrap()))
                        }
                        UintTy::U32 => {
                            ConstExpr::Numeric(NumericConst::U32(n.parse::<u32>().unwrap()))
                        }
                        UintTy::U64 => {
                            ConstExpr::Numeric(NumericConst::U64(n.parse::<u64>().unwrap()))
                        }
                        UintTy::U128 => {
                            ConstExpr::Numeric(NumericConst::U128(n.parse::<u128>().unwrap()))
                        }
                    },
                    MathicType::Numeric(NumericTy::Sint(sint_ty)) => match sint_ty {
                        SintTy::Isize => {
                            ConstExpr::Numeric(NumericConst::Isize(n.parse::<isize>().unwrap()))
                        }
                        SintTy::I8 => {
                            ConstExpr::Numeric(NumericConst::I8(n.parse::<i8>().unwrap()))
                        }
                        SintTy::I16 => {
                            ConstExpr::Numeric(NumericConst::I16(n.parse::<i16>().unwrap()))
                        }
                        SintTy::I32 => {
                            ConstExpr::Numeric(NumericConst::I32(n.parse::<i32>().unwrap()))
                        }
                        SintTy::I64 => {
                            ConstExpr::Numeric(NumericConst::I64(n.parse::<i64>().unwrap()))
                        }
                        SintTy::I128 => {
                            ConstExpr::Numeric(NumericConst::I128(n.parse::<i128>().unwrap()))
                        }
                    },
                    MathicType::Numeric(NumericTy::Float(float_ty)) => match float_ty {
                        FloatTy::F32 => {
                            ConstExpr::Numeric(NumericConst::F32(n.parse::<f32>().unwrap()))
                        }
                        FloatTy::F64 => {
                            ConstExpr::Numeric(NumericConst::F64(n.parse::<f64>().unwrap()))
                        }
                    },
                    MathicType::Bool
                    | MathicType::Void
                    | MathicType::Char
                    | MathicType::Str
                    | MathicType::SymbolicExpr(_)
                    | MathicType::Adt { .. } => {
                        unreachable!()
                    }
                }),
                ty,
            ),
            None => (
                Value::Const(ConstExpr::Numeric(NumericConst::I32(
                    n.parse::<i32>().unwrap(),
                ))),
                func.get_or_insert_global_type_idx(MathicType::Numeric(NumericTy::Sint(
                    SintTy::I32,
                ))),
            ),
        },
        PrimaryExpr::Bool(b) => (
            Value::Const(ConstExpr::Bool(*b)),
            func.get_or_insert_global_type_idx(MathicType::Bool),
        ),
        PrimaryExpr::Str(s) => (
            Value::Const(ConstExpr::Str(s.clone())),
            func.get_or_insert_global_type_idx(MathicType::Str),
        ),
        PrimaryExpr::Char(c) => (
            Value::Const(ConstExpr::Char(*c)),
            func.get_or_insert_global_type_idx(MathicType::Char),
        ),
    };

    Ok(RValInstruct {
        kind: RValueKind::Use {
            value,
            span: Some(span),
        },
        ty,
    })
}

/// Tries to infer the type of an expression.
///
/// A **ty_hint** may be provided to help guessing the type of expressions such
/// as numeric constants, whose type depend on the bit width declared. In such
/// cases, if no **ty_hint** was provided, the default type will be returned.
fn lower_expression_type(
    func: &mut FunctionBuilder,
    expr: &ExprStmtKind,
    ty_hint: Option<TypeIndex>,
    span: Span,
) -> Result<TypeIndex, LoweringError> {
    Ok(match expr {
        ExprStmtKind::Primary(primary_expr) => match primary_expr {
            PrimaryExpr::Ident(name) => func.sym_table.get_local_from_name(name, span)?.ty,
            PrimaryExpr::Num(_) => match ty_hint {
                Some(ty) => ty,
                None => func.get_or_insert_global_type_idx(MathicType::Numeric(NumericTy::Sint(
                    SintTy::I32,
                ))),
            },
            PrimaryExpr::Str(_) => func.get_or_insert_global_type_idx(MathicType::Str),
            PrimaryExpr::Char(_) => func.get_or_insert_global_type_idx(MathicType::Char),
            PrimaryExpr::Bool(_) => func.get_or_insert_global_type_idx(MathicType::Bool),
        },
        ExprStmtKind::Binary { lhs, op, rhs } => match op {
            BinaryOp::Compare(_) => func.get_or_insert_global_type_idx(MathicType::Bool),
            BinaryOp::Arithmetic(_) => {
                // We need to check if either of the operans is symbolic since
                // the distinction is done through the type.
                let lhs_ty_idx = lower_expression_type(func, &lhs.kind, ty_hint, span)?;
                if func.get_type(lhs_ty_idx, span)?.is_symbolic() {
                    return Ok(lhs_ty_idx);
                }

                let rhs_ty_idx = lower_expression_type(func, &rhs.kind, ty_hint, span)?;
                if func.get_type(rhs_ty_idx, span)?.is_symbolic() {
                    return Ok(rhs_ty_idx);
                }

                lhs_ty_idx
            }
        },
        ExprStmtKind::Call { callee, .. } => {
            let func_decl = func.get_function_decl(callee, span)?;
            match func_decl.return_ty {
                Some(ty) => lower_inner_ast_type(func, &ty, span)?,
                None => func.get_or_insert_global_type_idx(MathicType::Void),
            }
        }
        ExprStmtKind::Group(expr_stmt) => lower_expression_type(func, &expr_stmt.kind, None, span)?,
        ExprStmtKind::Index { .. } => todo!(),
        ExprStmtKind::Logical { .. } => func.get_or_insert_global_type_idx(MathicType::Bool),
        ExprStmtKind::Unary { rhs, .. } => lower_expression_type(func, &rhs.kind, None, span)?,
        ExprStmtKind::Assign { expr, .. } | ExprStmtKind::StructSet { rhs: expr, .. } => {
            lower_expression_type(func, &expr.kind, None, span)?
        }
        ExprStmtKind::StructInit { name, .. } => func.get_user_def_type(name, span)?,
        ExprStmtKind::StructGet { expr, field_name } => {
            let adt_ty = lower_expression_type(func, &expr.kind, ty_hint, span)?;
            let adt = func.get_adt(adt_ty, span)?;

            adt.get_field_ty(field_name).unwrap()
        }
    })
}