use tici_grpc::FragProperty;

use super::shard::{HandleRange, RangeBound};

pub(crate) fn weighted_u64(value: u64, p: f64) -> u64 {
    if value == 0 || !p.is_finite() || p <= 0.0 {
        return 0;
    }
    ((value as f64) * p).round() as u64
}

fn normalize_right_pad(input: &[u8], len: usize) -> Vec<u8> {
    let mut out = input.to_vec();
    if out.len() < len {
        out.resize(len, 0);
    }
    out
}

fn trim_trailing_zeros(mut input: Vec<u8>) -> Vec<u8> {
    if input.is_empty() {
        return input;
    }
    let mut end = input.len();
    while end > 1 && input[end - 1] == 0 {
        end -= 1;
    }
    input.truncate(end);
    input
}

fn trim_leading_zeros(input: Vec<u8>) -> Vec<u8> {
    if input.is_empty() {
        return input;
    }
    let mut start = 0;
    while start + 1 < input.len() && input[start] == 0 {
        start += 1;
    }
    input[start..].to_vec()
}

fn add_be_bytes(a: &[u8], b: &[u8]) -> Vec<u8> {
    let max_len = a.len().max(b.len());
    let mut out_rev = Vec::with_capacity(max_len + 1);
    let mut carry: u16 = 0;
    for i in 0..max_len {
        let ai = if a.len() > i {
            a[a.len() - 1 - i] as u16
        } else {
            0
        };
        let bi = if b.len() > i {
            b[b.len() - 1 - i] as u16
        } else {
            0
        };
        let sum = ai + bi + carry;
        out_rev.push((sum & 0xFF) as u8);
        carry = sum >> 8;
    }
    if carry > 0 {
        out_rev.push(carry as u8);
    }
    out_rev.reverse();
    trim_leading_zeros(out_rev)
}

fn add_be_bytes_fixed_len(a: &[u8], b: &[u8]) -> Vec<u8> {
    if a.len() != b.len() {
        return vec![];
    }
    let mut out = vec![0; a.len()];
    let mut carry: u16 = 0;
    for i in (0..a.len()).rev() {
        let sum = a[i] as u16 + b[i] as u16 + carry;
        out[i] = (sum & 0xFF) as u8;
        carry = sum >> 8;
    }
    out
}

fn sub_be_bytes(a: &[u8], b: &[u8]) -> Vec<u8> {
    let max_len = a.len().max(b.len());
    let mut out = vec![0; max_len];
    let mut borrow: i16 = 0;
    for i in 0..max_len {
        let ai = if a.len() > i {
            a[a.len() - 1 - i] as i16
        } else {
            0
        };
        let bi = if b.len() > i {
            b[b.len() - 1 - i] as i16
        } else {
            0
        };
        let mut val = ai - borrow;
        if val < bi {
            val += 256;
            borrow = 1;
        } else {
            borrow = 0;
        }
        out[max_len - 1 - i] = (val - bi) as u8;
    }
    out
}

fn mul_be_bytes_u64(value: &[u8], multiplier: u64) -> Vec<u8> {
    if multiplier == 0 {
        return vec![0];
    }
    if value.is_empty() {
        return vec![0];
    }
    let mut out_rev = Vec::with_capacity(value.len() + 8);
    let mut carry: u128 = 0;
    let multiplier = multiplier as u128;
    for &byte in value.iter().rev() {
        let prod = (byte as u128) * multiplier + carry;
        out_rev.push((prod & 0xFF) as u8);
        carry = prod >> 8;
    }
    while carry > 0 {
        out_rev.push((carry & 0xFF) as u8);
        carry >>= 8;
    }
    out_rev.reverse();
    trim_leading_zeros(out_rev)
}

fn div_be_bytes_u64(value: &[u8], divisor: u64) -> Vec<u8> {
    if divisor == 0 {
        return vec![0];
    }
    if value.is_empty() {
        return vec![0];
    }
    let mut remainder: u128 = 0;
    let divisor = divisor as u128;
    let mut out = Vec::with_capacity(value.len());
    for &byte in value {
        remainder = remainder * 256 + byte as u128;
        let quo = remainder / divisor;
        remainder %= divisor;
        out.push(quo as u8);
    }
    trim_leading_zeros(out)
}

fn pad_left_to_len(value: &[u8], len: usize) -> Vec<u8> {
    if value.len() >= len {
        let extra = value.len().saturating_sub(len);
        return value[extra..].to_vec();
    }
    let mut out = vec![0; len - value.len()];
    out.extend_from_slice(value);
    out
}

pub(crate) fn weighted_average_keys(candidates: &[(Vec<u8>, u64)]) -> Option<Vec<u8>> {
    // Produces a key no longer than the longest input: normalize to max length,
    // compute, then trim trailing 0x00 bytes (never to empty).
    if candidates.is_empty() {
        return None;
    }
    let total_weight: u64 = candidates.iter().map(|(_, weight)| *weight).sum();
    if total_weight == 0 {
        return None;
    }
    let max_len = candidates
        .iter()
        .map(|(key, _)| key.len())
        .max()
        .unwrap_or(0);
    if max_len == 0 {
        return None;
    }
    let normalized: Vec<Vec<u8>> = candidates
        .iter()
        .map(|(key, _)| normalize_right_pad(key, max_len))
        .collect();
    let mut base = normalized[0].clone();
    for key in normalized.iter().skip(1) {
        if key < &base {
            base = key.clone();
        }
    }
    let mut sum = vec![0u8];
    for (idx, (_, weight)) in candidates.iter().enumerate() {
        if *weight == 0 {
            continue;
        }
        let delta = sub_be_bytes(&normalized[idx], &base);
        let product = mul_be_bytes_u64(&delta, *weight);
        sum = add_be_bytes(&sum, &product);
    }
    let avg_delta = div_be_bytes_u64(&sum, total_weight);
    let avg_delta_padded = pad_left_to_len(&avg_delta, max_len);
    let result = add_be_bytes_fixed_len(&base, &avg_delta_padded);
    if result.is_empty() {
        return None;
    }
    Some(trim_trailing_zeros(result))
}

pub(crate) fn midpoint_key(start: &[u8], end: &[u8]) -> Option<Vec<u8>> {
    // Midpoint is computed as a base-256 average with fixed max length; result
    // never grows beyond the longest input and may shrink after trimming
    // trailing 0x00 bytes.
    weighted_average_keys(&[(start.to_vec(), 1), (end.to_vec(), 1)])
}

fn overlap_range_bounds(range: &HandleRange, prop: &FragProperty) -> Option<(Vec<u8>, Vec<u8>)> {
    if !range.range_overlaps_minmax(&prop.min_handle, &prop.max_handle) {
        return None;
    }
    let frag_start = RangeBound::make_start(prop.min_handle.clone());
    let frag_end = RangeBound::make_end(prop.max_handle.clone());
    let start = if range.sortable_start_ref() > &frag_start {
        range.sortable_start()
    } else {
        frag_start
    };
    let end = if range.sortable_end_ref() < &frag_end {
        range.sortable_end()
    } else {
        frag_end
    };
    Some((start.into(), end.into()))
}

pub(crate) fn fragment_split_handle(prop: &FragProperty, range: &HandleRange) -> Option<Vec<u8>> {
    if !range.range_overlaps_minmax(&prop.min_handle, &prop.max_handle) {
        return None;
    }
    let mut usable: Vec<&[u8]> = Vec::new();
    for handle in prop.split_handles.iter() {
        let bound = RangeBound::Key(handle.clone());
        if range.contains_key_not_included(&bound) {
            usable.push(handle.as_slice());
        }
    }
    if !usable.is_empty() {
        let mid_index = usable.len() / 2;
        if usable.len() % 2 == 1 {
            return Some(usable[mid_index].to_vec());
        }
        return midpoint_key(usable[mid_index - 1], usable[mid_index]);
    }
    let (start, end) = overlap_range_bounds(range, prop)?;
    midpoint_key(&start, &end)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_base256_arithmetic_basic() {
        // Verifies base-256 add/sub/mul/div helpers behave as expected.
        assert_eq!(add_be_bytes(&[0x01, 0xFF], &[0x00, 0x01]), vec![0x02, 0x00]);
        assert_eq!(sub_be_bytes(&[0x02, 0x00], &[0x00, 0x01]), vec![0x01, 0xFF]);
        assert_eq!(mul_be_bytes_u64(&[0x01, 0x00], 2), vec![0x02, 0x00]);
        assert_eq!(div_be_bytes_u64(&[0x02, 0x00], 2), vec![0x01, 0x00]);
    }

    #[test]
    fn test_weighted_average_right_pad() {
        // Ensures right-padding normalization keeps equivalent keys stable.
        let candidates = vec![(vec![0x01], 1), (vec![0x01, 0x00], 1)];
        let avg = match weighted_average_keys(&candidates) {
            Some(value) => value,
            None => panic!("expected weighted average"),
        };
        assert_eq!(avg, vec![0x01]);
    }

    #[test]
    fn test_weighted_average_weighted_shift() {
        // Ensures weighted average shifts toward heavier candidates.
        let candidates = vec![(vec![0x10], 1), (vec![0x30], 3)];
        let avg = match weighted_average_keys(&candidates) {
            Some(value) => value,
            None => panic!("expected weighted average"),
        };
        assert_eq!(avg, vec![0x28]);
    }

    #[test]
    fn test_midpoint_key_basic() {
        // Ensures midpoint_key computes the midpoint of two handles.
        let mid = match midpoint_key(&[0x30], &[0x40]) {
            Some(value) => value,
            None => panic!("expected midpoint"),
        };
        assert_eq!(mid, vec![0x38]);
    }
}