use super::gpu::Buffer;
use crate::model::ModelChoice;
use serde_json::{Value, json};
use std::cmp::Ordering;
use std::fs::{self, File};
use std::path::PathBuf;

const CACHE_CAPS: [usize; 10] = [1, 2, 4, 8, 16, 32, 64, 128, 256, 384];

#[derive(Default)]
struct LayerProfile {
    records: u64,
    counts: Vec<u64>,
    weights: Vec<f64>,
    caches: Vec<Vec<i32>>,
    cache_hits: Vec<u64>,
    cache_weight_hits: Vec<f64>,
    previous: Option<(u32, Vec<i32>)>,
    adjacent_pairs: u64,
    adjacent_overlap: f64,
    adjacent_jaccard: f64,
    hash_router: bool,
}

pub(super) struct ExpertProfile {
    path: PathBuf,
    model: ModelChoice,
    experts: usize,
    used: usize,
    caps: Vec<usize>,
    layers: Vec<LayerProfile>,
}

impl ExpertProfile {
    pub(super) fn new(
        path: Option<&str>,
        model: ModelChoice,
        layers: u32,
        experts: u64,
        used: u64,
    ) -> Result<Option<Self>, String> {
        let Some(path) = path else { return Ok(None) };
        if path.trim().is_empty() {
            return Err("Expert profile path cannot be empty".into());
        }
        let experts = usize::try_from(experts).map_err(|_| "expert count is too large")?;
        let used = usize::try_from(used).map_err(|_| "selected expert count is too large")?;
        let caps = CACHE_CAPS
            .into_iter()
            .filter(|cap| *cap <= experts)
            .collect::<Vec<_>>();
        let layers = (0..layers)
            .map(|_| LayerProfile {
                counts: vec![0; experts],
                weights: vec![0.0; experts],
                caches: vec![Vec::new(); caps.len()],
                cache_hits: vec![0; caps.len()],
                cache_weight_hits: vec![0.0; caps.len()],
                ..LayerProfile::default()
            })
            .collect();
        Ok(Some(Self {
            path: path.into(),
            model,
            experts,
            used,
            caps,
            layers,
        }))
    }

    pub(super) fn record(
        &mut self,
        layer: usize,
        pos: u32,
        selected: &Buffer,
        weights: &Buffer,
        rows: u32,
        hash_router: bool,
    ) -> Result<(), String> {
        let count = self
            .used
            .checked_mul(rows as usize)
            .ok_or("expert profile row count overflow")?;
        let mut ids = vec![0; count];
        let mut route_weights = vec![0.0; count];
        super::call(
            unsafe { super::gpu::ds4_gpu_end_commands() },
            "ending Metal commands for expert profiling",
        )?;
        let read = selected
            .read_i32(&mut ids)
            .and_then(|()| weights.read_f32(&mut route_weights));
        let resumed = super::call(
            unsafe { super::gpu::ds4_gpu_begin_commands() },
            "resuming Metal commands after expert profiling",
        );
        read.and(resumed)?;
        for (row, (ids, weights)) in ids
            .chunks_exact(self.used)
            .zip(route_weights.chunks_exact(self.used))
            .enumerate()
        {
            self.record_row(layer, pos + row as u32, ids, weights, hash_router)?;
        }
        Ok(())
    }

    fn record_row(
        &mut self,
        layer: usize,
        pos: u32,
        ids: &[i32],
        weights: &[f32],
        hash_router: bool,
    ) -> Result<(), String> {
        let profile = self
            .layers
            .get_mut(layer)
            .ok_or("expert profile layer is outside the model")?;
        profile.records += 1;
        profile.hash_router |= hash_router;
        if let Some((previous_pos, previous)) = &profile.previous
            && previous_pos.checked_add(1) == Some(pos)
        {
            let intersection = previous.iter().filter(|id| ids.contains(id)).count();
            profile.adjacent_pairs += 1;
            profile.adjacent_overlap += intersection as f64 / self.used as f64;
            profile.adjacent_jaccard += intersection as f64 / (2 * self.used - intersection) as f64;
        }
        profile.previous = Some((pos, ids.to_vec()));
        for (&id, &weight) in ids.iter().zip(weights) {
            let expert = usize::try_from(id)
                .ok()
                .filter(|expert| *expert < self.experts)
                .ok_or_else(|| format!("router selected invalid expert {id} at layer {layer}"))?;
            profile.counts[expert] += 1;
            profile.weights[expert] += f64::from(weight);
            for (index, (&cap, cache)) in self.caps.iter().zip(&mut profile.caches).enumerate() {
                if let Some(found) = cache.iter().position(|cached| *cached == id) {
                    profile.cache_hits[index] += 1;
                    profile.cache_weight_hits[index] += f64::from(weight);
                    cache.remove(found);
                } else if cache.len() == cap {
                    cache.pop();
                }
                cache.insert(0, id);
            }
        }
        Ok(())
    }

    pub(super) fn write(&self) -> Result<(), String> {
        let selections: u64 = self
            .layers
            .iter()
            .map(|layer| layer.counts.iter().sum::<u64>())
            .sum();
        let weight_total = self
            .layers
            .iter()
            .flat_map(|layer| &layer.weights)
            .sum::<f64>();
        let cache_summary = self
            .caps
            .iter()
            .enumerate()
            .map(|(index, cap)| {
                let hits = self
                    .layers
                    .iter()
                    .map(|layer| layer.cache_hits[index])
                    .sum::<u64>();
                let weighted = self
                    .layers
                    .iter()
                    .map(|layer| layer.cache_weight_hits[index])
                    .sum::<f64>();
                json!({
                    "n": cap,
                    "hits": hits,
                    "selections": selections,
                    "hit_rate": fraction(hits as f64, selections as f64),
                    "weighted_hit_rate": fraction(weighted, weight_total),
                })
            })
            .collect::<Vec<_>>();
        let layers = self
            .layers
            .iter()
            .enumerate()
            .map(|(index, layer)| self.layer_json(index, layer))
            .collect::<Vec<_>>();
        let value = json!({
            "source": "ds4 Metal expert locality profile",
            "model": self.model.id(),
            "layers": self.layers.len(),
            "experts": self.experts,
            "expert_used": self.used,
            "layer_records": self.layers.iter().map(|layer| layer.records).sum::<u64>(),
            "selections": selections,
            "cache_ns": self.caps,
            "cache_summary": cache_summary,
            "layers_detail": layers,
        });
        if let Some(parent) = self
            .path
            .parent()
            .filter(|parent| !parent.as_os_str().is_empty())
        {
            fs::create_dir_all(parent).map_err(|error| error.to_string())?;
        }
        let temporary = self.path.with_extension("tmp");
        let file = File::create(&temporary).map_err(|error| error.to_string())?;
        serde_json::to_writer_pretty(&file, &value).map_err(|error| error.to_string())?;
        file.sync_all().map_err(|error| error.to_string())?;
        fs::rename(temporary, &self.path).map_err(|error| error.to_string())
    }

    fn layer_json(&self, index: usize, layer: &LayerProfile) -> Value {
        let selections = layer.counts.iter().sum::<u64>();
        let total_weight = layer.weights.iter().sum::<f64>();
        let mut experts = (0..self.experts)
            .filter(|expert| layer.counts[*expert] != 0)
            .collect::<Vec<_>>();
        experts.sort_by(|a, b| {
            layer.counts[*b]
                .cmp(&layer.counts[*a])
                .then_with(|| {
                    layer.weights[*b]
                        .partial_cmp(&layer.weights[*a])
                        .unwrap_or(Ordering::Equal)
                })
                .then_with(|| a.cmp(b))
        });
        let top = experts
            .into_iter()
            .take(16)
            .map(|expert| {
                json!({
                    "id": expert,
                    "count": layer.counts[expert],
                    "pct": 100.0 * fraction(layer.counts[expert] as f64, selections as f64),
                    "weight": layer.weights[expert],
                    "weight_pct": 100.0 * fraction(layer.weights[expert], total_weight),
                })
            })
            .collect::<Vec<_>>();
        let cache = self
            .caps
            .iter()
            .enumerate()
            .map(|(cap_index, cap)| {
                json!({
                    "n": cap,
                    "hits": layer.cache_hits[cap_index],
                    "hit_rate": fraction(layer.cache_hits[cap_index] as f64, selections as f64),
                    "weighted_hit_rate": fraction(layer.cache_weight_hits[cap_index], total_weight),
                })
            })
            .collect::<Vec<_>>();
        json!({
            "layer": index,
            "hash_router": layer.hash_router,
            "records": layer.records,
            "selections": selections,
            "unique_experts": layer.counts.iter().filter(|count| **count != 0).count(),
            "avg_adjacent_overlap": fraction(layer.adjacent_overlap, layer.adjacent_pairs as f64),
            "avg_adjacent_jaccard": fraction(layer.adjacent_jaccard, layer.adjacent_pairs as f64),
            "top_experts": top,
            "cache": cache,
        })
    }
}

fn fraction(value: f64, total: f64) -> f64 {
    if total > 0.0 { value / total } else { 0.0 }
}

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

    #[test]
    fn profile_reports_locality_and_lru_hits() {
        let path = std::env::temp_dir().join(format!(
            "ds4-profile-{}-{}.json",
            std::process::id(),
            std::thread::current().name().unwrap_or("test")
        ));
        let mut profile = ExpertProfile::new(path.to_str(), ModelChoice::DeepSeekV4Flash, 1, 8, 2)
            .unwrap()
            .unwrap();
        profile
            .record_row(0, 10, &[1, 2], &[0.6, 0.4], false)
            .unwrap();
        profile
            .record_row(0, 11, &[1, 3], &[0.7, 0.3], false)
            .unwrap();
        profile.write().unwrap();
        let value: Value = serde_json::from_slice(&fs::read(&path).unwrap()).unwrap();
        assert_eq!(value["selections"], 4);
        assert_eq!(value["layers_detail"][0]["avg_adjacent_overlap"], 0.5);
        assert_eq!(value["cache_summary"][1]["hits"], 1);
        let _ = fs::remove_file(path);
    }
}
