use crate::bpf::*;
use crate::energy::{self, Request as EnergyRequest};
use crate::freq::{self, FrequencyKHZ, Request as FrequencyRequest};

use crate::e_core_selector::{ECoreSelector, RoundRobinSelector};
use anyhow::Result;
use dashmap::DashMap;
use libbpf_rs::OpenObject;
use scx_utils::{Topology, UserExitInfo};

use std::collections::{HashMap, VecDeque};
use std::mem::MaybeUninit;
use std::ops::{Range, RangeInclusive};
use std::process;
use std::sync::{mpsc, Arc, RwLock};
use std::time::{Duration, Instant};

use crate::Pid;

const SLICE_US: u64 = 50000;

pub struct Scheduler<'a> {
    bpf: BpfScheduler<'a>,
    task_queue: VecDeque<QueuedTask>,
    no_budget_task_queue: VecDeque<QueuedTask>,
    managed_tasks: HashMap<Pid, Instant>,
    maximum_budget: u64,
    own_pid: Pid,
    p_cores: Range<i32>,
    e_cores: Range<i32>,
    topology: Topology,
    to_remove: Vec<Pid>,
    e_core_selector: Box<dyn ECoreSelector>,
    energy_sender: mpsc::SyncSender<EnergyRequest>,
    shared_budgets: Arc<DashMap<Pid, u64>>,
    frequency_sender: mpsc::SyncSender<FrequencyRequest>,
    shared_cpu_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
    shared_policy_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>>,
    shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
}

impl<'a> Scheduler<'a> {
    pub fn init(
        open_object: &'a mut MaybeUninit<OpenObject>,
        use_mocking: bool,
        power_cap: u64,
    ) -> Result<Self> {
        let bpf = BpfScheduler::init(
            open_object,
            0,     // exit_dump_len (buffer size of exit info, 0 = default)
            false, // partial (false = include all tasks)
            false, // debug (false = debug mode off)
        )?;

        println!("Initializing energy-aware scheduler");

        let shared_cpu_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>> =
            Arc::new(RwLock::new(Vec::new()));
        let shared_policy_frequency_ranges: Arc<RwLock<Vec<RangeInclusive<FrequencyKHZ>>>> =
            Arc::new(RwLock::new(Vec::new()));
        let shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>> =
            Arc::new(RwLock::new(Vec::new()));

        // Shared budget map between energy service and scheduler
        let shared_budgets: Arc<DashMap<Pid, u64>> = Arc::new(DashMap::with_capacity(100000));

        // Start energy tracking service
        let energy_sender = energy::start_energy_service(
            use_mocking,
            power_cap,
            shared_budgets.clone(),
            shared_cpu_frequency_ranges.clone(),
            shared_policy_frequency_ranges.clone(),
            shared_cpu_current_frequencies.clone(),
        )?;

        let topology = Topology::new().unwrap();
        let mut e_core_ids = Vec::new();
        let mut p_core_ids = Vec::new();
        for (id, cpu) in &topology.all_cpus {
            match cpu.core_type {
                scx_utils::CoreType::Big { turbo: _ } => p_core_ids.push(*id as i32),
                scx_utils::CoreType::Little => e_core_ids.push(*id as i32),
            }
        }
        // We assume that the CPU IDs for each core type are assigned contiguously.
        e_core_ids.sort();
        p_core_ids.sort();
        let e_cores = *e_core_ids.first().unwrap_or(&0)..*e_core_ids.last().unwrap_or(&0);
        let p_cores = *p_core_ids.first().unwrap_or(&0)..*p_core_ids.last().unwrap_or(&0);
        let all_cores = 0..((e_cores.len() + p_cores.len()) as u32);

        let selector = Box::new(RoundRobinSelector::new(&e_cores));
        let to_remove = Vec::with_capacity(1000);

        let frequency_sender = freq::start_frequency_service(
            all_cores,
            shared_cpu_frequency_ranges.clone(),
            shared_policy_frequency_ranges.clone(),
            shared_cpu_current_frequencies.clone(),
            Duration::from_millis(1),
        )?;

        frequency_sender
            .try_send(FrequencyRequest::GetPolicyCPUFrequency)
            .unwrap();
        frequency_sender
            .try_send(FrequencyRequest::GetPossibleCPUFrequencyRange)
            .unwrap();
        frequency_sender
            .try_send(FrequencyRequest::GetCurrentFrequencies)
            .unwrap();

        Ok(Self {
            bpf,
            task_queue: VecDeque::new(),
            no_budget_task_queue: VecDeque::new(),
            managed_tasks: HashMap::new(),
            maximum_budget: u64::MAX,
            own_pid: process::id() as i32,
            p_cores,
            e_cores,
            topology,
            e_core_selector: selector,
            energy_sender,
            shared_budgets,
            to_remove,
            frequency_sender,
            shared_cpu_frequency_ranges,
            shared_policy_frequency_ranges,
            shared_cpu_current_frequencies,
        })
    }

    fn consume_all_tasks(&mut self) {
        while let Ok(Some(task)) = self.bpf.dequeue_task() {
            // The scheduler itself has to be scheduled regardless of its energy usage
            if task.pid == self.own_pid {
                self.task_queue.push_front(task);
                self.dispatch_next_task();
                continue;
            }

            // Check if we've seen this task before
            if let std::collections::hash_map::Entry::Vacant(e) = self.managed_tasks.entry(task.pid)
            {
                e.insert(Instant::now());
                // New task - register it with the energy service
                self.energy_sender
                    .try_send(EnergyRequest::NewTask(task.pid))
                    .unwrap();
            }

            // Get current budget for this task
            match self.shared_budgets.try_get(&task.pid) {
                dashmap::try_result::TryResult::Present(budget) => match *budget {
                    0 => self.task_queue.push_back(task),
                    _ => self.no_budget_task_queue.push_back(task),
                },
                _ => self.task_queue.push_back(task),
            }
        }
    }

    fn dispatch_next_task(&mut self) {
        if let Some(task) = self.task_queue.pop_front() {
            let mut dispatched_task = DispatchedTask::new(&task);

            self.managed_tasks.insert(task.pid, Instant::now());

            let cpu = self.bpf.select_cpu(task.pid, task.cpu, 0);
            if cpu >= 0 {
                dispatched_task.cpu = cpu;
            } else {
                dispatched_task.flags |= RL_CPU_ANY as u64;
            }

            dispatched_task.slice_ns = SLICE_US;

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

            self.bpf.notify_complete(
                self.task_queue.len() as u64 + self.no_budget_task_queue.len() as u64,
            );
        } else if let Some(task) = self.no_budget_task_queue.pop_front() {
            let mut dispatched_task = DispatchedTask::new(&task);

            // Low budget tasks go to e-cores
            let cpu = self.e_core_selector.next_core(task.cpu);
            if cpu >= 0 {
                dispatched_task.cpu = cpu;
            } else {
                dispatched_task.flags |= RL_CPU_ANY as u64;
            }

            // Scheduler tasks get longer slices
            if task.pid == self.own_pid {
                dispatched_task.slice_ns = SLICE_US * 1000;
            } else {
                dispatched_task.slice_ns = SLICE_US;
            }

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

            self.bpf.notify_complete(
                self.task_queue.len() as u64 + self.no_budget_task_queue.len() as u64,
            );
        }
    }

    fn cleanup_old_tasks(&mut self) {
        let current = Instant::now();
        for (pid, last_scheduled) in &self.managed_tasks {
            if current - *last_scheduled > Duration::from_secs(5) {
                self.to_remove.push(*pid);
            }
        }
        for pid in self.to_remove.drain(..) {
            self.managed_tasks.remove(&pid);
            self.energy_sender
                .try_send(EnergyRequest::RemoveTask(pid))
                .unwrap();
        }
    }

    fn dispatch_tasks(&mut self) {
        loop {
            self.consume_all_tasks();
            self.dispatch_next_task();

            if self.task_queue.is_empty() && self.no_budget_task_queue.is_empty() {
                self.bpf.notify_complete(0);
                break;
            }
        }
    }

    pub fn run(&mut self) -> Result<UserExitInfo> {
        let mut i = 0;
        while !self.bpf.exited() {
            i += 1;
            self.dispatch_tasks();
            if i % 100 == 0 {
                self.cleanup_old_tasks();
            }
        }

        // Clean up - signal the energy service to stop tracking all managed tasks
        for pid in self.managed_tasks.keys() {
            let _ = self.energy_sender.send(EnergyRequest::RemoveTask(*pid));
        }

        self.bpf.shutdown_and_report()
    }
}