use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::fs;
use std::path::Path;

use anyhow::{Context, Result};
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
use raw_cpuid::CpuId;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CpuGeneration {
    Zen2,
    Zen3,
    Unknown,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RssIrqAffinity {
    pub irq: u32,
    pub cpus: Vec<usize>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TopologyReport {
    pub cpu_vendor: String,
    pub family_model: String,
    pub cpu_generation: CpuGeneration,
    pub numa_nodes: Vec<Vec<usize>>,
    pub ccx_groups: Vec<Vec<usize>>,
    pub rss_irq_map: Vec<RssIrqAffinity>,
    pub selected_cores: Vec<usize>,
}

impl TopologyReport {
    pub fn to_json(&self) -> String {
        let numa_nodes = self
            .numa_nodes
            .iter()
            .map(|node| format!("[{}]", join_usize(node)))
            .collect::<Vec<_>>()
            .join(",");
        let ccx_groups = self
            .ccx_groups
            .iter()
            .map(|ccx| format!("[{}]", join_usize(ccx)))
            .collect::<Vec<_>>()
            .join(",");
        let rss_irq_map = self
            .rss_irq_map
            .iter()
            .map(|entry| {
                format!(
                    "{{\"irq\":{},\"cpus\":[{}]}}",
                    entry.irq,
                    join_usize(&entry.cpus)
                )
            })
            .collect::<Vec<_>>()
            .join(",");

        format!(
            "{{\"cpu_vendor\":\"{}\",\"family_model\":\"{}\",\"cpu_generation\":\"{}\",\"numa_nodes\":[{}],\"ccx_groups\":[{}],\"rss_irq_map\":[{}],\"selected_cores\":[{}]}}",
            self.cpu_vendor,
            self.family_model,
            cpu_generation_label(self.cpu_generation),
            numa_nodes,
            ccx_groups,
            rss_irq_map,
            join_usize(&self.selected_cores)
        )
    }
}

pub fn inspect_topology() -> Result<TopologyReport> {
    let (cpu_vendor, display_family, display_model) = read_cpuid_identity();

    let family_model = format!("0x{display_family:x}:0x{display_model:x}");
    let cpu_generation = detect_cpu_generation(&cpu_vendor, display_family);

    let numa_nodes = read_numa_nodes().unwrap_or_default();
    let ccx_groups = read_ccx_groups().unwrap_or_default();
    let rss_irq_map = read_rss_irq_map().unwrap_or_default();
    let cpuset = read_process_cpuset().unwrap_or_else(|_| default_cpuset());

    let selected_cores = select_cores(&cpuset, &ccx_groups, &rss_irq_map, 64);

    Ok(TopologyReport {
        cpu_vendor,
        family_model,
        cpu_generation,
        numa_nodes,
        ccx_groups,
        rss_irq_map,
        selected_cores,
    })
}

pub fn detect_cpu_generation(cpu_vendor: &str, display_family: u32) -> CpuGeneration {
    if cpu_vendor != "AuthenticAMD" {
        return CpuGeneration::Unknown;
    }
    match display_family {
        0x17 => CpuGeneration::Zen2,
        0x19 => CpuGeneration::Zen3,
        _ => CpuGeneration::Unknown,
    }
}

pub fn select_cores(
    cpuset: &[usize],
    ccx_groups: &[Vec<usize>],
    rss_irq_map: &[RssIrqAffinity],
    max_cores: usize,
) -> Vec<usize> {
    let cpuset_set: BTreeSet<usize> = cpuset.iter().copied().collect();
    let rss_union: BTreeSet<usize> = rss_irq_map
        .iter()
        .flat_map(|entry| entry.cpus.iter().copied())
        .collect();

    let target: BTreeSet<usize> = if rss_union.is_empty() {
        cpuset_set
    } else {
        cpuset_set
            .intersection(&rss_union)
            .copied()
            .collect::<BTreeSet<_>>()
    };

    if target.is_empty() {
        return Vec::new();
    }

    let mut selected = Vec::<usize>::new();
    let mut used = HashSet::<usize>::new();

    for group in ccx_groups {
        for &cpu in group {
            if target.contains(&cpu) && used.insert(cpu) {
                selected.push(cpu);
                if selected.len() >= max_cores {
                    return selected;
                }
            }
        }
    }

    for cpu in target {
        if used.insert(cpu) {
            selected.push(cpu);
            if selected.len() >= max_cores {
                break;
            }
        }
    }

    selected
}

pub fn parse_cpu_list(value: &str) -> Result<Vec<usize>> {
    let mut cpus = BTreeSet::<usize>::new();
    for token in value.trim().split(',').filter(|s| !s.trim().is_empty()) {
        let token = token.trim();
        if let Some((start, end)) = token.split_once('-') {
            let start: usize = start
                .trim()
                .parse()
                .with_context(|| format!("invalid cpu list start: {token}"))?;
            let end: usize = end
                .trim()
                .parse()
                .with_context(|| format!("invalid cpu list end: {token}"))?;
            if start > end {
                anyhow::bail!("invalid cpu range {}-{}", start, end);
            }
            for cpu in start..=end {
                cpus.insert(cpu);
            }
        } else {
            let cpu: usize = token
                .parse()
                .with_context(|| format!("invalid cpu list token: {token}"))?;
            cpus.insert(cpu);
        }
    }
    Ok(cpus.into_iter().collect())
}

fn read_numa_nodes() -> Result<Vec<Vec<usize>>> {
    let node_root = Path::new("/sys/devices/system/node");
    if !node_root.exists() {
        return Ok(Vec::new());
    }

    let mut nodes = Vec::<Vec<usize>>::new();
    for entry in fs::read_dir(node_root).context("failed reading numa node root")? {
        let entry = entry.context("failed reading numa node entry")?;
        let name = entry.file_name();
        let name = name.to_string_lossy();
        if !name.starts_with("node") {
            continue;
        }
        let cpulist = entry.path().join("cpulist");
        let text = match fs::read_to_string(&cpulist) {
            Ok(text) => text,
            Err(_) => continue,
        };
        let parsed = parse_cpu_list(text.trim())?;
        if !parsed.is_empty() {
            nodes.push(parsed);
        }
    }
    nodes.sort();
    Ok(nodes)
}

fn read_ccx_groups() -> Result<Vec<Vec<usize>>> {
    let cpu_root = Path::new("/sys/devices/system/cpu");
    if !cpu_root.exists() {
        return Ok(Vec::new());
    }

    let mut ccx_map = BTreeMap::<String, BTreeSet<usize>>::new();
    for entry in fs::read_dir(cpu_root).context("failed reading cpu root")? {
        let entry = entry.context("failed reading cpu entry")?;
        let file_name = entry.file_name();
        let file_name = file_name.to_string_lossy();
        if !file_name.starts_with("cpu") {
            continue;
        }
        let cpu_id = match file_name.trim_start_matches("cpu").parse::<usize>() {
            Ok(id) => id,
            Err(_) => continue,
        };

        let cache_index3 = entry.path().join("cache/index3");
        if !cache_index3.exists() {
            continue;
        }
        let cache_id = fs::read_to_string(cache_index3.join("id"))
            .unwrap_or_else(|_| "unknown".to_string())
            .trim()
            .to_string();
        let shared_cpu_list = fs::read_to_string(cache_index3.join("shared_cpu_list"))
            .unwrap_or_else(|_| cpu_id.to_string());
        let shared = parse_cpu_list(shared_cpu_list.trim()).unwrap_or_else(|_| vec![cpu_id]);
        let group = ccx_map.entry(cache_id).or_default();
        for cpu in shared {
            group.insert(cpu);
        }
    }

    let mut groups = ccx_map
        .into_values()
        .map(|set| set.into_iter().collect::<Vec<_>>())
        .collect::<Vec<_>>();
    groups.sort();
    Ok(groups)
}

fn read_rss_irq_map() -> Result<Vec<RssIrqAffinity>> {
    let interrupts =
        fs::read_to_string("/proc/interrupts").context("failed reading /proc/interrupts")?;
    let mut map = Vec::<RssIrqAffinity>::new();

    for line in interrupts.lines() {
        let trimmed = line.trim_start();
        let Some((irq_text, rest)) = trimmed.split_once(':') else {
            continue;
        };
        let irq = match irq_text.trim().parse::<u32>() {
            Ok(irq) => irq,
            Err(_) => continue,
        };

        let lower = rest.to_ascii_lowercase();
        if !(lower.contains("eth")
            || lower.contains("ens")
            || lower.contains("enp")
            || lower.contains("eno")
            || lower.contains("mlx")
            || lower.contains("net"))
        {
            continue;
        }

        let affinity_path = format!("/proc/irq/{irq}/smp_affinity_list");
        let affinity = match fs::read_to_string(&affinity_path) {
            Ok(text) => text,
            Err(_) => continue,
        };
        let cpus = parse_cpu_list(affinity.trim()).unwrap_or_default();
        if cpus.is_empty() {
            continue;
        }
        map.push(RssIrqAffinity { irq, cpus });
    }

    map.sort_by_key(|entry| entry.irq);
    Ok(map)
}

fn read_process_cpuset() -> Result<Vec<usize>> {
    let status =
        fs::read_to_string("/proc/self/status").context("failed reading /proc/self/status")?;
    for line in status.lines() {
        if let Some(value) = line.strip_prefix("Cpus_allowed_list:\t") {
            return parse_cpu_list(value.trim());
        }
    }
    anyhow::bail!("Cpus_allowed_list not found in /proc/self/status")
}

fn default_cpuset() -> Vec<usize> {
    let n = std::thread::available_parallelism()
        .map(|n| n.get())
        .unwrap_or(1);
    (0..n).collect()
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
fn read_cpuid_identity() -> (String, u32, u32) {
    let cpuid = CpuId::new();
    let cpu_vendor = cpuid
        .get_vendor_info()
        .map(|v| v.as_str().to_string())
        .unwrap_or_else(|| "unknown".to_string());

    let (display_family, display_model) = cpuid
        .get_feature_info()
        .map(|f| {
            let family = f.family_id() as u32;
            let model = f.model_id() as u32;
            let ext_family = f.extended_family_id() as u32;
            let ext_model = f.extended_model_id() as u32;
            let resolved_family = if family == 0xF {
                family + ext_family
            } else {
                family
            };
            let resolved_model = if family == 0x6 || family == 0xF {
                (ext_model << 4) + model
            } else {
                model
            };
            (resolved_family, resolved_model)
        })
        .unwrap_or((0, 0));

    (cpu_vendor, display_family, display_model)
}

#[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
fn read_cpuid_identity() -> (String, u32, u32) {
    ("unknown".to_string(), 0, 0)
}

fn cpu_generation_label(value: CpuGeneration) -> &'static str {
    match value {
        CpuGeneration::Zen2 => "Zen2",
        CpuGeneration::Zen3 => "Zen3",
        CpuGeneration::Unknown => "Unknown",
    }
}

fn join_usize(values: &[usize]) -> String {
    values
        .iter()
        .map(|v| v.to_string())
        .collect::<Vec<_>>()
        .join(",")
}

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

    #[test]
    fn parse_cpu_list_supports_ranges_and_singletons() {
        let parsed = parse_cpu_list("0-3,8,10-12").expect("cpu list parse should succeed");
        assert_eq!(parsed, vec![0, 1, 2, 3, 8, 10, 11, 12]);
    }

    #[test]
    fn detect_cpu_generation_maps_amd_families() {
        assert_eq!(
            detect_cpu_generation("AuthenticAMD", 0x17),
            CpuGeneration::Zen2
        );
        assert_eq!(
            detect_cpu_generation("AuthenticAMD", 0x19),
            CpuGeneration::Zen3
        );
        assert_eq!(
            detect_cpu_generation("AuthenticAMD", 0x1A),
            CpuGeneration::Unknown
        );
    }

    #[test]
    fn select_cores_prefers_rss_intersection() {
        let cpuset = vec![0, 1, 2, 3, 4, 5];
        let ccx_groups = vec![vec![0, 1, 2], vec![3, 4, 5]];
        let rss = vec![RssIrqAffinity {
            irq: 44,
            cpus: vec![1, 2, 4],
        }];
        let selected = select_cores(&cpuset, &ccx_groups, &rss, 64);
        assert_eq!(selected, vec![1, 2, 4]);
    }
}
