mclangc/src/parser/ast/statement.rs

use anyhow::bail;

use crate::{common::loc::LocBox, lerror, parser::ast::expr::Path};

use super::{expr::{Block, Expr}, literal::Literal, typ::Type, Char, Ident, Number, Program, TString};

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum Statement {
    Fn(Function),
    TypeAlias(TypeAlias),
    Struct(Struct),
    Enum(Enum),
    ConstVar(ConstVar),
    StaticVar(StaticVar),
    Let(Let),
}


#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Let {
    pub name: Ident,
    pub typ: Option<LocBox<Type>>,
    pub val: Option<LocBox<Expr>>,
}
#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct ConstVar {
    pub name: Ident,
    pub typ: LocBox<Type>,
    pub val: LocBox<Expr>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct StaticVar {
    pub name: Ident,
    pub typ: LocBox<Type>,
    pub val: LocBox<Expr>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct TypeAlias {
    pub name: Ident,
    pub typ: LocBox<Type>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Struct {
    pub name: Ident,
    pub fields: Vec<(Ident, LocBox<Type>)>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Enum {
    pub name: Ident,
    pub fields: Vec<Ident>,
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub struct Function {
    pub struct_name: Option<Ident>,
    pub name: Ident,
    pub params: Vec<(Ident, LocBox<Type>)>,
    pub ret_type: Option<LocBox<Type>>,
    pub qual_const: bool,
    pub qual_extern: Option<TString>, // abi
    pub body: Option<Block>, // If None then its a type declaration
}

impl Function {
    pub fn get_full_name(&self) -> String {
        if let Some(sn) = &self.struct_name {
            format!("{sn}__S__{}", self.name.0)
        } else {
            format!("{}", self.name.0)
        }
    }
    pub fn get_full_name_pretty(&self) -> String {
        if let Some(sn) = &self.struct_name {
            format!("{sn}::{}", self.name.0)
        } else {
            format!("{}", self.name.0)
        }
    }
    pub fn get_def_as_str(&self) -> String {
        let mut f = String::new();
        f.push_str(&format!("fn {} (", self.get_full_name_pretty()));
        let mut first = true;
        for (name, typ) in &self.params {
            if first {
                first = false;
            } else {
                f.push_str(", ");
            }
            f.push_str(&format!("{name}: {}", typ.inner()));
        }
        f.push_str(")");
        if let Some(ret) = &self.ret_type {
            f.push_str(&format!(" -> {}", ret.inner()));
        }
        f.push_str(";");
        f
    }
}


impl Struct {
    pub fn get_size(&self, program: &Program) -> anyhow::Result<usize> {
        let mut size = 0;
        for (_, v) in &self.fields {
            size += v.inner().size_of(program)?;
        }
        Ok(size)
    }

    pub fn get_offset_of(&self, program: &Program, field: &Ident) -> anyhow::Result<usize> {
        let mut size = 0;
        for (k, v) in &self.fields {
            if k == field {
                return Ok(size);
            }
            size += v.inner().size_of(program)?;
        }
        bail!("Field not found")
    }
}

#[derive(Debug, Clone, PartialEq, PartialOrd)]
pub enum ConstDataTyp {
    Byte(u8),
    AddrOfConst(Ident),
    AddrOfStatic(Ident),
    AddrOfFunc(Ident),
}

pub fn get_constant_data_as_bytes(program: &Program, value: LocBox<Expr>, is_little_endian: bool, must_be_number: bool) -> anyhow::Result<Vec<ConstDataTyp>> {
    match value.inner() {
        Expr::Literal(lit) => {
            match lit {
                Literal::Char(Char(c)) => {
                    if must_be_number {
                        lerror!(value.loc(), "Expected number got char");
                        bail!("")
                    }
                    Ok(vec![ConstDataTyp::Byte(*c as u8)])
                },
                Literal::String(TString { val, cstr }) => {
                    if must_be_number {
                        lerror!(value.loc(), "Expected number got string");
                        bail!("")
                    }
                    let mut val = val.chars().into_iter().map(|f| ConstDataTyp::Byte(f as u8)).collect::<Vec<_>>();
                    if *cstr {
                        val.push(ConstDataTyp::Byte(0));
                    }
                    Ok(val)
                }
                Literal::Number(Number { val, base: _, signed: _ }) => {
                    let mut buf = Vec::new();
                    let mut inc = 0;
                    while inc < 8 {
                        buf.push(ConstDataTyp::Byte(((val << 8*inc) & 0xff) as u8));
                        inc += 1;
                    }
                    
                    if is_little_endian {
                        buf.reverse();
                    }

                    Ok(buf)
                }
                Literal::Array(arr) => {
                    if must_be_number {
                        lerror!(value.loc(), "Expected number got array");
                        bail!("")
                    }
                    let mut bytes = Vec::new();
                    if arr.len() < 1 {
                        return Ok(vec![]);
                    }
                    for item in arr {
                        let mut data = get_constant_data_as_bytes(program, item.clone(), is_little_endian, must_be_number)?;
                        bytes.append(&mut data);
                    }
                    
                    Ok(bytes)
                }
                Literal::ArrayRepeat { val, count } => {
                    if must_be_number {
                        lerror!(value.loc(), "Expected number got repeating array");
                        bail!("")
                    }
                    let mut val = get_constant_data_as_bytes(program, (**val).clone(), is_little_endian, must_be_number)?;

                    let mut num = Vec::new();
                    let mut inc = 0;
                    while inc < 8 {
                        num.push(ConstDataTyp::Byte(((count << 8*inc) & 0xff) as u8));
                        inc += 1;
                    }
                    
                    if is_little_endian {
                        num.reverse();
                    }

                    let mut count = 0 as usize;
                    for b in num {
                        let ConstDataTyp::Byte(b) = b else {unreachable!()};
                        count = b as usize;
                        count <<= 8;
                    }
                    let orig = val.clone();
                    for _ in 0..count {
                        val.append(&mut orig.clone());
                    }
                    Ok(val)
                }
                Literal::Ident(name) => {
                    if let Some(var) = program.const_vars.get(name) {
                        Ok(get_constant_data_as_bytes(program, var.val.clone(), is_little_endian, must_be_number)?)
                    } else if let Some(_) = program.static_vars.get(name) {
                        lerror!(value.loc(), "Statics cannot be passed by value, use a reference");
                        bail!("")
                    } else if let Some(_) = program.functions.get(name) {
                        Ok(vec![ConstDataTyp::AddrOfFunc(name.clone())])
                    } else {
                        lerror!(value.loc(), "Unable to find ident '{name}'");
                        bail!("")
                    }
                }
                //Literal::Struct { name, fields } => {
                //    if must_be_number {
                //        lerror!(value.loc(), "Expected number got struct literal");
                //        bail!("")
                //    }
                //    let Some(strct) = program.structs.get(name) else {
                //        lerror!(value.loc(), "Could not find struct {name}");
                //        bail!("")
                //    };
                //    
                //    let mut strct_fields = HashMap::new();
                //    for (name, typ) in &strct.inner().fields {
                //        strct_fields.insert(name, typ);
                //    }
//
                //    for (name, val) in fields {
                //        if let Some(_fld) = strct_fields.get(name) {
                //            // TODO: Actually check if the fields are the right type
                //            
                //        }
                //    }
                //    todo!()
                //}

            }
        }
        _ => unreachable!()
    }
}