use crate::bpf::*;
use crate::energy::rapl;
use anyhow::Result;
use csv::Writer;
use libbpf_rs::OpenObject;
use perf_event::{
    events::{Cache, CacheOp, CacheResult, Event, Hardware, Software, WhichCache},
    Builder, Counter, Group,
};
use rand::seq::IteratorRandom;
use scx_utils::UserExitInfo;
use std::collections::HashMap;
use std::fs::File;
use std::mem::MaybeUninit;
use std::process;
use std::thread;
use std::time::{Duration, Instant};

const SLICE_US: u64 = 50000;
const LOG_INTERVAL_MS: u64 = 10; // Log every 1 second
                                 // const RESHUFFLE_ROUNDS: usize = 5; // Number of rounds before reshuffling counters
const RESHUFFLE_ROUNDS: usize = 1; // Number of rounds before reshuffling counters
const MAX_COUNTERS_AT_ONCE: usize = 7;

type Pid = i32;

pub struct BenchmarkScheduler<'a> {
    bpf: BpfScheduler<'a>,
    own_pid: Pid,
    log_path: String,
}

// Represents a single measurement point in time
struct Measurement {
    timestamp: Instant,
    energy: Option<f64>,
    frequency: Option<f64>,
    counter_values: Vec<Option<u64>>,
}

impl Measurement {
    fn new() -> Self {
        Self {
            timestamp: Instant::now(),
            energy: None,
            frequency: None,
            counter_values: Vec::new(),
        }
    }

    // Take a measurement with the given counter group
    fn take(counters: &[(String, Counter)], group: &mut Group) -> Result<Self> {
        let mut measurement = Self::new();

        // Read energy
        // Basline is 4.5W
        measurement.energy = rapl::read_package_energy().ok();

        // Read CPU frequency
        measurement.frequency = read_cpu_frequency(0);

        // Read performance counters
        let counts = group.read()?;
        // dbg!(&counts);

        let counters: HashMap<_, _> = counters.iter().map(|(a, b)| (a.clone(), b)).collect();

        // Extract counter values
        for (name, _) in define_available_events() {
            let Some(counter) = counters.get(&name) else {
                measurement.counter_values.push(None);
                continue;
            };
            measurement
                .counter_values
                .push(counts.get(counter).cloned());
        }

        Ok(measurement)
    }

    // Calculate the difference between two measurements
    fn diff(&self, previous: &Measurement) -> MeasurementDiff {
        let duration_ms = self
            .timestamp
            .duration_since(previous.timestamp)
            .as_millis() as u64;

        // Calculate energy delta
        let energy_delta = match (previous.energy, self.energy) {
            (Some(prev), Some(curr)) => curr - prev,
            _ => 0.0,
        };

        MeasurementDiff {
            timestamp: self.timestamp,
            duration_ms,
            energy_delta,
            frequency: self.frequency,
            counter_deltas: self.counter_values.clone(),
        }
    }
}

// Represents the difference between two measurements
struct MeasurementDiff {
    timestamp: Instant,
    duration_ms: u64,
    energy_delta: f64,
    frequency: Option<f64>,
    counter_deltas: Vec<Option<u64>>,
}

impl MeasurementDiff {
    // Write this diff as a CSV record
    fn write_csv_record(&self, writer: &mut Writer<File>) -> Result<()> {
        // Prepare CSV record
        let mut record = vec![
            self.timestamp.elapsed().as_secs_f64().to_string(),
            self.duration_ms.to_string(),
            self.energy_delta.to_string(),
            self.frequency
                .map(|f| f.to_string())
                .unwrap_or_else(|| "N/A".to_string()),
        ];
        record.extend(
            self.counter_deltas
                .iter()
                .map(|x| x.map(|x| x.to_string()).unwrap_or_default()),
        );

        // Write record
        writer.write_record(&record)?;
        // writer.flush()?;

        Ok(())
    }
}

impl<'a> BenchmarkScheduler<'a> {
    pub fn init(open_object: &'a mut MaybeUninit<OpenObject>, log_path: &str) -> Result<Self> {
        let bpf = BpfScheduler::init(
            open_object,
            0,     // exit_dump_len (default)
            false, // partial (include all tasks)
            false, // debug mode off
        )?;

        println!("Initializing benchmark scheduler (single-core profiling mode)");

        Ok(Self {
            bpf,
            own_pid: process::id() as i32,
            log_path: log_path.to_string(),
        })
    }

    fn consume_all_tasks(&mut self) -> Result<()> {
        while let Ok(Some(task)) = self.bpf.dequeue_task() {
            let mut dispatched_task = DispatchedTask::new(&task);

            // If it's our own process, schedule it to core 1
            if task.pid == self.own_pid {
                dispatched_task.cpu = 1;
                // dispatched_task.flags |= RL_CPU_ANY as u64;
            } else {
                // Schedule all other tasks on core 0
                // dispatched_task.flags |= RL_CPU_ANY as u64;
                dispatched_task.cpu = 0;
            }

            dispatched_task.slice_ns = SLICE_US;

            // Dispatch the task
            if let Err(e) = self.bpf.dispatch_task(&dispatched_task) {
                eprintln!("Failed to dispatch task: {}", e);
            }
        }

        // Notify BPF we're done processing tasks
        self.bpf.notify_complete(0);

        Ok(())
    }

    fn start_measurement_thread(&self) -> thread::JoinHandle<()> {
        let log_path = self.log_path.clone();

        thread::spawn(move || {
            if let Err(e) = run_measurement_loop(log_path) {
                eprintln!("Measurement thread error: {:?}", e);
            }
        })
    }

    pub fn run(&mut self) -> Result<UserExitInfo> {
        // Start the measurement thread
        self.start_measurement_thread();

        // Main scheduling loop
        while !self.bpf.exited() {
            // Process all tasks
            self.consume_all_tasks()?;
        }

        self.bpf.shutdown_and_report()
    }
}

// Main measurement loop
fn run_measurement_loop(log_path: String) -> Result<()> {
    // Define available hardware counters
    let available_events = define_available_events();

    // Initialize CSV writer with header
    let mut csv_writer = initialize_csv_writer(&log_path, &available_events)?;

    let mut rng = rand::rng();
    let mut round_counter = 0;

    // Main measurement loop
    loop {
        // println!("Starting new counter group (round {})", round_counter);
        round_counter += 1;

        let cpu = 0;
        // Create a new perf group
        let mut group = match Group::new_with_pid_and_cpu(-1, cpu) {
            Ok(g) => g,
            Err(e) => {
                eprintln!("Failed to create perf group: {}", e);
                thread::sleep(Duration::from_millis(100));
                continue;
            }
        };

        // Select random subset of counters
        let selected_events = available_events
            .iter()
            .choose_multiple(&mut rng, MAX_COUNTERS_AT_ONCE);

        // println!("Selected {} events for monitoring", selected_events.len());

        // Build counters
        let mut counters = Vec::new();
        for (name, event) in &selected_events {
            match Builder::new()
                .group(&mut group)
                .kind(event.clone())
                .observe_pid(-1)
                .one_cpu(cpu.try_into().unwrap()) // Core 0 is where we're scheduling tasks
                .build()
            {
                Ok(counter) => {
                    // println!("Successfully created counter for {}", name);
                    counters.push((name.clone(), counter));
                }
                Err(e) => {
                    eprintln!("Failed to create counter for {}: {}", name, e);
                }
            }
        }
        group.enable().unwrap();

        if counters.is_empty() {
            eprintln!("Failed to create any counters, retrying...");
            thread::sleep(Duration::from_millis(100));
            continue;
        }

        // Enable the counter group
        if let Err(e) = group.enable() {
            eprintln!("Failed to enable perf group: {}", e);
            thread::sleep(Duration::from_millis(100));
            continue;
        }

        // println!(
        //     "Successfully enabled counter group with {} counters",
        //     counters.len()
        // );

        // Take initial measurement
        let mut prev_measurement = match Measurement::take(&counters, &mut group) {
            Ok(m) => m,
            Err(e) => {
                eprintln!("Failed to take initial measurement: {}", e);
                thread::sleep(Duration::from_millis(100));
                continue;
            }
        };

        // println!("Took initial measurement");

        // Monitor for several rounds before reshuffling
        for round in 0..RESHUFFLE_ROUNDS {
            group.enable().unwrap();
            group.reset().unwrap();
            // Wait for the sampling interval
            thread::sleep(Duration::from_millis(LOG_INTERVAL_MS));

            // Take current measurement
            let curr_measurement = match Measurement::take(&counters, &mut group) {
                Ok(m) => m,
                Err(e) => {
                    eprintln!("Failed to take measurement in round {}: {}", round, e);
                    continue;
                }
            };

            // Calculate difference and write to CSV
            let diff = curr_measurement.diff(&prev_measurement);
            // println!(
            //     "Measurement diff: duration={}ms, energy={}J",
            //     diff.duration_ms, diff.energy_delta
            // );

            if let Err(e) = diff.write_csv_record(&mut csv_writer) {
                eprintln!("Failed to write CSV record: {}", e);
            }

            // Current becomes previous for next iteration
            prev_measurement = curr_measurement;
        }

        let _ = group.disable();
        // panic!();
    }
}

fn initialize_csv_writer(
    log_path: &str,
    available_events: &[(String, Event)],
) -> Result<Writer<File>> {
    let file = File::create(log_path)?;
    let mut csv_writer = Writer::from_writer(file);

    // Write header with all possible counter names
    let mut header = vec![
        "timestamp".to_string(),
        "duration_ms".to_string(),
        "package_power_j".to_string(),
        "cpu_frequency_mhz".to_string(),
    ];

    // Add counter deltas
    for (name, _) in available_events {
        header.push(name.into());
    }

    csv_writer.write_record(&header)?;
    csv_writer.flush()?;

    Ok(csv_writer)
}

pub fn read_cpu_frequency(cpu_id: u32) -> Option<f64> {
    // Try to read frequency from sysfs
    let freq_path = format!(
        "/sys/devices/system/cpu/cpu{}/cpufreq/scaling_cur_freq",
        cpu_id
    );

    match std::fs::read_to_string(freq_path) {
        Ok(content) => {
            // Convert from kHz to MHz
            content.trim().parse::<f64>().ok().map(|freq| freq / 1000.0)
        }
        Err(_) => None,
    }
}

fn define_available_events() -> Vec<(String, Event)> {
    let mut events = Vec::new();

    // Hardware events
    events.extend([
        (
            "cpu_cycles".to_string(),
            Event::Hardware(Hardware::CPU_CYCLES),
        ),
        (
            "instructions".to_string(),
            Event::Hardware(Hardware::INSTRUCTIONS),
        ),
        (
            "cache_references".to_string(),
            Event::Hardware(Hardware::CACHE_REFERENCES),
        ),
        (
            "cache_misses".to_string(),
            Event::Hardware(Hardware::CACHE_MISSES),
        ),
        (
            "branch_instructions".to_string(),
            Event::Hardware(Hardware::BRANCH_INSTRUCTIONS),
        ),
        (
            "branch_misses".to_string(),
            Event::Hardware(Hardware::BRANCH_MISSES),
        ),
        (
            "ref_cpu_cycles".to_string(),
            Event::Hardware(Hardware::REF_CPU_CYCLES),
        ),
        (
            "task_clock".to_string(),
            Event::Software(Software::TASK_CLOCK),
        ),
        // (
        //     "stalled-cycles-frontend".to_string(),
        //     Event::Hardware(Hardware::STALLED_CYCLES_FRONTEND),
        // ),
    ]);

    // L1 Data Cache events
    events.extend([
        // (
        //     "l1d_read_access".to_string(),
        //     Event::Cache(Cache {
        //         which: WhichCache::L1D,
        //         operation: CacheOp::READ,
        //         result: CacheResult::ACCESS,
        //     }),
        // ),
        // (
        //     "l1d_read_miss".to_string(),
        //     Event::Cache(Cache {
        //         which: WhichCache::L1D,
        //         operation: CacheOp::READ,
        //         result: CacheResult::MISS,
        //     }),
        // ),
    ]);

    // events.push((
    //     "fp_ops_retired_by_type.all".to_string(),
    //     Event::Raw(4, 0xFF0A), // EventCode 0x0A with UMask 0xFF
    // ));
    // events.push((
    //     "sse_avx_ops_retired.all".to_string(),
    //     Event::Raw(4, 0xFF0B), // EventCode 0x0B with UMask 0xFF
    // ));

    // L1 Instruction Cache events
    events.extend([
        // (
        //     "l1i_read_access".to_string(),
        //     Event::Cache(Cache {
        //         which: WhichCache::L1I,
        //         operation: CacheOp::READ,
        //         result: CacheResult::ACCESS,
        //     }),
        // ),
        // (
        //     "l1i_read_miss".to_string(),
        //     Event::Cache(Cache {
        //         which: WhichCache::L1I,
        //         operation: CacheOp::READ,
        //         result: CacheResult::MISS,
        //     }),
        // ),
    ]);

    // Data TLB events
    events.extend([
        // (
        //     "dtlb_read_access".to_string(),
        //     Event::Cache(Cache {
        //         which: WhichCache::DTLB,
        //         operation: CacheOp::READ,
        //         result: CacheResult::ACCESS,
        //     }),
        // ),
        // (
        //     "dtlb_read_miss".to_string(),
        //     Event::Cache(Cache {
        //         which: WhichCache::DTLB,
        //         operation: CacheOp::READ,
        //         result: CacheResult::MISS,
        //     }),
        // ),
    ]);

    // Instruction TLB events
    events.extend([
        // (
        //     "itlb_read_access".to_string(),
        //     Event::Cache(Cache {
        //         which: WhichCache::ITLB,
        //         operation: CacheOp::READ,
        //         result: CacheResult::ACCESS,
        //     }),
        // ),
        // (
        //     "itlb_read_miss".to_string(),
        //     Event::Cache(Cache {
        //         which: WhichCache::ITLB,
        //         operation: CacheOp::READ,
        //         result: CacheResult::MISS,
        //     }),
        // ),
    ]);

    // Branch Prediction Unit events
    events.extend([
        // (
        //     "bpu_read_access".to_string(),
        //     Event::Cache(Cache {
        //         which: WhichCache::BPU,
        //         operation: CacheOp::READ,
        //         result: CacheResult::ACCESS,
        //     }),
        // ),
        // (
        //     "bpu_read_miss".to_string(),
        //     Event::Cache(Cache {
        //         which: WhichCache::BPU,
        //         operation: CacheOp::READ,
        //         result: CacheResult::MISS,
        //     }),
        // ),
    ]);

    // Sort events by name for consistent ordering
    events.sort_unstable_by_key(|(name, _)| name.clone());

    events
}