use crate::patio::types::{
    PatioBoolType, PatioByteType, PatioDataLayout, PatioDataPtrType, PatioFunctionPtrType,
    PatioInt32Type, PatioIntType, PatioType,
};

pub(crate) fn patio_type_to_cranelift_type(
    patio_type: &PatioType,
    target_triple: &target_lexicon::Triple,
) -> cranelift::prelude::Type {
    match patio_type {
        PatioType::PatioBoolType(_) => CRANELIFT_BOOL_TYPE,

        PatioType::PatioByteType(_) => CRANELIFT_BYTE_TYPE,

        PatioType::PatioDataPtrType(_) => {
            cranelift::prelude::Type::triple_pointer_type(&target_triple)
        }

        PatioType::PatioFunctionPtrType(_) => {
            cranelift::prelude::Type::triple_pointer_type(&target_triple)
        }

        PatioType::PatioIntType(_) => cranelift_int_type(target_triple),

        PatioType::PatioInt32Type(_) => CRANELIFT_INT_32_TYPE,
    }
}

pub(crate) const CRANELIFT_BOOL_TYPE: cranelift_codegen::ir::Type =
    cranelift_codegen::ir::types::I8;
pub(crate) const CRANELIFT_BYTE_TYPE: cranelift_codegen::ir::Type =
    cranelift_codegen::ir::types::I8;

pub(crate) fn cranelift_int_type(
    target_triple: &target_lexicon::Triple,
) -> cranelift_codegen::ir::Type {
    // TODO: this assumes that pointers are at least as wide as i32s:
    // probably safe, but we should check.
    // TODO: this relies on cranelift pointer types just being integers.
    cranelift::prelude::Type::triple_pointer_type(&target_triple)
}

pub(crate) const CRANELIFT_INT_32_TYPE: cranelift_codegen::ir::Type =
    cranelift_codegen::ir::types::I32;

pub(crate) fn data_layout_byte_size(
    layout: &PatioDataLayout,
    target_triple: &target_lexicon::Triple,
) -> u32 {
    match layout {
        PatioDataLayout::PatioDataLayoutArray(array) => {
            data_layout_byte_size(&array.element, target_triple)
                * u32::try_from(array.size).unwrap()
        }

        PatioDataLayout::PatioDataLayoutBool(_) => CRANELIFT_BOOL_TYPE.bytes(),

        PatioDataLayout::PatioDataLayoutByte(_) => CRANELIFT_BYTE_TYPE.bytes(),

        PatioDataLayout::PatioDataLayoutDataPtr(_) => {
            cranelift::prelude::Type::triple_pointer_type(&target_triple).bytes()
        }

        PatioDataLayout::PatioDataLayoutFunctionPtr(_) => {
            cranelift::prelude::Type::triple_pointer_type(&target_triple).bytes()
        }

        PatioDataLayout::PatioDataLayoutInt(_) => cranelift_int_type(target_triple).bytes(),

        PatioDataLayout::PatioDataLayoutInt32(_) => CRANELIFT_INT_32_TYPE.bytes(),

        PatioDataLayout::PatioDataLayoutStruct(layout_struct) => layout_struct
            .elements
            .iter()
            .map(|element| data_layout_byte_size(element, target_triple))
            .sum(),
    }
}

#[cfg(test)]
mod data_layout_byte_size_test {
    use std::str::FromStr;

    use crate::patio::types::{
        PatioDataLayout, PatioDataLayoutArray, PatioDataLayoutBool, PatioDataLayoutByte,
        PatioDataLayoutDataPtr, PatioDataLayoutFunctionPtr, PatioDataLayoutInt,
        PatioDataLayoutInt32, PatioDataLayoutStruct,
    };

    fn data_layout_byte_size(layout: &PatioDataLayout) -> u32 {
        super::data_layout_byte_size(
            layout,
            &target_lexicon::Triple::from_str("x86_64-unknown-linux-gnu").unwrap(),
        )
    }

    #[test]
    fn bool_is_one_byte() {
        assert_eq!(data_layout_byte_size(&PatioDataLayoutBool.into()), 1);
    }

    #[test]
    fn byte_is_one_byte() {
        assert_eq!(data_layout_byte_size(&PatioDataLayoutByte.into()), 1);
    }

    #[test]
    fn data_ptr_is_width_of_pointer() {
        assert_eq!(data_layout_byte_size(&PatioDataLayoutDataPtr.into()), 8);
    }

    #[test]
    fn function_ptr_is_width_of_pointer() {
        assert_eq!(data_layout_byte_size(&PatioDataLayoutFunctionPtr.into()), 8);
    }

    #[test]
    fn int_is_width_of_pointer_and_at_least_32_bits() {
        assert_eq!(data_layout_byte_size(&PatioDataLayoutInt.into()), 8);
    }

    #[test]
    fn int32_is_four_bytes() {
        assert_eq!(data_layout_byte_size(&PatioDataLayoutInt32.into()), 4);
    }

    #[test]
    fn array_size_is_size_of_each_element_multiplied_by_number_of_elements() {
        assert_eq!(
            data_layout_byte_size(
                &PatioDataLayoutArray {
                    element: PatioDataLayoutInt32.into(),
                    size: 7
                }
                .into()
            ),
            28
        );
    }

    #[test]
    fn struct_size_is_sum_of_element_sizes() {
        assert_eq!(
            data_layout_byte_size(
                &PatioDataLayoutStruct {
                    elements: vec![
                        PatioDataLayout::PatioDataLayoutInt32(PatioDataLayoutInt32),
                        PatioDataLayout::PatioDataLayoutByte(PatioDataLayoutByte)
                    ]
                }
                .into()
            ),
            5
        );
    }
}

pub(crate) fn data_layout_byte_offset(
    layout: &PatioDataLayout,
    element_offsets: &[i32],
    target_triple: &target_lexicon::Triple,
) -> u32 {
    let mut byte_offset = 0;
    let mut current_element = layout;

    for element_offset in element_offsets {
        match current_element {
            PatioDataLayout::PatioDataLayoutArray(array) => {
                byte_offset += data_layout_byte_size(&array.element, target_triple)
                    * u32::try_from(*element_offset).unwrap();
                current_element = &array.element;
            }

            PatioDataLayout::PatioDataLayoutBool(_) => panic!("offset not allowed for bool"),
            PatioDataLayout::PatioDataLayoutByte(_) => panic!("offset not allowed for byte"),
            PatioDataLayout::PatioDataLayoutDataPtr(_) => panic!("offset not allowed for data ptr"),
            PatioDataLayout::PatioDataLayoutFunctionPtr(_) => {
                panic!("offset not allowed for function ptr")
            }
            PatioDataLayout::PatioDataLayoutInt(_) => panic!("offset not allowed for int"),
            PatioDataLayout::PatioDataLayoutInt32(_) => panic!("offset not allowed for int32"),

            PatioDataLayout::PatioDataLayoutStruct(layout_struct) => {
                for element in layout_struct
                    .elements
                    .iter()
                    .take((*element_offset).try_into().unwrap())
                {
                    byte_offset += data_layout_byte_size(element, target_triple);
                }

                current_element =
                    &layout_struct.elements[usize::try_from(*element_offset).unwrap()];
            }
        }
    }

    byte_offset
}

#[cfg(test)]
mod data_layout_byte_offset_test {
    use std::str::FromStr;

    use crate::patio::types::{
        PatioDataLayout, PatioDataLayoutArray, PatioDataLayoutByte, PatioDataLayoutDataPtr,
        PatioDataLayoutInt32, PatioDataLayoutStruct,
    };

    fn data_layout_byte_offset(layout: &PatioDataLayout, element_offsets: Vec<i32>) -> u32 {
        super::data_layout_byte_offset(
            layout,
            &element_offsets,
            &target_lexicon::Triple::from_str("x86_64-unknown-linux-gnu").unwrap(),
        )
    }

    #[test]
    fn byte_offset_of_first_element_of_array_is_zero() {
        let layout = PatioDataLayoutArray {
            element: PatioDataLayoutInt32.into(),
            size: 2,
        }
        .into();

        let byte_offset = data_layout_byte_offset(&layout, vec![0]);

        assert_eq!(byte_offset, 0);
    }

    #[test]
    fn byte_offset_of_element_of_array_is_size_of_skipped_elements() {
        let layout = PatioDataLayoutArray {
            element: PatioDataLayoutInt32.into(),
            size: 5,
        }
        .into();

        let byte_offset = data_layout_byte_offset(&layout, vec![2]);

        assert_eq!(byte_offset, 8);
    }

    #[test]
    fn byte_offset_of_element_of_element_of_array_is_size_of_skipped_elements() {
        let layout = PatioDataLayoutArray {
            element: PatioDataLayoutArray {
                element: PatioDataLayoutByte.into(),
                size: 7,
            }
            .into(),
            size: 5,
        }
        .into();

        let byte_offset = data_layout_byte_offset(&layout, vec![2, 3]);

        assert_eq!(byte_offset, 17);
    }

    #[test]
    fn byte_offset_of_first_element_of_struct_is_zero() {
        let layout = PatioDataLayoutStruct {
            elements: vec![
                PatioDataLayoutInt32.into(),
                PatioDataLayoutByte.into(),
                PatioDataLayoutDataPtr.into(),
            ],
        }
        .into();

        let byte_offset = data_layout_byte_offset(&layout, vec![0]);

        assert_eq!(byte_offset, 0);
    }

    #[test]
    fn byte_offset_of_element_of_struct_is_size_of_skipped_elements() {
        let layout = PatioDataLayoutStruct {
            elements: vec![
                PatioDataLayoutInt32.into(),
                PatioDataLayoutByte.into(),
                PatioDataLayoutDataPtr.into(),
            ],
        }
        .into();

        let byte_offset = data_layout_byte_offset(&layout, vec![2]);

        assert_eq!(byte_offset, 5);
    }

    #[test]
    fn byte_offset_of_element_of_element_of_struct_is_size_of_skipped_elements() {
        let layout = PatioDataLayoutStruct {
            elements: vec![
                PatioDataLayoutInt32.into(),
                PatioDataLayoutStruct {
                    elements: vec![
                        PatioDataLayoutByte.into(),
                        PatioDataLayoutDataPtr.into(),
                        PatioDataLayoutByte.into(),
                    ],
                }
                .into(),
                PatioDataLayoutInt32.into(),
            ],
        }
        .into();

        let byte_offset = data_layout_byte_offset(&layout, vec![1, 2]);

        assert_eq!(byte_offset, 13);
    }
}

pub(crate) fn data_layout_element_at<'a>(
    layout: &'a PatioDataLayout,
    element_offsets: &[i32],
) -> &'a PatioDataLayout {
    let mut current_element = layout;

    for element_offset in element_offsets {
        match current_element {
            PatioDataLayout::PatioDataLayoutArray(array) => {
                current_element = &array.element;
            }

            PatioDataLayout::PatioDataLayoutBool(_) => panic!("offset not allowed for bool"),
            PatioDataLayout::PatioDataLayoutByte(_) => panic!("offset not allowed for byte"),
            PatioDataLayout::PatioDataLayoutDataPtr(_) => panic!("offset not allowed for data ptr"),
            PatioDataLayout::PatioDataLayoutFunctionPtr(_) => {
                panic!("offset not allowed for function ptr")
            }
            PatioDataLayout::PatioDataLayoutInt(_) => panic!("offset not allowed for int"),
            PatioDataLayout::PatioDataLayoutInt32(_) => panic!("offset not allowed for int32"),

            PatioDataLayout::PatioDataLayoutStruct(layout_struct) => {
                current_element =
                    &layout_struct.elements[usize::try_from(*element_offset).unwrap()];
            }
        }
    }

    current_element
}

#[cfg(test)]
mod data_layout_element_at_test {
    use crate::patio::types::{
        PatioDataLayout, PatioDataLayoutArray, PatioDataLayoutByte, PatioDataLayoutDataPtr,
        PatioDataLayoutInt32, PatioDataLayoutStruct,
    };

    fn data_layout_element_at(
        layout: &PatioDataLayout,
        element_offsets: Vec<i32>,
    ) -> &PatioDataLayout {
        super::data_layout_element_at(layout, &element_offsets)
    }

    #[test]
    fn element_of_array_is_element_layout() {
        let layout = PatioDataLayoutArray {
            element: PatioDataLayoutInt32.into(),
            size: 2,
        }
        .into();

        let element = data_layout_element_at(&layout, vec![0]);

        assert_eq!(element, &PatioDataLayoutInt32.into());
    }

    #[test]
    fn element_of_element_of_array() {
        let layout = PatioDataLayoutArray {
            element: PatioDataLayoutArray {
                element: PatioDataLayoutByte.into(),
                size: 7,
            }
            .into(),
            size: 5,
        }
        .into();

        let element = data_layout_element_at(&layout, vec![2, 3]);

        assert_eq!(element, &PatioDataLayoutByte.into());
    }

    #[test]
    fn first_element_of_struct() {
        let layout = PatioDataLayoutStruct {
            elements: vec![
                PatioDataLayoutInt32.into(),
                PatioDataLayoutByte.into(),
                PatioDataLayoutDataPtr.into(),
            ],
        }
        .into();

        let element = data_layout_element_at(&layout, vec![0]);

        assert_eq!(element, &PatioDataLayoutInt32.into());
    }

    #[test]
    fn non_first_element_of_struct() {
        let layout = PatioDataLayoutStruct {
            elements: vec![
                PatioDataLayoutInt32.into(),
                PatioDataLayoutByte.into(),
                PatioDataLayoutDataPtr.into(),
            ],
        }
        .into();

        let element = data_layout_element_at(&layout, vec![2]);

        assert_eq!(element, &PatioDataLayoutDataPtr.into());
    }

    #[test]
    fn element_of_element_of_struct() {
        let layout = PatioDataLayoutStruct {
            elements: vec![
                PatioDataLayoutInt32.into(),
                PatioDataLayoutStruct {
                    elements: vec![
                        PatioDataLayoutByte.into(),
                        PatioDataLayoutDataPtr.into(),
                        PatioDataLayoutByte.into(),
                    ],
                }
                .into(),
                PatioDataLayoutInt32.into(),
            ],
        }
        .into();

        let element = data_layout_element_at(&layout, vec![1, 2]);

        assert_eq!(element, &PatioDataLayoutByte.into());
    }
}

pub(crate) fn data_layout_to_type(layout: &PatioDataLayout) -> PatioType {
    match layout {
        PatioDataLayout::PatioDataLayoutArray(_) => panic!("array cannot be represented as type"),
        PatioDataLayout::PatioDataLayoutBool(_) => PatioBoolType.into(),
        PatioDataLayout::PatioDataLayoutByte(_) => PatioByteType.into(),
        PatioDataLayout::PatioDataLayoutDataPtr(_) => PatioDataPtrType.into(),
        PatioDataLayout::PatioDataLayoutFunctionPtr(_) => PatioFunctionPtrType.into(),
        PatioDataLayout::PatioDataLayoutInt(_) => PatioIntType.into(),
        PatioDataLayout::PatioDataLayoutInt32(_) => PatioInt32Type.into(),
        PatioDataLayout::PatioDataLayoutStruct(_) => panic!("struct cannot be represented as type"),
    }
}