use std::collections::{HashMap, VecDeque};

use crate::parallel::types::RepoIdentity;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct CaInstanceId(pub u64);

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct PublicationPointId(pub u64);

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PublicationPointState {
    pub ca_instance_id: CaInstanceId,
    pub pending_roa_tasks: usize,
    pub child_discovery_released: bool,
    pub finalized: bool,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Phase2CompletionSnapshot {
    pub ca_ready_queue_empty: bool,
    pub ca_waiting_repo_empty: bool,
    pub repo_tasks_idle: bool,
    pub pending_roa_dispatch_empty: bool,
    pub worker_queues_empty: bool,
    pub object_result_queue_empty: bool,
    pub object_workers_idle: bool,
    pub inflight_publication_points_empty: bool,
}

impl Phase2CompletionSnapshot {
    pub fn is_complete(&self) -> bool {
        self.ca_ready_queue_empty
            && self.ca_waiting_repo_empty
            && self.repo_tasks_idle
            && self.pending_roa_dispatch_empty
            && self.worker_queues_empty
            && self.object_result_queue_empty
            && self.object_workers_idle
            && self.inflight_publication_points_empty
    }
}

#[derive(Default)]
pub struct Phase2SchedulerState<T> {
    ca_waiting_repo_by_identity: HashMap<RepoIdentity, Vec<CaInstanceId>>,
    ca_ready_queue: VecDeque<CaInstanceId>,
    inflight_publication_points: HashMap<PublicationPointId, PublicationPointState>,
    pending_roa_dispatch: VecDeque<T>,
}

impl<T> Phase2SchedulerState<T> {
    pub fn new() -> Self {
        Self {
            ca_waiting_repo_by_identity: HashMap::new(),
            ca_ready_queue: VecDeque::new(),
            inflight_publication_points: HashMap::new(),
            pending_roa_dispatch: VecDeque::new(),
        }
    }

    pub fn wait_for_repo(&mut self, identity: RepoIdentity, ca_id: CaInstanceId) {
        self.ca_waiting_repo_by_identity
            .entry(identity)
            .or_default()
            .push(ca_id);
    }

    pub fn release_repo_waiters(&mut self, identity: &RepoIdentity) -> Vec<CaInstanceId> {
        let released = self
            .ca_waiting_repo_by_identity
            .remove(identity)
            .unwrap_or_default();
        for ca_id in &released {
            self.ca_ready_queue.push_back(*ca_id);
        }
        released
    }

    pub fn push_ready_ca(&mut self, ca_id: CaInstanceId) {
        self.ca_ready_queue.push_back(ca_id);
    }

    pub fn pop_ready_ca(&mut self) -> Option<CaInstanceId> {
        self.ca_ready_queue.pop_front()
    }

    pub fn start_publication_point(
        &mut self,
        pp_id: PublicationPointId,
        ca_id: CaInstanceId,
        pending_roa_tasks: usize,
    ) {
        self.inflight_publication_points.insert(
            pp_id,
            PublicationPointState {
                ca_instance_id: ca_id,
                pending_roa_tasks,
                child_discovery_released: false,
                finalized: false,
            },
        );
    }

    pub fn mark_child_discovery_released(&mut self, pp_id: PublicationPointId) {
        if let Some(state) = self.inflight_publication_points.get_mut(&pp_id) {
            state.child_discovery_released = true;
        }
    }

    pub fn enqueue_roa_task(&mut self, task: T) {
        self.pending_roa_dispatch.push_back(task);
    }

    pub fn pop_pending_roa_dispatch(&mut self) -> Option<T> {
        self.pending_roa_dispatch.pop_front()
    }

    pub fn push_front_pending_roa_dispatch(&mut self, task: T) {
        self.pending_roa_dispatch.push_front(task);
    }

    pub fn record_roa_result(&mut self, pp_id: PublicationPointId) -> Option<PublicationPointId> {
        let state = self.inflight_publication_points.get_mut(&pp_id)?;
        state.pending_roa_tasks = state.pending_roa_tasks.saturating_sub(1);
        if state.pending_roa_tasks == 0 {
            state.finalized = true;
            self.inflight_publication_points.remove(&pp_id);
            Some(pp_id)
        } else {
            None
        }
    }

    pub fn waiting_repo_len(&self) -> usize {
        self.ca_waiting_repo_by_identity
            .values()
            .map(Vec::len)
            .sum()
    }

    pub fn ready_queue_len(&self) -> usize {
        self.ca_ready_queue.len()
    }

    pub fn inflight_len(&self) -> usize {
        self.inflight_publication_points.len()
    }

    pub fn pending_roa_dispatch_len(&self) -> usize {
        self.pending_roa_dispatch.len()
    }

    pub fn publication_point_state(
        &self,
        pp_id: PublicationPointId,
    ) -> Option<&PublicationPointState> {
        self.inflight_publication_points.get(&pp_id)
    }

    pub fn completion_snapshot(
        &self,
        repo_tasks_idle: bool,
        worker_queues_empty: bool,
        object_result_queue_empty: bool,
        object_workers_idle: bool,
    ) -> Phase2CompletionSnapshot {
        Phase2CompletionSnapshot {
            ca_ready_queue_empty: self.ca_ready_queue.is_empty(),
            ca_waiting_repo_empty: self.ca_waiting_repo_by_identity.is_empty(),
            repo_tasks_idle,
            pending_roa_dispatch_empty: self.pending_roa_dispatch.is_empty(),
            worker_queues_empty,
            object_result_queue_empty,
            object_workers_idle,
            inflight_publication_points_empty: self.inflight_publication_points.is_empty(),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{CaInstanceId, Phase2SchedulerState, PublicationPointId};
    use crate::parallel::object_worker::{ObjectTaskExecutor, ObjectWorkerPool};
    use crate::parallel::types::RepoIdentity;
    use std::time::Duration;

    fn identity(name: &str) -> RepoIdentity {
        RepoIdentity::new(
            Some(format!("https://example.test/{name}/notification.xml")),
            format!("rsync://example.test/{name}/"),
        )
    }

    #[test]
    fn scheduler_repo_ready_moves_waiting_ca_to_ready_queue() {
        let mut state = Phase2SchedulerState::<u64>::new();
        let repo = identity("arin");
        state.wait_for_repo(repo.clone(), CaInstanceId(1));
        state.wait_for_repo(repo.clone(), CaInstanceId(2));
        assert_eq!(state.waiting_repo_len(), 2);

        let released = state.release_repo_waiters(&repo);
        assert_eq!(released, vec![CaInstanceId(1), CaInstanceId(2)]);
        assert_eq!(state.waiting_repo_len(), 0);
        assert_eq!(state.ready_queue_len(), 2);
        assert_eq!(state.pop_ready_ca(), Some(CaInstanceId(1)));
        assert_eq!(state.pop_ready_ca(), Some(CaInstanceId(2)));
    }

    #[test]
    fn scheduler_releases_child_before_roa_results_finalize_parent() {
        let mut state = Phase2SchedulerState::<u64>::new();
        let pp = PublicationPointId(10);
        state.start_publication_point(pp, CaInstanceId(1), 2);
        state.mark_child_discovery_released(pp);
        state.push_ready_ca(CaInstanceId(2));

        let pp_state = state.publication_point_state(pp).expect("inflight pp");
        assert!(pp_state.child_discovery_released);
        assert_eq!(pp_state.pending_roa_tasks, 2);
        assert_eq!(state.pop_ready_ca(), Some(CaInstanceId(2)));
        assert_eq!(state.record_roa_result(pp), None);
        assert_eq!(state.inflight_len(), 1);
        assert_eq!(state.record_roa_result(pp), Some(pp));
        assert_eq!(state.inflight_len(), 0);
    }

    #[test]
    fn scheduler_completion_requires_all_queues_and_inflight_to_be_empty() {
        let mut state = Phase2SchedulerState::new();
        assert!(
            state
                .completion_snapshot(true, true, true, true)
                .is_complete()
        );

        state.enqueue_roa_task(1u64);
        assert!(
            !state
                .completion_snapshot(true, true, true, true)
                .is_complete()
        );
        assert_eq!(state.pop_pending_roa_dispatch(), Some(1));
        assert!(
            state
                .completion_snapshot(true, true, true, true)
                .is_complete()
        );
    }

    #[test]
    fn scheduler_can_retry_pending_roa_task_at_front() {
        let mut state = Phase2SchedulerState::new();
        state.enqueue_roa_task(1u64);
        state.enqueue_roa_task(2u64);

        assert_eq!(state.pop_pending_roa_dispatch(), Some(1));
        state.push_front_pending_roa_dispatch(3);
        assert_eq!(state.pop_pending_roa_dispatch(), Some(3));
        assert_eq!(state.pop_pending_roa_dispatch(), Some(2));
        assert_eq!(state.pop_pending_roa_dispatch(), None);
    }

    #[derive(Clone, Debug, PartialEq, Eq)]
    struct TestRoaTask {
        pp_id: PublicationPointId,
        value: u64,
    }

    #[derive(Clone, Debug, PartialEq, Eq)]
    struct TestRoaResult {
        worker_index: usize,
        pp_id: PublicationPointId,
        value: u64,
    }

    #[derive(Clone)]
    struct TestRoaExecutor;

    impl ObjectTaskExecutor<TestRoaTask, TestRoaResult> for TestRoaExecutor {
        fn execute(&self, worker_index: usize, task: TestRoaTask) -> TestRoaResult {
            TestRoaResult {
                worker_index,
                pp_id: task.pp_id,
                value: task.value,
            }
        }
    }

    #[test]
    fn scheduler_dispatches_pending_roa_tasks_to_workers_and_finalizes_on_results() {
        let pp = PublicationPointId(42);
        let mut state = Phase2SchedulerState::new();
        state.start_publication_point(pp, CaInstanceId(7), 3);
        state.mark_child_discovery_released(pp);
        for value in 0..3 {
            state.enqueue_roa_task(TestRoaTask { pp_id: pp, value });
        }

        let mut pool = ObjectWorkerPool::new(2, 4, TestRoaExecutor).expect("object pool");
        while let Some(task) = state.pop_pending_roa_dispatch() {
            pool.try_submit_round_robin(task).expect("submit task");
        }
        assert_eq!(state.pending_roa_dispatch_len(), 0);

        let mut results = Vec::new();
        for _ in 0..3 {
            let result = pool
                .recv_result_timeout(Duration::from_secs(1))
                .expect("result channel")
                .expect("result");
            let finalized = state.record_roa_result(result.pp_id);
            results.push(result);
            if results.len() < 3 {
                assert_eq!(finalized, None);
            } else {
                assert_eq!(finalized, Some(pp));
            }
        }
        results.sort_by_key(|result| result.value);
        assert_eq!(results[0].worker_index, 0);
        assert_eq!(results[1].worker_index, 1);
        assert_eq!(results[2].worker_index, 0);
        assert_eq!(state.inflight_len(), 0);
        assert!(
            state
                .completion_snapshot(true, true, true, true)
                .is_complete()
        );
    }
}
