use super::{Type, UInstrInst, UInstruction, UProgram};
use crate::common::{CompilerMsg, CompilerOutput, FileSpan};

impl UProgram {
    pub fn validate(&self, output: &mut CompilerOutput) {
        for (id, f) in self.iter_fns() {
            self.validate_fn(
                &f.instructions,
                self.origins.get(id),
                &f.ret,
                output,
                true,
                false,
            );
        }
    }

    pub fn validate_fn(
        &self,
        instructions: &[UInstrInst],
        origin: FileSpan,
        ret: &Type,
        output: &mut CompilerOutput,
        needs_ret: bool,
        breakable: bool,
    ) {
        let mut no_ret = true;
        for i in instructions {
            match &i.i {
                UInstruction::Mv { dst: dest, src } => {
                    let dest = self.expect(dest.id);
                    let src = self.expect(src.id);
                    output.check_assign(self, &src.ty, &dest.ty, i.origin);
                }
                UInstruction::Ref { dst: dest, src } => {
                    let dest = self.expect(dest.id);
                    let src = self.expect(src.id);
                    output.check_assign(self, &src.ty.clone().rf(), &dest.ty, i.origin);
                }
                UInstruction::LoadData { dst: dest, src } => {
                    let dest = self.expect(dest.id);
                    let src = self.expect(*src);
                    output.check_assign(self, &src.ty, &dest.ty, i.origin);
                }
                UInstruction::LoadSlice { dst: dest, src } => {
                    let dest = self.expect(dest.id);
                    let src = self.expect(*src);
                    let Type::Array(srcty, ..) = &src.ty else {
                        todo!()
                    };
                    output.check_assign(self, &Type::Slice(srcty.clone()), &dest.ty, i.origin);
                }
                UInstruction::LoadFn { dst: dest, src } => todo!(),
                UInstruction::Call { dst: dest, f, args } => {
                    let destty = &self.expect(dest.id).ty;
                    let f = self.expect(f.id);
                    let Type::Fn { args: argtys, ret } = &f.ty else {
                        output.err(CompilerMsg {
                            msg: format!("Type {} is not callable", self.type_name(&f.ty)),
                            spans: vec![dest.origin],
                        });
                        continue;
                    };
                    output.check_assign(self, ret, destty, dest.origin);
                    if args.len() != argtys.len() {
                        output.err(CompilerMsg {
                            msg: "Wrong number of arguments to function".to_string(),
                            spans: vec![dest.origin],
                        });
                    }
                    for (dst_ty, src) in argtys.iter().zip(args) {
                        let src_var = self.expect(src.id);
                        output.check_assign(self, &src_var.ty, dst_ty, src.origin);
                    }
                }
                UInstruction::AsmBlock { instructions, args } => {
                    for arg in args {
                        // TODO: validate size with enabled targets
                        // maybe should happen in lowering? but I think it could happen here
                        // if let Some(size) = self.size_of_var(arg.var.id)
                        //     && size != 64
                        // {
                        //     output.err(CompilerMsg {
                        //         msg: format!("asm block args must be size 64, is size {}", size),
                        //         spans: vec![arg.var.span],
                        //     });
                        // }
                    }
                }
                UInstruction::Ret { src } => {
                    let srcty = &self.expect(src.id).ty;
                    output.check_assign(self, srcty, ret, src.origin);
                    no_ret = false;
                }
                UInstruction::Construct { dst: dest, fields } => {
                    let dest_def = self.expect(dest.id);
                    let sty = match &dest_def.ty {
                        Type::Struct(sty) => sty,
                        _ => {
                            output.err(CompilerMsg {
                                msg: format!(
                                    "Type {} cannot be constructed",
                                    self.type_name(&dest_def.ty)
                                ),
                                spans: vec![dest.origin],
                            });
                            continue;
                        }
                    };
                    let def = self.expect(sty.id);
                    for (name, field) in &def.fields {
                        if let Some(var) = fields.get(name) {
                            let mut fty = &field.ty;
                            if let Type::Generic { id } = fty {
                                for (g, a) in def.generics.iter().zip(&sty.args) {
                                    if *g == *id {
                                        fty = a;
                                    }
                                }
                            }
                            let ety = &self.expect(var.id).ty;
                            output.check_assign(self, ety, fty, var.origin);
                        } else {
                            output.err(CompilerMsg {
                                msg: format!("field '{}' missing from struct", name),
                                spans: vec![dest.origin],
                            });
                        }
                    }
                }
                UInstruction::If { cond, body } => {
                    let cond = self.expect(cond.id);
                    output.check_assign(self, &cond.ty, &Type::Bits(64), i.origin);
                    self.validate_fn(body, origin, ret, output, false, breakable);
                }
                UInstruction::Loop { body } => {
                    self.validate_fn(body, origin, ret, output, false, true);
                }
                UInstruction::Break => {
                    if !breakable {
                        output.err(CompilerMsg {
                            msg: "Can't break here (outside of loop)".to_string(),
                            spans: vec![i.origin],
                        });
                    }
                    // TODO
                }
                UInstruction::Continue => {
                    if !breakable {
                        output.err(CompilerMsg {
                            msg: "Can't continue here (outside of loop)".to_string(),
                            spans: vec![i.origin],
                        });
                    }
                    // TODO
                }
            }
        }
        if needs_ret && no_ret && *ret != Type::Unit {
            output.err(CompilerMsg {
                msg: format!(
                    "Function implicitly returns () at the end, must return {}",
                    self.type_name(ret)
                ),
                spans: vec![origin],
            });
        }
    }
}