use std::{
    borrow::{Cow, ToOwned as _},
    collections::BTreeMap,
    format,
    num::ParseIntError,
    ops::Bound::Included,
    ptr,
    string::ToString as _,
};

use object::SectionIndex;

use crate::{
    Function, Object,
    btf::{
        Array, Btf, BtfError, BtfKind, BtfMember, BtfType, IntEncoding, MAX_SPEC_LEN, Struct,
        Union, fields_are_compatible, types_are_compatible,
    },
    generated::{
        BPF_ALU, BPF_ALU64, BPF_B, BPF_CALL, BPF_DW, BPF_H, BPF_JMP, BPF_K, BPF_LD, BPF_LDX,
        BPF_ST, BPF_STX, BPF_W, BTF_INT_SIGNED, bpf_core_relo, bpf_core_relo_kind, bpf_insn,
    },
    util::HashMap,
};

/// The error type returned by [`Object::relocate_btf`].
#[derive(thiserror::Error, Debug)]
#[error("error relocating `{section}`")]
pub struct BtfRelocationError {
    /// The function name
    pub section: String,
    #[source]
    /// The original error
    error: RelocationError,
}

/// Relocation failures
#[derive(thiserror::Error, Debug)]
enum RelocationError {
    /// I/O error
    #[error(transparent)]
    IOError(#[from] std::io::Error),

    /// Section not found
    #[error("section not found")]
    SectionNotFound,

    /// Function not found
    #[error("function not found")]
    FunctionNotFound,

    /// Invalid relocation access string
    #[error("invalid relocation access string {access_str}")]
    InvalidAccessString {
        /// The access string
        access_str: String,
    },

    /// Invalid instruction index referenced by relocation
    #[error(
        "invalid instruction index #{index} referenced by relocation #{relocation_number}, the program contains {num_instructions} instructions"
    )]
    InvalidInstructionIndex {
        /// The invalid instruction index
        index: usize,
        /// Number of instructions in the program
        num_instructions: usize,
        /// The relocation number
        relocation_number: usize,
    },

    /// Multiple candidate target types found with different memory layouts
    #[error(
        "error relocating {type_name}, multiple candidate target types found with different memory layouts: {candidates:?}"
    )]
    ConflictingCandidates {
        /// The type name
        type_name: String,
        /// The candidates
        candidates: Vec<String>,
    },

    /// Maximum nesting level reached evaluating candidate type
    #[error("maximum nesting level reached evaluating candidate type `{}`", err_type_name(.type_name))]
    MaximumNestingLevelReached {
        /// The type name
        type_name: Option<String>,
    },

    /// Invalid access string
    #[error("invalid access string `{spec}` for type `{}`: {error}", err_type_name(.type_name))]
    InvalidAccessIndex {
        /// The type name
        type_name: Option<String>,
        /// The access string
        spec: String,
        /// The index
        index: usize,
        /// The max index
        max_index: usize,
        /// The error message
        error: &'static str,
    },

    /// Relocation not valid for type
    #[error(
        "relocation #{relocation_number} of kind `{relocation_kind}` not valid for type `{type_kind}`: {error}"
    )]
    InvalidRelocationKindForType {
        /// The relocation number
        relocation_number: usize,
        /// The relocation kind
        relocation_kind: String,
        /// The type kind
        type_kind: String,
        /// The error message
        error: &'static str,
    },

    /// Invalid instruction referenced by relocation
    #[error(
        "instruction #{index} referenced by relocation #{relocation_number} is invalid: {error}"
    )]
    InvalidInstruction {
        /// The relocation number
        relocation_number: usize,
        /// The instruction index
        index: usize,
        /// The error message
        error: Cow<'static, str>,
    },

    #[error(
        "applying relocation `{kind:?}` missing target BTF info for type `{type_id}` at instruction #{ins_index}"
    )]
    MissingTargetDefinition {
        kind: RelocationKind,
        type_id: u32,
        ins_index: usize,
    },

    /// BTF error
    #[error("invalid BTF")]
    BtfError(#[from] BtfError),
}

#[expect(clippy::ref_option, reason = "more ergonomic here")]
fn err_type_name(name: &Option<String>) -> &str {
    name.as_deref().unwrap_or("[unknown name]")
}

#[derive(Copy, Clone, Debug)]
#[repr(u32)]
enum RelocationKind {
    FieldByteOffset = bpf_core_relo_kind::BPF_CORE_FIELD_BYTE_OFFSET,
    FieldByteSize = bpf_core_relo_kind::BPF_CORE_FIELD_BYTE_SIZE,
    FieldExists = bpf_core_relo_kind::BPF_CORE_FIELD_EXISTS,
    FieldSigned = bpf_core_relo_kind::BPF_CORE_FIELD_SIGNED,
    FieldLShift64 = bpf_core_relo_kind::BPF_CORE_FIELD_LSHIFT_U64,
    FieldRShift64 = bpf_core_relo_kind::BPF_CORE_FIELD_RSHIFT_U64,
    TypeIdLocal = bpf_core_relo_kind::BPF_CORE_TYPE_ID_LOCAL,
    TypeIdTarget = bpf_core_relo_kind::BPF_CORE_TYPE_ID_TARGET,
    TypeExists = bpf_core_relo_kind::BPF_CORE_TYPE_EXISTS,
    TypeSize = bpf_core_relo_kind::BPF_CORE_TYPE_SIZE,
    EnumVariantExists = bpf_core_relo_kind::BPF_CORE_ENUMVAL_EXISTS,
    EnumVariantValue = bpf_core_relo_kind::BPF_CORE_ENUMVAL_VALUE,
}

impl TryFrom<u32> for RelocationKind {
    type Error = BtfError;

    fn try_from(v: u32) -> Result<Self, Self::Error> {
        Ok(match v {
            bpf_core_relo_kind::BPF_CORE_FIELD_BYTE_OFFSET => Self::FieldByteOffset,
            bpf_core_relo_kind::BPF_CORE_FIELD_BYTE_SIZE => Self::FieldByteSize,
            bpf_core_relo_kind::BPF_CORE_FIELD_EXISTS => Self::FieldExists,
            bpf_core_relo_kind::BPF_CORE_FIELD_SIGNED => Self::FieldSigned,
            bpf_core_relo_kind::BPF_CORE_FIELD_LSHIFT_U64 => Self::FieldLShift64,
            bpf_core_relo_kind::BPF_CORE_FIELD_RSHIFT_U64 => Self::FieldRShift64,
            bpf_core_relo_kind::BPF_CORE_TYPE_ID_LOCAL => Self::TypeIdLocal,
            bpf_core_relo_kind::BPF_CORE_TYPE_ID_TARGET => Self::TypeIdTarget,
            bpf_core_relo_kind::BPF_CORE_TYPE_EXISTS => Self::TypeExists,
            bpf_core_relo_kind::BPF_CORE_TYPE_SIZE => Self::TypeSize,
            bpf_core_relo_kind::BPF_CORE_ENUMVAL_EXISTS => Self::EnumVariantExists,
            bpf_core_relo_kind::BPF_CORE_ENUMVAL_VALUE => Self::EnumVariantValue,
            kind => return Err(BtfError::InvalidRelocationKind { kind }),
        })
    }
}

#[derive(Debug, Copy, Clone)]
pub(crate) struct Relocation {
    kind: RelocationKind,
    ins_offset: usize,
    type_id: u32,
    access_str_offset: u32,
    number: usize,
}

impl Relocation {
    pub(crate) unsafe fn parse(data: &[u8], number: usize) -> Result<Self, BtfError> {
        if size_of::<bpf_core_relo>() > data.len() {
            return Err(BtfError::InvalidRelocationInfo);
        }

        let rel = unsafe { ptr::read_unaligned::<bpf_core_relo>(data.as_ptr().cast()) };

        Ok(Self {
            kind: rel.kind.try_into()?,
            ins_offset: rel.insn_off as usize,
            type_id: rel.type_id,
            access_str_offset: rel.access_str_off,
            number,
        })
    }
}

impl Object {
    /// Relocates programs inside this object file with loaded BTF info.
    pub fn relocate_btf(&mut self, target_btf: &Btf) -> Result<(), BtfRelocationError> {
        let (Some(local_btf), Some(btf_ext)) = (&self.btf, &self.btf_ext) else {
            return Ok(());
        };

        let mut candidates_cache = HashMap::<u32, Vec<Candidate<'_>>>::new();
        for (sec_name_off, relos) in btf_ext.relocations() {
            let section_name =
                local_btf
                    .string_at(*sec_name_off)
                    .map_err(|e| BtfRelocationError {
                        section: format!("section@{sec_name_off}"),
                        error: RelocationError::BtfError(e),
                    })?;

            let (section_index, _) = self
                .section_infos
                .get(&section_name.to_string())
                .ok_or_else(|| BtfRelocationError {
                    section: section_name.to_string(),
                    error: RelocationError::SectionNotFound,
                })?;

            match relocate_btf_functions(
                *section_index,
                &mut self.functions,
                relos,
                local_btf,
                target_btf,
                &mut candidates_cache,
            ) {
                Ok(()) => {}
                Err(error) => {
                    return Err(BtfRelocationError {
                        section: section_name.to_string(),
                        error,
                    });
                }
            }
        }

        Ok(())
    }
}

fn is_relocation_inside_function(
    section_index: SectionIndex,
    func: &Function,
    rel: &Relocation,
) -> bool {
    if section_index.0 != func.section_index.0 {
        return false;
    }

    let ins_offset = rel.ins_offset / size_of::<bpf_insn>();
    let func_offset = func.section_offset / size_of::<bpf_insn>();
    let func_size = func.instructions.len();

    (func_offset..func_offset + func_size).contains(&ins_offset)
}

fn function_by_relocation<'a>(
    section_index: SectionIndex,
    functions: &'a mut BTreeMap<(usize, u64), Function>,
    rel: &Relocation,
) -> Option<&'a mut Function> {
    functions
        .range_mut((
            Included(&(section_index.0, 0)),
            Included(&(section_index.0, u64::MAX)),
        ))
        .map(|(_, func)| func)
        .find(|func| is_relocation_inside_function(section_index, func, rel))
}

fn relocate_btf_functions<'target>(
    section_index: SectionIndex,
    functions: &mut BTreeMap<(usize, u64), Function>,
    relos: &[Relocation],
    local_btf: &Btf,
    target_btf: &'target Btf,
    candidates_cache: &mut HashMap<u32, Vec<Candidate<'target>>>,
) -> Result<(), RelocationError> {
    let mut last_function_opt: Option<&mut Function> = None;

    for rel in relos {
        let function = match last_function_opt.take() {
            Some(func) if is_relocation_inside_function(section_index, func, rel) => func,
            _ => function_by_relocation(section_index, functions, rel)
                .ok_or(RelocationError::FunctionNotFound)?,
        };

        let instructions = &mut function.instructions;
        let ins_index = (rel.ins_offset - function.section_offset) / size_of::<bpf_insn>();
        if ins_index >= instructions.len() {
            return Err(RelocationError::InvalidInstructionIndex {
                index: ins_index,
                num_instructions: instructions.len(),
                relocation_number: rel.number,
            });
        }

        let local_ty = local_btf.type_by_id(rel.type_id)?;
        let local_name = &*local_btf.type_name(local_ty)?;
        let access_str = &*local_btf.string_at(rel.access_str_offset)?;
        let local_spec = AccessSpec::new(local_btf, rel.type_id, access_str, *rel)?;

        let matches = if matches!(rel.kind, RelocationKind::TypeIdLocal) {
            // we don't need to look at target types to relocate this value
            Vec::new()
        } else {
            let candidates = if let Some(cands) = candidates_cache.get(&rel.type_id) {
                cands
            } else {
                candidates_cache.insert(
                    rel.type_id,
                    find_candidates(local_ty, local_name, target_btf)?,
                );
                candidates_cache.get(&rel.type_id).unwrap()
            };

            let mut matches = Vec::new();
            for candidate in candidates {
                if let Some(candidate_spec) = match_candidate(&local_spec, candidate)? {
                    let comp_rel =
                        ComputedRelocation::new(rel, &local_spec, Some(&candidate_spec))?;
                    matches.push((candidate.name.clone(), candidate_spec, comp_rel));
                }
            }

            matches
        };

        let comp_rel = if matches.is_empty() {
            // there are no candidate matches and therefore no target_spec. This might mean
            // that matching failed, or that the relocation can be applied looking at local
            // types only (eg with EnumVariantExists, FieldExists etc)
            ComputedRelocation::new(rel, &local_spec, None)?
        } else {
            let mut matches = matches.into_iter();
            let (_, target_spec, target_comp_rel) = matches.next().unwrap();

            // if there's more than one candidate, make sure that they all resolve to the
            // same value, else the relocation is ambiguous and can't be applied
            let conflicts = matches
                .filter_map(|(cand_name, cand_spec, cand_comp_rel)| {
                    if cand_spec.bit_offset != target_spec.bit_offset {
                        return Some(cand_name);
                    } else if let (Some(cand_comp_rel_target), Some(target_comp_rel_target)) = (
                        cand_comp_rel.target.as_ref(),
                        target_comp_rel.target.as_ref(),
                    ) {
                        if cand_comp_rel_target.value != target_comp_rel_target.value {
                            return Some(cand_name);
                        }
                    }

                    None
                })
                .collect::<Vec<_>>();
            if !conflicts.is_empty() {
                return Err(RelocationError::ConflictingCandidates {
                    type_name: local_name.to_string(),
                    candidates: conflicts,
                });
            }
            target_comp_rel
        };

        comp_rel.apply(function, rel, local_btf, target_btf)?;

        last_function_opt = Some(function);
    }

    Ok(())
}

fn flavorless_name(name: &str) -> &str {
    name.split_once("___").map_or(name, |x| x.0)
}

fn find_candidates<'target>(
    local_ty: &BtfType,
    local_name: &str,
    target_btf: &'target Btf,
) -> Result<Vec<Candidate<'target>>, BtfError> {
    let mut candidates = Vec::new();
    let local_name = flavorless_name(local_name);
    let local_kind = local_ty.kind();

    // When we downgrade an ENUM64 to a UNION we still want to match the enum
    // definition recorded in the target BTF. If the sanitized type has a
    // fallback, allow ENUM64 candidates through the kind check.
    //
    // Note that we do not sanitize the target BTF so the kinds will not
    // naturally match!
    let allow_enum_match = matches!(
        local_ty,
        BtfType::Union(Union {
            enum64_fallback: Some(_),
            ..
        })
    );

    for (type_id, ty) in target_btf.types().enumerate().filter(|(_, ty)| {
        let candidate_kind = ty.kind();
        candidate_kind == local_kind || (allow_enum_match && candidate_kind == BtfKind::Enum64)
    }) {
        let name = &*target_btf.type_name(ty)?;
        if local_name != flavorless_name(name) {
            continue;
        }

        candidates.push(Candidate {
            name: name.to_owned(),
            btf: target_btf,
            _ty: ty,
            type_id: type_id as u32,
        });
    }

    Ok(candidates)
}

fn match_candidate<'target>(
    local_spec: &AccessSpec<'_>,
    candidate: &'target Candidate<'_>,
) -> Result<Option<AccessSpec<'target>>, RelocationError> {
    let mut target_spec = AccessSpec {
        btf: candidate.btf,
        root_type_id: candidate.type_id,
        relocation: local_spec.relocation,
        parts: Vec::new(),
        accessors: Vec::new(),
        bit_offset: 0,
    };

    match local_spec.relocation.kind {
        RelocationKind::TypeIdLocal
        | RelocationKind::TypeIdTarget
        | RelocationKind::TypeExists
        | RelocationKind::TypeSize => {
            if types_are_compatible(
                local_spec.btf,
                local_spec.root_type_id,
                candidate.btf,
                candidate.type_id,
            )? {
                Ok(Some(target_spec))
            } else {
                Ok(None)
            }
        }
        RelocationKind::EnumVariantExists | RelocationKind::EnumVariantValue => {
            let target_id = candidate.btf.resolve_type(candidate.type_id)?;

            let match_enum = |iterator: &mut dyn Iterator<Item = u32>| {
                // the first accessor is guaranteed to have a name by construction
                let local_variant_name = local_spec.accessors[0].name.as_ref().unwrap();

                for (index, name_offset) in iterator.enumerate() {
                    let target_variant_name = candidate.btf.string_at(name_offset)?;
                    if flavorless_name(local_variant_name) == flavorless_name(&target_variant_name)
                    {
                        target_spec.parts.push(index);
                        target_spec.accessors.push(Accessor {
                            index,
                            type_id: target_id,
                            name: None,
                        });
                        return Ok(Some(target_spec));
                    }
                }
                Ok(None)
            };

            match candidate.btf.type_by_id(target_id)? {
                BtfType::Enum(en) => {
                    match_enum(&mut en.variants.iter().map(|member| member.name_offset))
                }
                BtfType::Enum64(en) => {
                    match_enum(&mut en.variants.iter().map(|member| member.name_offset))
                }
                BtfType::Union(Union {
                    enum64_fallback: Some(fallback),
                    ..
                }) => {
                    // Local ENUM64 types become UNION types during sanitisation; the fallback retains
                    // their original variant names so we can line them up with target enums.
                    match_enum(&mut fallback.variants.iter().map(|variant| variant.name_offset))
                }
                _ => Ok(None),
            }
        }
        RelocationKind::FieldByteOffset
        | RelocationKind::FieldByteSize
        | RelocationKind::FieldExists
        | RelocationKind::FieldSigned
        | RelocationKind::FieldLShift64
        | RelocationKind::FieldRShift64 => {
            let mut target_id = candidate.type_id;
            for (i, accessor) in local_spec.accessors.iter().enumerate() {
                target_id = candidate.btf.resolve_type(target_id)?;

                if accessor.name.is_some() {
                    if let Some(next_id) = match_member(
                        local_spec.btf,
                        local_spec,
                        accessor,
                        candidate.btf,
                        target_id,
                        &mut target_spec,
                    )? {
                        target_id = next_id;
                    } else {
                        return Ok(None);
                    }
                } else {
                    // i = 0 is the base struct. for i > 0, we need to potentially do bounds checking
                    if i > 0 {
                        let target_ty = candidate.btf.type_by_id(target_id)?;
                        let BtfType::Array(Array { array, .. }) = target_ty else {
                            return Ok(None);
                        };

                        let var_len = array.len == 0 && {
                            // an array is potentially variable length if it's the last field
                            // of the parent struct and has 0 elements
                            let parent = target_spec.accessors.last().unwrap();
                            let parent_ty = candidate.btf.type_by_id(parent.type_id)?;
                            match parent_ty {
                                BtfType::Struct(s) => parent.index == s.members.len() - 1,
                                _ => false,
                            }
                        };
                        if !var_len && accessor.index >= array.len as usize {
                            return Ok(None);
                        }
                        target_id = candidate.btf.resolve_type(array.element_type)?;
                    }

                    if target_spec.parts.len() == MAX_SPEC_LEN {
                        return Err(RelocationError::MaximumNestingLevelReached {
                            type_name: Some(candidate.name.clone()),
                        });
                    }

                    target_spec.parts.push(accessor.index);
                    target_spec.accessors.push(Accessor {
                        index: accessor.index,
                        type_id: target_id,
                        name: None,
                    });
                    target_spec.bit_offset +=
                        accessor.index * candidate.btf.type_size(target_id)? * 8;
                }
            }
            Ok(Some(target_spec))
        }
    }
}

fn match_member<'target>(
    local_btf: &Btf,
    local_spec: &AccessSpec<'_>,
    local_accessor: &Accessor,
    target_btf: &'target Btf,
    target_id: u32,
    target_spec: &mut AccessSpec<'target>,
) -> Result<Option<u32>, RelocationError> {
    let local_ty = local_btf.type_by_id(local_accessor.type_id)?;
    let local_member = match local_ty {
        // this won't panic, bounds are checked when local_spec is built in AccessSpec::new
        BtfType::Struct(s) => &s.members[local_accessor.index],
        BtfType::Union(u) => &u.members[local_accessor.index],
        #[expect(clippy::unreachable, reason = "accessors are validated upfront")]
        local_ty => unreachable!("unexpected type {:?}", local_ty),
    };

    let local_name = &*local_btf.string_at(local_member.name_offset)?;
    let target_id = target_btf.resolve_type(target_id)?;
    let target_ty = target_btf.type_by_id(target_id)?;

    let target_members: Vec<&BtfMember> = match target_ty.members() {
        Some(members) => members.collect(),
        // not a fields type, no match
        None => return Ok(None),
    };

    for (index, target_member) in target_members.iter().enumerate() {
        if target_spec.parts.len() == MAX_SPEC_LEN {
            let root_ty = target_spec.btf.type_by_id(target_spec.root_type_id)?;
            return Err(RelocationError::MaximumNestingLevelReached {
                type_name: target_spec.btf.err_type_name(root_ty),
            });
        }

        // this will not panic as we've already established these are fields types
        let bit_offset = target_ty.member_bit_offset(target_member).unwrap();
        let target_name = &*target_btf.string_at(target_member.name_offset)?;

        if target_name.is_empty() {
            let ret = match_member(
                local_btf,
                local_spec,
                local_accessor,
                target_btf,
                target_member.btf_type,
                target_spec,
            )?;
            if ret.is_some() {
                target_spec.bit_offset += bit_offset;
                target_spec.parts.push(index);
                return Ok(ret);
            }
        } else if local_name == target_name {
            if fields_are_compatible(
                local_spec.btf,
                local_member.btf_type,
                target_btf,
                target_member.btf_type,
            )? {
                target_spec.bit_offset += bit_offset;
                target_spec.parts.push(index);
                target_spec.accessors.push(Accessor {
                    type_id: target_id,
                    index,
                    name: Some(target_name.to_owned()),
                });
                return Ok(Some(target_member.btf_type));
            }
            return Ok(None);
        }
    }

    Ok(None)
}

#[derive(Debug)]
struct AccessSpec<'a> {
    btf: &'a Btf,
    root_type_id: u32,
    parts: Vec<usize>,
    accessors: Vec<Accessor>,
    relocation: Relocation,
    bit_offset: usize,
}

impl<'a> AccessSpec<'a> {
    fn new(
        btf: &'a Btf,
        root_type_id: u32,
        spec: &str,
        relocation: Relocation,
    ) -> Result<Self, RelocationError> {
        let parts = spec
            .split(':')
            .map(str::parse::<usize>)
            .collect::<Result<Vec<_>, _>>()
            .map_err(
                |ParseIntError { .. }| RelocationError::InvalidAccessString {
                    access_str: spec.to_string(),
                },
            )?;

        let mut type_id = btf.resolve_type(root_type_id)?;
        let ty = btf.type_by_id(type_id)?;

        let spec = match relocation.kind {
            RelocationKind::TypeIdLocal
            | RelocationKind::TypeIdTarget
            | RelocationKind::TypeExists
            | RelocationKind::TypeSize => {
                if parts != [0] {
                    return Err(RelocationError::InvalidAccessString {
                        access_str: spec.to_string(),
                    });
                }
                AccessSpec {
                    btf,
                    root_type_id,
                    relocation,
                    parts,
                    accessors: Vec::new(),
                    bit_offset: 0,
                }
            }
            RelocationKind::EnumVariantExists | RelocationKind::EnumVariantValue => {
                let index = || match parts.as_slice() {
                    [index] => Ok(*index),
                    _ => Err(RelocationError::InvalidAccessString {
                        access_str: spec.to_string(),
                    }),
                };

                let (n_variants, name_offset, index) = match ty {
                    BtfType::Enum(en) => {
                        let index = index()?;
                        (
                            en.variants.len(),
                            en.variants.get(index).map(|v| v.name_offset),
                            index,
                        )
                    }
                    BtfType::Enum64(en) => {
                        let index = index()?;
                        (
                            en.variants.len(),
                            en.variants.get(index).map(|v| v.name_offset),
                            index,
                        )
                    }
                    BtfType::Union(Union {
                        enum64_fallback: Some(fallback),
                        ..
                    }) => {
                        let index = index()?;
                        (
                            fallback.variants.len(),
                            fallback.variants.get(index).map(|v| v.name_offset),
                            index,
                        )
                    }
                    _ => {
                        return Err(RelocationError::InvalidRelocationKindForType {
                            relocation_number: relocation.number,
                            relocation_kind: format!("{:?}", relocation.kind),
                            type_kind: format!("{:?}", ty.kind()),
                            error: "enum relocation on non-enum type",
                        });
                    }
                };
                let name_offset =
                    name_offset.ok_or_else(|| RelocationError::InvalidAccessIndex {
                        type_name: btf.err_type_name(ty),
                        spec: spec.to_string(),
                        index,
                        max_index: n_variants,
                        error: "tried to access nonexistent enum variant",
                    })?;
                let name = btf.string_at(name_offset)?;
                let accessors = vec![Accessor {
                    type_id,
                    index,
                    name: Some(name.to_string()),
                }];

                AccessSpec {
                    btf,
                    root_type_id,
                    relocation,
                    parts,
                    accessors,
                    bit_offset: 0,
                }
            }

            RelocationKind::FieldByteOffset
            | RelocationKind::FieldByteSize
            | RelocationKind::FieldExists
            | RelocationKind::FieldSigned
            | RelocationKind::FieldLShift64
            | RelocationKind::FieldRShift64 => {
                let mut accessors = vec![Accessor {
                    type_id,
                    index: parts[0],
                    name: None,
                }];
                let mut bit_offset = accessors[0].index * btf.type_size(type_id)?;
                for index in parts.iter().skip(1).copied() {
                    type_id = btf.resolve_type(type_id)?;
                    let ty = btf.type_by_id(type_id)?;

                    match ty {
                        BtfType::Struct(Struct { members, .. })
                        | BtfType::Union(Union { members, .. }) => {
                            if index >= members.len() {
                                return Err(RelocationError::InvalidAccessIndex {
                                    type_name: btf.err_type_name(ty),
                                    spec: spec.to_string(),
                                    index,
                                    max_index: members.len(),
                                    error: "out of bounds struct or union access",
                                });
                            }

                            let member = &members[index];
                            bit_offset += ty.member_bit_offset(member).unwrap();

                            if member.name_offset != 0 {
                                accessors.push(Accessor {
                                    type_id,
                                    index,
                                    name: Some(btf.string_at(member.name_offset)?.to_string()),
                                });
                            }

                            type_id = member.btf_type;
                        }

                        BtfType::Array(Array { array, .. }) => {
                            type_id = btf.resolve_type(array.element_type)?;
                            let var_len = array.len == 0 && {
                                // an array is potentially variable length if it's the last field
                                // of the parent struct and has 0 elements
                                let parent = accessors.last().unwrap();
                                let parent_ty = btf.type_by_id(parent.type_id)?;
                                match parent_ty {
                                    BtfType::Struct(s) => index == s.members.len() - 1,
                                    _ => false,
                                }
                            };
                            if !var_len && index >= array.len as usize {
                                return Err(RelocationError::InvalidAccessIndex {
                                    type_name: btf.err_type_name(ty),
                                    spec: spec.to_string(),
                                    index,
                                    max_index: array.len as usize,
                                    error: "array index out of bounds",
                                });
                            }
                            accessors.push(Accessor {
                                type_id,
                                index,
                                name: None,
                            });
                            let size = btf.type_size(type_id)?;
                            bit_offset += index * size * 8;
                        }
                        rel_kind => {
                            return Err(RelocationError::InvalidRelocationKindForType {
                                relocation_number: relocation.number,
                                relocation_kind: format!("{rel_kind:?}"),
                                type_kind: format!("{:?}", ty.kind()),
                                error: "field relocation on a type that doesn't have fields",
                            });
                        }
                    }
                }

                AccessSpec {
                    btf,
                    root_type_id,
                    parts,
                    accessors,
                    relocation,
                    bit_offset,
                }
            }
        };

        Ok(spec)
    }
}

#[derive(Debug)]
struct Accessor {
    type_id: u32,
    index: usize,
    name: Option<String>,
}

#[derive(Debug)]
struct Candidate<'a> {
    name: String,
    btf: &'a Btf,
    _ty: &'a BtfType,
    type_id: u32,
}

#[derive(Debug)]
struct ComputedRelocation {
    local: ComputedRelocationValue,
    target: Option<ComputedRelocationValue>,
}

#[derive(Clone, Copy, Debug)]
struct ComputedRelocationValue {
    value: u64,
    size: u32,
    type_id: Option<u32>,
}

fn poison_insn(ins: &mut bpf_insn) {
    ins.code = (BPF_JMP | BPF_CALL) as u8;
    ins.set_dst_reg(0);
    ins.set_src_reg(0);
    ins.off = 0;
    ins.imm = 0xBAD2310;
}

impl ComputedRelocation {
    fn new(
        rel: &Relocation,
        local_spec: &AccessSpec<'_>,
        target_spec: Option<&AccessSpec<'_>>,
    ) -> Result<Self, RelocationError> {
        let ret = match rel.kind {
            RelocationKind::FieldByteOffset
            | RelocationKind::FieldByteSize
            | RelocationKind::FieldExists
            | RelocationKind::FieldSigned
            | RelocationKind::FieldLShift64
            | RelocationKind::FieldRShift64 => Self {
                local: Self::compute_field_relocation(rel, Some(local_spec))?,
                target: Self::compute_field_relocation(rel, target_spec).ok(),
            },
            RelocationKind::TypeIdLocal
            | RelocationKind::TypeIdTarget
            | RelocationKind::TypeExists
            | RelocationKind::TypeSize => Self {
                local: Self::compute_type_relocation(rel, local_spec, target_spec)?,
                target: Self::compute_type_relocation(rel, local_spec, target_spec).ok(),
            },
            RelocationKind::EnumVariantExists | RelocationKind::EnumVariantValue => Self {
                local: Self::compute_enum_relocation(rel, Some(local_spec))?,
                target: Self::compute_enum_relocation(rel, target_spec).ok(),
            },
        };

        Ok(ret)
    }

    fn apply(
        &self,
        function: &mut Function,
        rel: &Relocation,
        local_btf: &Btf,
        target_btf: &Btf,
    ) -> Result<(), RelocationError> {
        let instructions = &mut function.instructions;
        let num_instructions = instructions.len();
        let ins_index = (rel.ins_offset - function.section_offset) / size_of::<bpf_insn>();
        let ins =
            instructions
                .get_mut(ins_index)
                .ok_or(RelocationError::InvalidInstructionIndex {
                    index: rel.ins_offset,
                    num_instructions,
                    relocation_number: rel.number,
                })?;

        let Some(target) = self.target.as_ref() else {
            let is_ld_imm64 = ins.code == (BPF_LD | BPF_DW) as u8;

            poison_insn(ins);

            if is_ld_imm64 {
                let next_ins = instructions.get_mut(ins_index + 1).ok_or(
                    RelocationError::InvalidInstructionIndex {
                        index: (ins_index + 1) * size_of::<bpf_insn>(),
                        num_instructions,
                        relocation_number: rel.number,
                    },
                )?;

                poison_insn(next_ins);
            }

            return Ok(());
        };

        let class = u32::from(ins.code & 0x07);

        let target_value = target.value;

        match class {
            BPF_ALU | BPF_ALU64 => {
                let src_reg = ins.src_reg();
                if src_reg != BPF_K as u8 {
                    return Err(RelocationError::InvalidInstruction {
                        relocation_number: rel.number,
                        index: ins_index,
                        error: format!("invalid src_reg={src_reg:x} expected {BPF_K:x}").into(),
                    });
                }

                ins.imm = target_value as i32;
            }
            BPF_LDX | BPF_ST | BPF_STX => {
                if target_value > i16::MAX as u64 {
                    return Err(RelocationError::InvalidInstruction {
                        relocation_number: rel.number,
                        index: ins_index,
                        error: format!("value `{target_value}` overflows 16 bits offset field")
                            .into(),
                    });
                }

                ins.off = target_value as i16;

                if self.local.size != target.size {
                    let local_ty = local_btf.type_by_id(self.local.type_id.unwrap())?;
                    let target_ty = target_btf.type_by_id(target.type_id.unwrap())?;
                    let unsigned = |info: u32| ((info >> 24) & 0x0F) & BTF_INT_SIGNED == 0;
                    match (local_ty, target_ty) {
                        (BtfType::Ptr(_), BtfType::Ptr(_)) => {}
                        (BtfType::Int(local), BtfType::Int(target))
                            if unsigned(local.data) && unsigned(target.data) => {}
                        _ => {
                            return Err(RelocationError::InvalidInstruction {
                                relocation_number: rel.number,
                                index: ins_index,
                                error: format!(
                                    "original type {} has size {} but target type {} has size {}",
                                    err_type_name(&local_btf.err_type_name(local_ty)),
                                    self.local.size,
                                    err_type_name(&target_btf.err_type_name(target_ty)),
                                    target.size,
                                )
                                .into(),
                            });
                        }
                    }

                    let size = match target.size {
                        8 => BPF_DW,
                        4 => BPF_W,
                        2 => BPF_H,
                        1 => BPF_B,
                        size => {
                            return Err(RelocationError::InvalidInstruction {
                                relocation_number: rel.number,
                                index: ins_index,
                                error: format!("invalid target size {size}").into(),
                            });
                        }
                    } as u8;
                    ins.code = ins.code & 0xE0 | size | ins.code & 0x07;
                }
            }
            BPF_LD => {
                ins.imm = target_value as i32;
                let next_ins = instructions.get_mut(ins_index + 1).ok_or(
                    RelocationError::InvalidInstructionIndex {
                        index: ins_index + 1,
                        num_instructions,
                        relocation_number: rel.number,
                    },
                )?;

                next_ins.imm = (target_value >> 32) as i32;
            }
            class => {
                return Err(RelocationError::InvalidInstruction {
                    relocation_number: rel.number,
                    index: ins_index,
                    error: format!("invalid instruction class {class:x}").into(),
                });
            }
        }

        Ok(())
    }

    fn compute_enum_relocation(
        rel: &Relocation,
        spec: Option<&AccessSpec<'_>>,
    ) -> Result<ComputedRelocationValue, RelocationError> {
        let value = match (rel.kind, spec) {
            (RelocationKind::EnumVariantExists, spec) => u64::from(spec.is_some()),
            (RelocationKind::EnumVariantValue, Some(spec)) => {
                let accessor = &spec.accessors[0];
                match spec.btf.type_by_id(accessor.type_id)? {
                    BtfType::Enum(en) => {
                        let value = en.variants[accessor.index].value;
                        if en.is_signed() {
                            value as i32 as u64
                        } else {
                            u64::from(value)
                        }
                    }
                    BtfType::Enum64(en) => {
                        let variant = &en.variants[accessor.index];
                        (u64::from(variant.value_high) << 32) | u64::from(variant.value_low)
                    }
                    BtfType::Union(Union {
                        enum64_fallback: Some(fallback),
                        ..
                    }) => {
                        let variant = &fallback.variants[accessor.index];
                        if fallback.signed {
                            (variant.value as i64) as u64
                        } else {
                            variant.value
                        }
                    }
                    // candidate selection ensures that rel_kind == local_kind == target_kind
                    #[expect(clippy::unreachable, reason = "candidate selection enforces this")]
                    typ => unreachable!("unexpected type {typ:?}"),
                }
            }
            _ => {
                return Err(RelocationError::MissingTargetDefinition {
                    kind: rel.kind,
                    type_id: rel.type_id,
                    ins_index: rel.ins_offset / size_of::<bpf_insn>(),
                });
            }
        };

        Ok(ComputedRelocationValue {
            value,
            size: 0,
            type_id: None,
        })
    }

    fn compute_field_relocation(
        rel: &Relocation,
        spec: Option<&AccessSpec<'_>>,
    ) -> Result<ComputedRelocationValue, RelocationError> {
        if matches!(rel.kind, RelocationKind::FieldExists) {
            // this is the bpf_preserve_field_info(member_access, FIELD_EXISTENCE) case. If we
            // managed to build a spec, it means the field exists.
            return Ok(ComputedRelocationValue {
                value: u64::from(spec.is_some()),
                size: 0,
                type_id: None,
            });
        }

        let Some(spec) = spec else {
            return Err(RelocationError::MissingTargetDefinition {
                kind: rel.kind,
                type_id: rel.type_id,
                ins_index: rel.ins_offset / size_of::<bpf_insn>(),
            });
        };

        let accessor = spec.accessors.last().unwrap();
        if accessor.name.is_none() {
            // the last accessor is unnamed, meaning that this is an array access
            return match rel.kind {
                RelocationKind::FieldByteOffset => Ok(ComputedRelocationValue {
                    value: (spec.bit_offset / 8) as u64,
                    size: spec.btf.type_size(accessor.type_id)? as u32,
                    type_id: Some(accessor.type_id),
                }),
                RelocationKind::FieldByteSize => Ok(ComputedRelocationValue {
                    value: spec.btf.type_size(accessor.type_id)? as u64,
                    size: 0,
                    type_id: Some(accessor.type_id),
                }),
                rel_kind => {
                    let ty = spec.btf.type_by_id(accessor.type_id)?;
                    return Err(RelocationError::InvalidRelocationKindForType {
                        relocation_number: rel.number,
                        relocation_kind: format!("{rel_kind:?}"),
                        type_kind: format!("{:?}", ty.kind()),
                        error: "invalid relocation kind for array type",
                    });
                }
            };
        }

        let ty = spec.btf.type_by_id(accessor.type_id)?;
        let (ll_ty, member) = match ty {
            BtfType::Struct(t) => (ty, &t.members[accessor.index]),
            BtfType::Union(t) => (ty, &t.members[accessor.index]),
            _ => {
                return Err(RelocationError::InvalidRelocationKindForType {
                    relocation_number: rel.number,
                    relocation_kind: format!("{:?}", rel.kind),
                    type_kind: format!("{:?}", ty.kind()),
                    error: "field relocation on a type that doesn't have fields",
                });
            }
        };

        let bit_off = spec.bit_offset as u32;
        let member_type_id = spec.btf.resolve_type(member.btf_type)?;
        let member_ty = spec.btf.type_by_id(member_type_id)?;

        let mut byte_size;
        let mut byte_off;
        let mut bit_size = ll_ty.member_bit_field_size(member).unwrap() as u32;
        let is_bitfield = bit_size > 0;
        if is_bitfield {
            // find out the smallest int size to load the bitfield
            byte_size = member_ty.size().unwrap();
            byte_off = bit_off / 8 / byte_size * byte_size;
            while bit_off + bit_size - byte_off * 8 > byte_size * 8 {
                if byte_size >= 8 {
                    // the bitfield is larger than 8 bytes!?
                    return Err(BtfError::InvalidTypeInfo.into());
                }
                byte_size *= 2;
                byte_off = bit_off / 8 / byte_size * byte_size;
            }
        } else {
            byte_size = spec.btf.type_size(member_type_id)? as u32;
            bit_size = byte_size * 8;
            byte_off = spec.bit_offset as u32 / 8;
        }

        let mut value = ComputedRelocationValue {
            value: 0,
            size: 0,
            type_id: None,
        };

        match rel.kind {
            RelocationKind::FieldByteOffset => {
                value.value = u64::from(byte_off);
                if !is_bitfield {
                    value.size = byte_size;
                    value.type_id = Some(member_type_id);
                }
            }
            RelocationKind::FieldByteSize => {
                value.value = u64::from(byte_size);
            }
            RelocationKind::FieldSigned => match member_ty {
                BtfType::Enum(en) => value.value = u64::from(en.is_signed()),
                BtfType::Enum64(en) => value.value = u64::from(en.is_signed()),
                BtfType::Int(i) => value.value = i.encoding() as u64 & IntEncoding::Signed as u64,
                _ => (),
            },
            RelocationKind::FieldLShift64 => {
                value.value = if cfg!(target_endian = "little") {
                    64 - u64::from(bit_off + bit_size - byte_off * 8)
                } else {
                    u64::from((8 - byte_size) * 8 + (bit_off - byte_off * 8))
                }
            }
            RelocationKind::FieldRShift64 => {
                value.value = 64 - u64::from(bit_size);
            }
            kind @ (RelocationKind::FieldExists
            | RelocationKind::TypeIdLocal
            | RelocationKind::TypeIdTarget
            | RelocationKind::TypeExists
            | RelocationKind::TypeSize
            | RelocationKind::EnumVariantExists
            | RelocationKind::EnumVariantValue) => {
                #[expect(
                    clippy::unreachable,
                    reason = "unexpected relocation kinds indicate an internal bug"
                )]
                {
                    unreachable!("unexpected relocation kind {:?}", kind)
                }
            }
        }

        Ok(value)
    }

    fn compute_type_relocation(
        rel: &Relocation,
        local_spec: &AccessSpec<'_>,
        target_spec: Option<&AccessSpec<'_>>,
    ) -> Result<ComputedRelocationValue, RelocationError> {
        let value = match (rel.kind, target_spec) {
            (RelocationKind::TypeIdLocal, _) => u64::from(local_spec.root_type_id),
            (RelocationKind::TypeIdTarget, Some(target_spec)) => {
                u64::from(target_spec.root_type_id)
            }
            (RelocationKind::TypeExists, target_spec) => u64::from(target_spec.is_some()),
            (RelocationKind::TypeSize, Some(target_spec)) => {
                target_spec.btf.type_size(target_spec.root_type_id)? as u64
            }
            _ => {
                return Err(RelocationError::MissingTargetDefinition {
                    kind: rel.kind,
                    type_id: rel.type_id,
                    ins_index: rel.ins_offset / size_of::<bpf_insn>(),
                });
            }
        };

        Ok(ComputedRelocationValue {
            value,
            size: 0,
            type_id: None,
        })
    }
}
