use crate::dominators::compute_dominators;
use crate::graph::ObjectGraph;
use crate::progress::ProgressGroup;
use rayon::prelude::*;

pub struct ClassRetainedRow {
    pub class_name: String,
    pub instance_count: u64,
    pub shallow_bytes: u64,
    pub retained_bytes: u64,
}

pub struct ObjectRetainedRow {
    pub addr: u64,
    pub class_name: String,
    pub shallow_bytes: u64,
    pub retained_bytes: u64,
}

pub struct RetainedAnalysis {
    pub class_rows: Vec<ClassRetainedRow>,
    pub top_objects: Vec<ObjectRetainedRow>,
    pub total_shallow: u64,
    pub total_retained: u64,
    pub reachable_objects: u64,
    pub unreachable_objects: u64,
    pub idom: Vec<u32>,
}

pub struct RetainerRow {
    pub retainer_class: String,
    pub instance_count: u64,
    pub shallow_bytes: u64,
}

pub struct ExplainedInstance {
    pub addr: u64,
    pub shallow_bytes: u64,
    pub retainer_class: String,
    pub retainer_addr: u64,
}

pub struct ClassExplanation {
    pub class_name: String,
    pub instance_count: u64,
    pub total_shallow: u64,
    pub top_instances: Vec<ExplainedInstance>,
    pub top_retainers: Vec<RetainerRow>,
}

pub fn class_matches(class_name: &str, filter: &str) -> bool {
    class_name == filter
        || class_name.ends_with(&format!("/{filter}"))
        || (filter.ends_with("[]") && class_name == filter)
}

pub fn explain_class(
    graph: &ObjectGraph,
    class_filter: &str,
    top_instances: usize,
    top_retainers: usize,
) -> Option<ClassExplanation> {
    let mut class_name = None;
    let mut instance_count = 0u64;
    let mut total_shallow = 0u64;
    let mut instances = Vec::new();
    let mut retainer_counts: rustc_hash::FxHashMap<String, (u64, u64)> =
        rustc_hash::FxHashMap::default();

    let mut preds: Vec<Vec<u32>> = vec![Vec::new(); graph.num_nodes];
    for v in 0..graph.num_nodes {
        let start = graph.offsets[v] as usize;
        let end = graph.offsets[v + 1] as usize;
        for &w in &graph.targets[start..end] {
            preds[w as usize].push(v as u32);
        }
    }
    for &r in &graph.roots {
        preds[r as usize].push(graph.super_root);
    }

    for v in 0..graph.num_nodes {
        let name = &graph.class_names[graph.object_class[v] as usize];
        if !class_matches(name, class_filter) {
            continue;
        }
        class_name = Some(name.clone());
        instance_count += 1;
        let shallow = graph.shallow[v] as u64;
        total_shallow += shallow;

        let (retainer_class, retainer_addr) = if preds[v].is_empty() {
            ("(no incoming refs)".to_string(), 0)
        } else {
            // Use the shallowest predecessor as the representative retainer for this instance.
            let &dom = preds[v]
                .iter()
                .min_by_key(|&&u| {
                    if u == graph.super_root {
                        u64::MAX
                    } else {
                        graph.shallow[u as usize] as u64
                    }
                })
                .unwrap_or(&preds[v][0]);
            if dom == graph.super_root {
                ("GC root".to_string(), 0)
            } else {
                (
                    graph.class_names[graph.object_class[dom as usize] as usize].clone(),
                    graph.addrs[dom as usize],
                )
            }
        };

        instances.push(ExplainedInstance {
            addr: graph.addrs[v],
            shallow_bytes: shallow,
            retainer_class: retainer_class.clone(),
            retainer_addr,
        });

        for &pred in &preds[v] {
            let rc = if pred == graph.super_root {
                "GC root".to_string()
            } else {
                graph.class_names[graph.object_class[pred as usize] as usize].clone()
            };
            let entry = retainer_counts.entry(rc).or_insert((0, 0));
            entry.0 += 1;
            entry.1 += shallow;
        }
    }

    let class_name = class_name?;
    instances.sort_by(|a, b| b.shallow_bytes.cmp(&a.shallow_bytes));
    instances.truncate(top_instances);

    let mut retainer_rows: Vec<RetainerRow> = retainer_counts
        .into_iter()
        .map(|(retainer_class, (instance_count, shallow_bytes))| RetainerRow {
            retainer_class,
            instance_count,
            shallow_bytes,
        })
        .collect();
    retainer_rows.sort_by(|a, b| b.shallow_bytes.cmp(&a.shallow_bytes));
    retainer_rows.truncate(top_retainers);

    Some(ClassExplanation {
        class_name,
        instance_count,
        total_shallow,
        top_instances: instances,
        top_retainers: retainer_rows,
    })
}

pub fn compute_retained(graph: &ObjectGraph, quiet: bool) -> RetainedAnalysis {
    let group = ProgressGroup::new("Computing retained sizes", 3, quiet);

    let progress = group.begin(1, "dominator tree");
    let n = graph.num_nodes;
    let super_root = graph.super_root as usize;
    let idom = compute_dominators(graph);
    progress.finish(format!("Dominator tree for {} nodes", n));

    let mut progress = group.begin(2, "accumulating retained sizes");
    let mut dom_children: Vec<Vec<u32>> = vec![Vec::new(); n + 1];
    for v in 0..n {
        let d = idom[v] as usize;
        if d != v && d != super_root {
            dom_children[d].push(v as u32);
        } else if d == super_root && v != super_root {
            dom_children[super_root].push(v as u32);
        }
    }

    let mut depth = vec![0u32; n + 1];
    let mut stack = vec![graph.super_root];
    while let Some(v) = stack.pop() {
        for &c in &dom_children[v as usize] {
            depth[c as usize] = depth[v as usize] + 1;
            stack.push(c);
        }
    }

    let mut order: Vec<usize> = (0..n).collect();
    order.sort_by_key(|&v| std::cmp::Reverse(depth[v]));

    let mut retained = vec![0u64; n];
    for &v in &order {
        retained[v] = graph.shallow[v] as u64;
        for &c in &dom_children[v] {
            retained[v] += retained[c as usize];
        }
        progress.add_nodes(1);
    }
    progress.finish(format!(
        "Retained sizes for {} objects",
        format_count(n as u64)
    ));

    let progress = group.begin(3, "aggregating by class");
    let k = graph.class_names.len();
    let mut class_count = vec![0u64; k];
    let mut class_shallow = vec![0u64; k];
    let mut class_retained = vec![0u64; k];

    let mut reachable = 0u64;
    let mut unreachable = 0u64;
    let mut total_shallow = 0u64;

    for v in 0..n {
        total_shallow += graph.shallow[v] as u64;
        let c = graph.object_class[v] as usize;
        class_count[c] += 1;
        class_shallow[c] += graph.shallow[v] as u64;
        class_retained[c] += retained[v];

        if depth[v] > 0 {
            reachable += 1;
        } else {
            unreachable += 1;
        }
    }

    let mut class_rows: Vec<ClassRetainedRow> = graph
        .class_names
        .iter()
        .enumerate()
        .map(|(i, name)| ClassRetainedRow {
            class_name: name.clone(),
            instance_count: class_count[i],
            shallow_bytes: class_shallow[i],
            retained_bytes: class_retained[i],
        })
        .collect();
    class_rows.par_sort_unstable_by(|a, b| b.retained_bytes.cmp(&a.retained_bytes));

    let mut object_rows: Vec<ObjectRetainedRow> = (0..n)
        .map(|v| ObjectRetainedRow {
            addr: graph.addrs[v],
            class_name: graph.class_names[graph.object_class[v] as usize].clone(),
            shallow_bytes: graph.shallow[v] as u64,
            retained_bytes: retained[v],
        })
        .collect();
    object_rows.par_sort_unstable_by(|a, b| b.retained_bytes.cmp(&a.retained_bytes));

    let total_retained = if !graph.roots.is_empty() {
        retained
            .iter()
            .enumerate()
            .filter(|(v, _)| depth[*v] > 0)
            .map(|(_, &r)| r)
            .max()
            .unwrap_or(0)
    } else {
        0
    };

    progress.finish(format!(
        "{} classes, {} reachable objects",
        format_count(k as u64),
        format_count(reachable)
    ));

    RetainedAnalysis {
        class_rows,
        top_objects: object_rows,
        total_shallow,
        total_retained,
        reachable_objects: reachable,
        unreachable_objects: unreachable,
        idom,
    }
}

fn format_count(n: u64) -> String {
    crate::progress::format_count(n)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::testutil::hprof::OwnedFixture;

    #[test]
    fn retained_per_object_reports_shallow_sizes() {
        let fixture = OwnedFixture::linked_list();
        let hprof = fixture.parse();
        let index = crate::index::HeapIndex::build(&hprof, true).unwrap();
        let graph = crate::graph::ObjectGraph::build(&hprof, &index, true).unwrap();
        let analysis = compute_retained(&graph, true);

        let by_addr: std::collections::HashMap<_, _> = analysis
            .top_objects
            .iter()
            .map(|o| (o.addr, o.retained_bytes))
            .collect();
        assert_eq!(by_addr[&0x3002], 24);
        assert_eq!(by_addr[&0x3001], 48);
        assert_eq!(by_addr[&0x3000], 72);
        assert_eq!(analysis.total_shallow, 72);
    }

    #[test]
    fn class_rows_aggregate_instance_counts() {
        let fixture = OwnedFixture::linked_list();
        let hprof = fixture.parse();
        let index = crate::index::HeapIndex::build(&hprof, true).unwrap();
        let graph = crate::graph::ObjectGraph::build(&hprof, &index, true).unwrap();
        let analysis = compute_retained(&graph, true);

        assert_eq!(analysis.class_rows.len(), 1);
        assert_eq!(analysis.class_rows[0].instance_count, 3);
        assert_eq!(analysis.class_rows[0].shallow_bytes, 72);
    }

    #[test]
    fn explain_class_reports_retainers() {
        let fixture = OwnedFixture::linked_list();
        let hprof = fixture.parse();
        let index = crate::index::HeapIndex::build(&hprof, true).unwrap();
        let graph = crate::graph::ObjectGraph::build(&hprof, &index, true).unwrap();
        let explanation = explain_class(&graph, "Node", 10, 10).unwrap();
        assert_eq!(explanation.instance_count, 3);
        assert!(explanation
            .top_retainers
            .iter()
            .any(|r| r.retainer_class == "GC root" || r.retainer_class.contains("Node")));
    }
}
