device-driver-lir 2.0.0

use std::ops::Add;

use convert_case::Case;
use device_driver_common::{
    identifier::{All, Identifier},
    span::{SpanExt, Spanned},
    specifiers::{BaseType, Integer, Repeat, RepeatSource},
};
use device_driver_diagnostics::{DynError, ResultExt};

use crate::model as lir;
use device_driver_mir::{
    find_min_max_addresses,
    model::{self as mir, Object},
    search_object,
};

pub fn transform_devices(manifest: &mir::Manifest) -> Result<Vec<lir::Device>, DynError> {
    manifest
        .iter_devices_with_config()
        .map(|(device, device_config)| {
            // Create a root block and pass the device objects to it
            let blocks = collect_into_blocks(
                BorrowedBlock {
                    description: &format!(
                        "{}Root block of the {} driver",
                        if device.description.is_empty() {
                            String::new()
                        } else {
                            format!("{}\n\n", device.description)
                        },
                        device.name.to_case(Case::Pascal),
                    ),
                    // Cast unchecked is fine here since this is a root block and the identifier is never used as an operation
                    name: &device.name.value.clone().cast_unchecked(),
                    address_offset: &0,
                    repeat: &None,
                    objects: &device.objects,
                },
                true,
                &device_config,
                manifest,
            )
            .with_message(|| "could not collect into blocks")?;

            Ok(lir::Device {
                internal_address_type: find_best_internal_address_type(manifest, device),
                blocks,
            })
        })
        .collect()
}

fn collect_into_blocks(
    block: BorrowedBlock,
    is_root: bool,
    device_config: &mir::DeviceConfig,
    manifest: &mir::Manifest,
) -> Result<Vec<lir::Block>, DynError> {
    let mut blocks = Vec::new();

    let BorrowedBlock {
        description,
        name,
        address_offset: _,
        repeat: _,
        objects,
    } = block;

    let mut methods = Vec::new();

    for object in objects {
        let Some(method) =
            get_method(object, &mut blocks, device_config, manifest).with_message(|| {
                format!(
                    "could not get method for object {}",
                    object.name().original()
                )
            })?
        else {
            continue;
        };

        methods.push(method);
    }

    let new_block = lir::Block {
        description: description.clone(),
        root: is_root,
        name: name.clone().cast(),
        register_address_type: device_config
            .register_address_type
            .map(|v| v.value)
            .unwrap_or(Integer::U8),
        command_address_type: device_config
            .command_address_type
            .map(|v| v.value)
            .unwrap_or(Integer::U8),
        buffer_address_type: device_config
            .buffer_address_type
            .map(|v| v.value)
            .unwrap_or(Integer::U8),
        register_address_mode: device_config.register_address_mode.map(|v| v.value),
        methods,
    };

    blocks.insert(0, new_block);

    Ok(blocks)
}

fn get_method(
    object: &mir::Object,
    blocks: &mut Vec<lir::Block>,
    device_config: &mir::DeviceConfig,
    manifest: &mir::Manifest,
) -> Result<Option<lir::BlockMethod>, DynError> {
    let method = match object {
        mir::Object::Device(_) => None,
        mir::Object::Block(
            b @ mir::Block {
                description,
                name,
                address_offset,
                repeat,
                ..
            },
        ) => {
            blocks.extend(collect_into_blocks(
                b.into(),
                false,
                device_config,
                manifest,
            )?);

            Some(lir::BlockMethod {
                description: description.clone(),
                name: name.value.clone().cast(),
                address: address_offset.value,
                repeat: repeat_to_method_kind(repeat, manifest),
                method_type: lir::BlockMethodType::Block {
                    name: name.value.clone().cast(),
                },
            })
        }
        mir::Object::Register(mir::Register {
            description,
            name,
            address,
            access,
            repeat,
            field_set_ref,
            reset_value,
            ..
        }) => {
            let field_set = search_object(manifest, field_set_ref).ok_or(DynError::new(
                format!("fieldset {} could not be found", field_set_ref.original()),
            ))?;

            Some(lir::BlockMethod {
                description: description.clone(),
                name: name.value.clone(),
                address: address.value,
                repeat: repeat_to_method_kind(repeat, manifest),
                method_type: lir::BlockMethodType::Register {
                    field_set_name: field_set.name().clone().cast_assert(),
                    access: access.ok_or_else(|| DynError::new("access is not set"))?,
                    reset_value: reset_value.as_ref().map(|rv| {
                        rv.as_array().cloned().map(|array| array.with_span(rv.span)).ok_or_else(
                            || DynError::new("reset value is not an array while it should have been converted to array a mir pass"),
                        )
                    })
                    .transpose()?,
                },
            })
        }
        mir::Object::Command(mir::Command {
            description,
            name,
            address,
            repeat,
            field_set_ref_in,
            field_set_ref_out,
            ..
        }) => {
            let field_set_in = field_set_ref_in
                .as_ref()
                .map(|id_ref| {
                    search_object(manifest, id_ref).ok_or(DynError::new(format!(
                        "fieldset {} could not be found",
                        id_ref.original()
                    )))
                })
                .transpose()?;
            let field_set_out = field_set_ref_out
                .as_ref()
                .map(|id_ref| {
                    search_object(manifest, id_ref).ok_or(DynError::new(format!(
                        "fieldset {} could not be found",
                        id_ref.original()
                    )))
                })
                .transpose()?;

            Some(lir::BlockMethod {
                description: description.clone(),
                name: name.value.clone(),
                address: address.value,
                repeat: repeat_to_method_kind(repeat, manifest),
                method_type: lir::BlockMethodType::Command {
                    field_set_name_in: field_set_in.map(|fs_in| fs_in.name().clone().cast_assert()),
                    field_set_name_out: field_set_out
                        .map(|fs_out| fs_out.name().clone().cast_assert()),
                },
            })
        }
        mir::Object::Buffer(mir::Buffer {
            description,
            name,
            access,
            address,
            short_properties_span: _,
            properties_span: _,
            span: _,
        }) => Some(lir::BlockMethod {
            description: description.clone(),
            name: name.value.clone(),
            address: address.value,
            repeat: lir::Repeat::None, // Buffers can't be repeated (for now?)
            method_type: lir::BlockMethodType::Buffer {
                access: access.ok_or_else(|| DynError::new("access is not set"))?,
            },
        }),
        mir::Object::FieldSet(_) => None,
        mir::Object::Enum(_) => None,
        mir::Object::Extern(_) => None,
        mir::Object::Field(_) => None,
    };

    Ok(method)
}

pub fn transform_field_sets(manifest: &mir::Manifest) -> Result<Vec<lir::FieldSet>, DynError> {
    manifest
        .iter_objects()
        .filter_map(|o| {
            if let Object::FieldSet(fs) = o {
                Some(
                    transform_field_set(manifest, fs)
                        .with_message(|| format!("transforming fieldset {}", fs.name.original())),
                )
            } else {
                None
            }
        })
        .collect()
}

fn transform_field_set(
    manifest: &mir::Manifest,
    field_set: &mir::FieldSet,
) -> Result<lir::FieldSet, DynError> {
    let fields = field_set
        .fields
        .iter()
        .map(|field| {
            transform_field(manifest, field)
                .with_message(|| format!("transforming field {}", field.name.original()))
        })
        .collect::<Result<_, _>>()
        .with_message(|| "transforming fields")?;

    Ok(lir::FieldSet {
        description: field_set.description.clone(),
        name: field_set.name.value.clone(),
        byte_order: field_set.byte_order.ok_or_else(|| {
            DynError::new("Byte order should never be none at this point after the MIR passes")
        })?,
        size_bytes: field_set.size_bytes.value,
        fields,
    })
}

fn transform_field(manifest: &mir::Manifest, field: &mir::Field) -> Result<lir::Field, DynError> {
    let mir::Field {
        description,
        name,
        access,
        base_type,
        field_conversion,
        field_address,
        repeat,
        short_properties_span: _,
        properties_span: _,
        span: _,
    } = field;

    let (base_type, conversion_method) = match (base_type.value, field_conversion) {
        (BaseType::Unspecified | BaseType::Uint | BaseType::Int, _) => {
            return Err(DynError::new(
                "base type cannot be left unspecified or unsized after the mir passes",
            ));
        }
        (BaseType::Bool, None) if field_address.len() == 1 => {
            ("u8".to_string(), lir::FieldConversionMethod::Bool)
        }
        (BaseType::Bool, _) => {
            return Err(DynError::new(
                "bools can only be 1 bit and have no conversion. Should have been checked in a MIR pass.",
            ));
        }
        (BaseType::FixedSize(integer), None) => {
            (integer.to_string(), lir::FieldConversionMethod::None)
        }
        (BaseType::FixedSize(integer), Some(fc)) => (integer.to_string(), {
            let field_bits = field.field_address.len() as u32;

            let fc_identifier = search_object(manifest, &fc.type_name)
                .ok_or_else(|| {
                    DynError::new(format!(
                        "{} existence checked in MIR pass",
                        fc.type_name.original()
                    ))
                })?
                .name()
                .clone();

            // Always use try if that's specified
            if fc.fallible {
                lir::FieldConversionMethod::TryInto(fc_identifier.cast_assert())
            }
            // Are we pointing at a potentially infallible enum and do we fulfil the requirements?
            else if let Some(mir::Enum {
                generation_style: Some(mir::EnumGenerationStyle::InfallibleWithinRange),
                size_bits,
                ..
            }) = manifest
                .iter_enums()
                .find(|e| e.name.take_ref() == fc.type_name.value)
                && field_bits
                    <= size_bits.ok_or_else(|| {
                        DynError::new(format!(
                            "enum {} size_bits must have been set in an earlier mir pass",
                            fc.type_name.original()
                        ))
                    })?
            {
                // This field is equal or smaller in bits than the infallible enum. So we can do the unsafe into
                lir::FieldConversionMethod::UnsafeInto(fc_identifier.cast_assert())
            } else {
                // Fallback is to use the into trait.
                // This is correct because in the field_conversion_valid mir pass we've already exited if we need a try and didn't specify it.
                // The only other option is the unsafe into and we've just checked that.
                lir::FieldConversionMethod::Into(fc_identifier.cast_assert())
            }
        }),
    };

    Ok(lir::Field {
        description: description.clone(),
        name: name.value.clone(),
        address: field_address.value,
        base_type,
        conversion_method,
        access: access.ok_or_else(|| DynError::new("access is not set"))?,
        repeat: repeat_to_method_kind(repeat, manifest),
    })
}

pub fn transform_enums(manifest: &mir::Manifest) -> Vec<lir::Enum> {
    manifest.iter_enums().map(|e| {
        let mir::Enum {
            description,
            name,
            variants: _,
            base_type,
            size_bits: _,
            generation_style: _,
            short_properties_span: _,
            properties_span: _,
            span: _
        } = e;

        let base_type = match base_type.value {
            BaseType::FixedSize(integer) => integer.to_string(),
            _ => {
                panic!("Enum base type should be set to fixed size integer in a mir pass at this point")
            }
        };

        let variants = e
            .iter_variants_with_discriminant()
            .map(|(discriminant, v)| {
                let mir::EnumVariant {
                    description,
                    name,
                    value,
                    span: _
               } = v;

                lir::EnumVariant {
                    description: description.clone(),
                    name: name.value.clone(),
                    discriminant,
                    default: matches!(value, mir::EnumValue::Default(_)),
                    catch_all: matches!(value, mir::EnumValue::CatchAll(_)),
                }
            })
            .collect();

        lir::Enum {
            description: description.clone(),
            name: name.value.clone(),
            base_type,
            variants,
        }
    }).collect()
}

fn repeat_to_method_kind(repeat: &Option<Repeat>, manifest: &mir::Manifest) -> lir::Repeat {
    match repeat {
        Some(Repeat {
            source:
                Spanned {
                    value: RepeatSource::Count(count),
                    ..
                },
            stride,
            span: _,
        }) => lir::Repeat::Count {
            count: count.get(),
            stride: stride.value,
        },
        Some(Repeat {
            source:
                Spanned {
                    value: RepeatSource::Enum(enum_name),
                    ..
                },
            stride,
            span: _,
        }) => {
            let target_enum = search_object(manifest, enum_name)
                .expect("Existence checked in MIR pass")
                .as_enum()
                .expect("checked in MIR pass");
            lir::Repeat::Enum {
                enum_name: target_enum.name.value.clone(),
                enum_variants: target_enum
                    .variants
                    .iter()
                    .map(|variant| variant.name.value.clone())
                    .collect(),
                stride: stride.value,
            }
        }
        None => lir::Repeat::None,
    }
}

#[derive(Debug, Clone)]
pub struct BorrowedBlock<'o> {
    pub description: &'o String,
    pub name: &'o Identifier<All>,
    #[expect(unused, reason = "included for completeness")]
    pub address_offset: &'o i128,
    #[expect(unused, reason = "included for completeness")]
    pub repeat: &'o Option<Repeat>,
    pub objects: &'o [mir::Object],
}

impl<'o> From<&'o mir::Block> for BorrowedBlock<'o> {
    fn from(value: &'o mir::Block) -> Self {
        let mir::Block {
            description,
            name,
            address_offset,
            repeat,
            objects,
            default_access: _,
            short_properties_span: _,
            properties_span: _,
            span: _,
        } = value;

        Self {
            description,
            name,
            address_offset,
            repeat,
            objects,
        }
    }
}

fn find_best_internal_address_type(manifest: &mir::Manifest, device: &mir::Device) -> Integer {
    let (min_address_found, max_address_found) = find_min_max_addresses(manifest, device, |_| true)
        .map(|((min, _), (max, _))| (min, max))
        .unwrap_or_default();

    let needs_signed = min_address_found < 0;
    let needs_bits = (min_address_found
        .unsigned_abs()
        .max(max_address_found.unsigned_abs())
        .add(1)
        .next_power_of_two()
        .ilog2()
        + u32::from(needs_signed))
    .next_power_of_two()
    .max(8);

    if needs_signed {
        match needs_bits {
            8 => Integer::I8,
            16 => Integer::I16,
            32 => Integer::I32,
            64 => Integer::I64,
            _ => unreachable!(),
        }
    } else {
        match needs_bits {
            8 => Integer::U8,
            16 => Integer::U16,
            32 => Integer::U32,
            64 => Integer::U64,
            _ => unreachable!(),
        }
    }
}