use super::TaskData;
use super::segment_tree::{FreeChainTree, Chain};

pub struct Solution;

/// Fields: file_id, block_start, block_len
pub type MemoryMap = Vec<(usize, usize, usize)>;

impl Solution {
    /// Part I: Move individual blocks (1 sector at a time) to the leftmost free block,
    /// eliminating gaps between file blocks.
    pub fn solve_simple(data: &TaskData) -> usize {
        let mut mgr = Self::init_chain_tree(data);

        // Build a new map by moving blocks one by one, from last file to first
        let defrag_map: MemoryMap = data.disk_dump.iter().rev()
            .flat_map(|&(file_id, start, len)| {
                // Free the original file segments
                mgr.mark_free(start .. start + len);

                // Move one sector at a time to the next free position
                let mut moved = Vec::new();
                let mut remaining = len;
                while remaining > 0 {
                    // find a free block of size >= 1
                    let Chain { start: new_pos, len: free_len } =
                        mgr.find_free_chain(1)
                            .expect("No free space during simple defrag");
                    let block_len = remaining.min(free_len);
                    // occupy the block
                    mgr.mark_busy(new_pos .. new_pos + block_len);
                    moved.push((file_id, new_pos, block_len));
                    remaining -= block_len;
                }
                moved
            }).collect();

        Self::check_sum(&defrag_map)
    }

    /// Part II: Move whole files once, in order of decreasing file ID,
    /// to the leftmost free span large enough for the file (if any).
    pub fn solve_hard(data: &TaskData) -> usize {
        let mut mgr = Self::init_chain_tree(data);

        let defrag_map = data.disk_dump.iter().rev().map(|&(file_id, start, len)| 
            if let Some(Chain{start:new_pos,  .. }) = mgr.find_free_chain(len)
                .filter(|chain| chain.start < start) 
            {
                // move whole file
                mgr.mark_free(start .. start + len);
                mgr.mark_busy(new_pos .. new_pos + len);
                (file_id, new_pos, len)
            } else {
                (file_id, start, len) // no move possible
            }).collect();

        Self::check_sum(&defrag_map)
    }

    /// Initialize a FreeChainTree by marking all occupied ranges busy
    fn init_chain_tree(data: &TaskData) -> FreeChainTree {
        let mut mgr = FreeChainTree::new(data.disk_size);
        // Mark all occupied sectors busy
        for &(_, start, len) in &data.disk_dump {
            mgr.mark_busy(start .. start + len);
        }
        mgr
    }

    /// Checksum: sum over each segment: file_id * sum_{i=0 to len-1}(offset + i)
    fn check_sum(map: &MemoryMap) -> usize {
        map.iter()
            .map(|&(file_id, offset, len)| {
                // sum of arithmetic sequence:
                let seq_sum = offset * len + len * (len - 1) / 2;
                seq_sum * file_id
            }).sum()
    }
}

use std::ops::Range;

/// A contiguous chain of free sectors: starting index and length.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Chain {
    pub start: usize,
    pub len: usize,
}

/// Node storing prefix, suffix, and max free lengths for a segment.
#[derive(Clone, Debug, Default)]
struct MetricsNode {
    prefix_free: usize,
    suffix_free: usize,
    max_free: usize,
}

/// Segment tree managing free-sector chains over `n` sectors.
pub struct FreeChainTree {
    n: usize,             // original disk size
    cap: usize,           // rounded-up power-of-two capacity
    tree: Vec<MetricsNode>,
}

impl FreeChainTree {
    /// Build a new tree for `n` sectors (all free initially).
    pub fn new(n: usize) -> Self {
        let cap = n.next_power_of_two();
        let mut tree = vec![MetricsNode::default(); 2 * cap];
        // Initialize leaves corresponding to real sectors
        tree[cap..cap+n].fill(MetricsNode::new(1));
        // Build internal nodes recursively
        Self::build_internal_nodes(&mut tree, 1, 0, cap);
        FreeChainTree { n, cap, tree }
    }

    /// Mark all sectors in `range` as free.
    pub fn mark_free(&mut self, range: Range<usize>) {
        self.update_range(1, 0..self.cap, range, true);
    }

    /// Mark all sectors in `range` as busy.
    pub fn mark_busy(&mut self, range: Range<usize>) {
        self.update_range(1, 0..self.cap, range, false);
    }

    /// Find the leftmost free chain of length ≥ `k`.
    pub fn find_free_chain(&self, k: usize) -> Option<Chain> {
        if self.tree[1].max_free < k {
            return None;
        }
        self.find_rec(1, 0..self.cap, k)
    }

    /// Recursively build internal nodes from leaves.
    fn build_internal_nodes(tree: &mut [MetricsNode], idx: usize, l: usize, r: usize) {
        if r - l == 1 { return; } // leaf
        let mid = Self::mid(l, r);
        let (left, right) = Self::child_indices(idx);
        Self::build_internal_nodes(tree, left, l, mid);
        Self::build_internal_nodes(tree, right, mid, r);
        let child_len = mid - l;
        tree[idx] = MetricsNode::merge(&tree[left], &tree[right], child_len);
    }

    /// Update a segment `seg` with `range`, setting free or busy.
    fn update_range(&mut self, idx: usize, seg: Range<usize>, range: Range<usize>, free: bool) {
        if range.end <= seg.start || range.start >= seg.end {
            return;
        }
        if seg.end == seg.start + 1 {
            let len = if free { 1 } else { 0 };
            self.tree[idx] = MetricsNode::new(len);
            return;
        }
        let mid = Self::mid(seg.start, seg.end);
        let (left, right) = Self::child_indices(idx);
        self.update_range(left, seg.start..mid, range.clone(), free);
        self.update_range(right, mid..seg.end, range, free);
        let child_len = mid - seg.start;
        self.tree[idx] = MetricsNode::merge(&self.tree[left], &self.tree[right], child_len);
    }

    /// Internal recursive search for a free chain of size `≥k`.
    fn find_rec(&self, idx: usize, seg: Range<usize>, k: usize) -> Option<Chain> {
        let node = &self.tree[idx];
        if node.max_free < k {
            return None;
        }
        if node.prefix_free >= k {
            return Some(Chain { start: seg.start, len: node.prefix_free });
        }
        let mid = Self::mid(seg.start, seg.end); 
        let (left, right) = Self::child_indices(idx);

        let middle_space = self.tree[left].suffix_free + self.tree[right].prefix_free;
        if self.tree[left].max_free >= k {
            self.find_rec(left, seg.start..mid, k)
        } else if middle_space >= k {
            Some(Chain { start: mid - self.tree[left].suffix_free, len: middle_space })
        } else {
            self.find_rec(right, mid..seg.end, k)
        }
    }

    /// Compute midpoint of [l..r).
    #[inline]
    fn mid(l: usize, r: usize) -> usize {
        (l + r) >> 1
    }

    /// Compute (left, right) child indices for `idx`.
    #[inline]
    fn child_indices(idx: usize) -> (usize, usize) {
        let left = idx << 1;
        (left, left + 1)
    }
}

impl MetricsNode {
    /// Leaf constructor: all free of length `free_len`.
    fn new(free_len: usize) -> Self {
        MetricsNode { prefix_free: free_len, suffix_free: free_len, max_free: free_len }
    }

    /// Merge two children summaries into parent summary.
    fn merge(left: &MetricsNode, right: &MetricsNode, child_len: usize) -> MetricsNode {
        let prefix_free = left.prefix_free
            + if left.prefix_free == child_len { right.prefix_free } else { 0 };
        let suffix_free = right.suffix_free
            + if right.suffix_free == child_len { left.suffix_free } else { 0 };
        let middle_space = left.suffix_free + right.prefix_free;
        // max of left, right, or crossing span
        let max_free = left.max_free.max(right.max_free).max(middle_space);
        MetricsNode { prefix_free, suffix_free, max_free }
    }
}