use std::fmt::Display;

use crate::{ast::types::Type as AstType, ir::types::IRType};

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum SemaType {
    I32,
    I1,
    Void,
    Array(Box<SemaType>, Vec<usize>),
    Ptr(Box<SemaType>),
}

impl SemaType {
    pub fn get_elevate_result(lhs: &SemaType, rhs: &SemaType) -> Option<SemaType> {
        if lhs == rhs {
            Some(lhs.clone())
        } else if (*lhs == SemaType::I32 && *rhs == SemaType::I1) || (*lhs == SemaType::I1 && *rhs == SemaType::I32) {
            Some(SemaType::I32)
        } else {
            None
        }
    }

    pub fn is_scalar(&self) -> bool {
        matches!(self, SemaType::I32 | SemaType::I1)
    }

    pub fn is_array_like(&self) -> bool {
        matches!(self, SemaType::Array(_, _) | SemaType::Ptr(_))
    }

    pub fn indexed_type(&self) -> Option<SemaType> {
        match self {
            SemaType::Array(elem, dims) => {
                if dims.len() == 1 {
                    Some((**elem).clone())
                } else if dims.first().is_some_and(|dim| *dim == 0) {
                    Some(SemaType::Array(elem.clone(), dims[1..].to_vec()))
                } else {
                    Some(SemaType::Array(elem.clone(), dims[1..].to_vec()))
                }
            }
            SemaType::Ptr(elem) => Some((**elem).clone()),
            _ => None,
        }
    }

    pub fn element_size_in_bytes(&self) -> usize {
        match self {
            SemaType::I32 => 4,
            SemaType::I1 => 1,
            SemaType::Void => 0,
            SemaType::Array(elem, dims) => elem.element_size_in_bytes() * dims.iter().product::<usize>(),
            SemaType::Ptr(_) => 4,
        }
    }
}

impl Display for SemaType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SemaType::I32 => write!(f, "i32"),
            SemaType::I1 => write!(f, "i1"),
            SemaType::Void => write!(f, "void"),
            SemaType::Array(elem, dims) => {
                write!(f, "{}", elem)?;
                for dim in dims {
                    write!(f, "[{}]", dim)?;
                }
                Ok(())
            }
            SemaType::Ptr(elem) => write!(f, "{}*", elem),
        }
    }
}

impl From<AstType> for SemaType {
    fn from(value: AstType) -> Self {
        match value {
            AstType::Int => SemaType::I32,
            AstType::Void => SemaType::Void,
        }
    }
}

impl From<SemaType> for IRType {
    fn from(value: SemaType) -> Self {
        match value {
            SemaType::I32 => IRType::I32,
            SemaType::I1 => IRType::I1,
            SemaType::Void => IRType::Void,
            SemaType::Array(elem, dims) => IRType::Array(Box::new((*elem).into()), dims),
            SemaType::Ptr(elem) => IRType::Ptr(Box::new((*elem).into())),
        }
    }
}
