//! Planning: split a file into weighted byte-range segments across exits.

use crate::error::FetchError;
use crate::types::{
    ByteRange, ExitId, ExitWeight, FileSize, MinSegmentBytes, SegmentIndex, SegmentsPerExit,
    TargetUrl,
};
use std::num::NonZeroUsize;

/// One planned segment: a byte range assigned to one exit.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SegmentPlan {
    pub index: SegmentIndex,
    pub range: ByteRange,
    pub exit: ExitId,
}

/// A validated download plan.
#[derive(Debug, Clone)]
pub struct FetchPlan {
    pub url: TargetUrl,
    pub file_size: FileSize,
    pub segments: Vec<SegmentPlan>,
}

/// Exits with their capacity weights.
pub type ExitSpecs = Vec<(ExitId, ExitWeight)>;

impl FetchPlan {
    /// Build a validated plan.
    ///
    /// Exits are weighted by capacity; each exit's share is further split
    /// into `segments_per_exit` sub-segments. Segments smaller than
    /// `min_segment` are merged into a neighbour so we never spawn
    /// absurdly tiny requests — unless the whole file is that small.
    pub fn build(
        url: TargetUrl,
        file_size: FileSize,
        exits: &ExitSpecs,
        segments_per_exit: SegmentsPerExit,
        min_segment: MinSegmentBytes,
    ) -> Result<Self, FetchError> {
        if exits.is_empty() {
            return Err(FetchError::InvalidPlan("no exits given".to_string()));
        }

        let mut exits = exits.clone();

        // Drop the weakest exits (smallest weight; ties: higher id goes
        // first) until every remaining exit can hold at least one
        // minimum segment.
        while exits.len() > 1
            && u128::from(file_size.value())
                < u128::from(exits.len() as u64) * u128::from(min_segment.value())
        {
            let weakest = exits
                .iter()
                .enumerate()
                .min_by_key(|(_, (id, weight))| (weight.value(), std::cmp::Reverse(id.value())))
                .map(|(pos, _)| pos)
                .unwrap_or(0);
            exits.remove(weakest);
        }
        // Deterministic boundary assignment order.
        exits.sort_by_key(|(id, _)| *id);

        // If even one segment per exit would be too small, reduce to one
        // segment per exit.
        let spe_limit = usize::try_from(
            u128::from(file_size.value())
                / (u128::from(exits.len() as u64) * u128::from(min_segment.value())),
        )
        .unwrap_or(0);
        let spe = if spe_limit >= segments_per_exit.value() {
            segments_per_exit
        } else {
            SegmentsPerExit::new(NonZeroUsize::MIN)
        };

        let mut assignments: Vec<(ByteRange, ExitId)> = Vec::new();

        if exits.len() == 1 && spe.value() == 1 {
            let only = ByteRange::try_new(0, file_size.value() - 1)?;
            assignments.push((only, exits[0].0));
        } else {
            // Weighted split across exits, then even sub-split per exit.
            let units: Vec<u128> = exits
                .iter()
                .map(|(_, w)| u128::from(w.value()) * spe.value() as u128)
                .collect();
            let total_units: u128 = units.iter().sum();

            let size = u128::from(file_size.value());
            let mut cursor: u128 = 0;
            let mut start: u64 = 0;

            for (i, units_i) in units.iter().enumerate() {
                let is_last = i == exits.len() - 1;
                let end_exclusive: u128 = if is_last {
                    size
                } else {
                    cursor + (size * units_i) / total_units
                };

                let exit_end: u64 = end_exclusive
                    .try_into()
                    .map_err(|_| FetchError::InvalidPlan("boundary overflow".to_string()))?;
                let exit_end_inclusive = exit_end.saturating_sub(1);

                if exit_end_inclusive >= start {
                    for range in even_subsplit(start, exit_end_inclusive, spe.value()) {
                        assignments.push((range, exits[i].0));
                    }
                    start = exit_end;
                }

                cursor = end_exclusive;
            }

            // The last exit covers to the end, so nothing can be left
            // over; stay honest and extend the last assignment anyway.
            if start < file_size.value() {
                match assignments.last_mut() {
                    Some((last_range, _)) => {
                        *last_range =
                            ByteRange::try_new(last_range.start(), file_size.value() - 1)?;
                    }
                    None => {
                        let only = ByteRange::try_new(0, file_size.value() - 1)?;
                        assignments.push((only, exits[0].0));
                    }
                }
            }
        }

        // Merge-fix: no segment below the minimum size (except a lone
        // tiny segment for a tiny file).
        let mut assignments = merge_tiny_segments(assignments, min_segment);

        let segments: Vec<SegmentPlan> = assignments
            .drain(..)
            .enumerate()
            .map(|(i, (range, exit))| {
                let index = u32::try_from(i)
                    .map(SegmentIndex::new)
                    .map_err(|_| FetchError::InvalidPlan("more than 2^32 segments".to_string()))?;
                Ok::<SegmentPlan, FetchError>(SegmentPlan { index, range, exit })
            })
            .collect::<Result<Vec<_>, _>>()?;

        let plan = Self {
            url,
            file_size,
            segments,
        };
        plan.validate()?;
        Ok(plan)
    }

    /// Check invariants: contiguous, gap-free, complete coverage of the
    /// file, sequential indices, non-empty ranges.
    pub fn validate(&self) -> Result<(), FetchError> {
        let mut cursor: u64 = 0;
        for (expected_index, seg) in self.segments.iter().enumerate() {
            let expected_index = u32::try_from(expected_index)
                .map_err(|_| FetchError::InvalidPlan("more than 2^32 segments".to_string()))?;
            if seg.index.value() != expected_index {
                return Err(FetchError::InvalidPlan(format!(
                    "segment index out of order: {} != {expected_index}",
                    seg.index
                )));
            }
            if seg.range.start() != cursor {
                return Err(FetchError::InvalidPlan(format!(
                    "gap or overlap at segment {}: starts at {} but cursor is at {cursor}",
                    seg.index,
                    seg.range.start()
                )));
            }
            if seg.range.is_empty() {
                return Err(FetchError::InvalidPlan(format!(
                    "segment {} is empty",
                    seg.index
                )));
            }
            cursor = seg
                .range
                .end_inclusive()
                .checked_add(1)
                .ok_or_else(|| FetchError::InvalidPlan("range end overflow".to_string()))?;
        }
        if cursor != self.file_size.value() {
            return Err(FetchError::InvalidPlan(format!(
                "plan covers up to {cursor} but file size is {}",
                self.file_size
            )));
        }
        Ok(())
    }
}

/// Split `[start, end_inclusive]` into `parts` nearly-even sub-ranges,
/// ordered; the first `remainder` parts get one extra byte.
///
/// The caller guarantees `parts >= 1` and `end_inclusive >= start`, so
/// every part is non-empty and `try_new` cannot fail.
fn even_subsplit(start: u64, end_inclusive: u64, parts: usize) -> Vec<ByteRange> {
    let total_len = end_inclusive - start + 1;
    let base = total_len / parts as u64;
    let remainder = total_len % parts as u64;

    let mut out = Vec::with_capacity(parts);
    let mut cursor = start;
    for i in 0..parts {
        let len = base + if (i as u64) < remainder { 1 } else { 0 };
        let range_end = cursor + len - 1;
        match ByteRange::try_new(cursor, range_end) {
            Ok(range) => out.push(range),
            // Unreachable by construction; keep the sub-range out instead
            // of panicking (validate() would then reject the plan).
            Err(_) => return Vec::new(),
        }
        cursor = range_end + 1;
    }
    out
}

/// Merge any segment smaller than `min_segment` into a neighbour
/// (prefer the next, fall back to the previous). A single remaining
/// segment is left as-is, however small.
fn merge_tiny_segments(
    mut assignments: Vec<(ByteRange, ExitId)>,
    min_segment: MinSegmentBytes,
) -> Vec<(ByteRange, ExitId)> {
    while let Some(pos) = assignments
        .iter()
        .position(|(range, _)| range.len() < min_segment.value())
    {
        if assignments.len() == 1 {
            break; // nothing left to merge into
        }

        let merged: (ByteRange, ExitId) = if pos + 1 < assignments.len() {
            // Merge into the next segment, keeping its exit.
            let next = &assignments[pos + 1];
            let range = match ByteRange::try_new(assignments[pos].0.start(), next.0.end_inclusive())
            {
                Ok(range) => range,
                // Unreachable: next.start == pos.end + 1, so the merged
                // range is always well-formed. Drop both segments instead
                // of panicking; validate() would reject a broken plan.
                Err(_) => break,
            };
            (range, next.1)
        } else {
            // Last segment is too small: merge into the previous one.
            let prev_pos = pos - 1;
            let prev = &assignments[prev_pos];
            let range = match ByteRange::try_new(prev.0.start(), assignments[pos].0.end_inclusive())
            {
                Ok(range) => range,
                Err(_) => break,
            };
            (range, prev.1)
        };

        match pos + 1 < assignments.len() {
            true => {
                assignments[pos + 1] = merged;
                assignments.remove(pos);
            }
            false => {
                assignments[pos - 1] = merged;
                assignments.remove(pos);
            }
        }
    }
    assignments
}