use std::collections::{HashMap, HashSet};
use super::TaskData;

pub struct Solution;

impl Solution {
    /// Part I: Sum the middle page numbers of all correctly ordered updates.
    pub fn solve_simple(data: &TaskData) -> i32 {
        let dependencies = Self::build_dependency_graph(&data.rules);
        data.sequences.iter()
            .filter(|seq| Self::is_valid_sequence(seq, &dependencies))
            .map(|seq| seq[seq.len() / 2]) // Take the middle page
            .sum()
    }

    /// Part II: Reorder the invalid updates according to the rules,
    /// then sum their middle page numbers.
    pub fn solve_hard(data: &TaskData) -> i32 {
        let dependencies = Self::build_dependency_graph(&data.rules);
        data.sequences.iter()
            .filter(|seq| !Self::is_valid_sequence(seq, &dependencies))
            .map(|seq| {
                let sorted = Self::topo_sort_subgraph(seq, &dependencies);
                sorted[sorted.len() / 2] // Take the middle of the corrected sequence
            })
            .sum()
    }

    /// Builds a graph where for every rule (A, B), A must appear before B.
    /// The graph stores for each page the list of pages that must appear before it.
    fn build_dependency_graph(rules: &[(i32, i32)]) -> HashMap<i32, Vec<i32>> {
        let mut graph: HashMap<i32, Vec<i32>> = HashMap::new();
        for &(a, b) in rules {
            graph.entry(b).or_default().push(a);
        }
        graph
    }

    /// Checks whether a given sequence respects all dependency rules.
    /// Returns false if any page appears after one of its required predecessors.
    fn is_valid_sequence(seq: &[i32], graph: &HashMap<i32, Vec<i32>>) -> bool {
        let mut fail_pages = HashSet::new(); // Pages that should have already appeared

        for &page in seq {
            if fail_pages.contains(&page) {
                return false; // This page appears too late
            }
            if let Some(dependencies) = graph.get(&page) {
                for &dep in dependencies {
                    fail_pages.insert(dep); // These pages must appear before `page`
                }
            }
        }
        true
    }

    /// Performs a topological sort of the pages in the sequence based on the rules.
    /// Ensures that all required pages appear before the pages that depend on them.
    fn topo_sort_subgraph(seq: &[i32], graph: &HashMap<i32, Vec<i32>>) -> Vec<i32> {
        struct Context<'a> {
            topo_order: Vec<i32>,               // Final sorted sequence
            visited: HashMap<i32, bool>,        // Tracks which pages have been visited
            graph: &'a HashMap<i32, Vec<i32>>,  // Dependency graph
        }

        /// DFS to ensure all required pages appear before the current page.
        fn dfs(page: i32, ctx: &mut Context) {
            ctx.visited.insert(page, true);
            if let Some(dependencies) = ctx.graph.get(&page) {
                for &prev in dependencies {
                    if ctx.visited.get(&prev) == Some(&false) {
                        dfs(prev, ctx);
                    }
                }
            }
            ctx.topo_order.push(page); // Add after its dependencies
        }

        // Set up the traversal context.
        let mut ctx = Context {
            topo_order: Vec::new(),
            visited: seq.iter().map(|&p| (p, false)).collect(),
            graph,
        };

        // Start DFS from each page in the sequence.
        for &page in seq {
            if !ctx.visited[&page] {
                dfs(page, &mut ctx);
            }
        }

        ctx.topo_order
    }
}