use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};

use crate::parallel::types::{
    RepoDedupKey, RepoIdentity, RepoRequester, RepoTransportMode, RepoTransportTask,
};
use crate::policy::SyncPreference;

pub const TRANSPORT_PREFETCH_SCHEMA_VERSION: u32 = 1;

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct TransportPrefetchSnapshot {
    pub schema_version: u32,
    pub generated_at_unix: i64,
    pub sync_preference: SyncPreference,
    pub requests: Vec<TransportPrefetchRequest>,
}

impl TransportPrefetchSnapshot {
    pub fn new(sync_preference: SyncPreference, requests: Vec<TransportPrefetchRequest>) -> Self {
        Self {
            schema_version: TRANSPORT_PREFETCH_SCHEMA_VERSION,
            generated_at_unix: time::OffsetDateTime::now_utc().unix_timestamp(),
            sync_preference,
            requests,
        }
    }

    pub fn is_compatible_with(&self, sync_preference: SyncPreference) -> bool {
        self.schema_version == TRANSPORT_PREFETCH_SCHEMA_VERSION
            && self.sync_preference == sync_preference
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TransportPrefetchRequest {
    pub dedup_key: TransportPrefetchDedupKey,
    pub rsync_scope_uri: String,
    pub rsync_failure_scope_uri: Option<String>,
    pub repo_identity: TransportPrefetchRepoIdentity,
    pub mode: TransportPrefetchMode,
    #[serde(default)]
    pub last_result: Option<TransportPrefetchLastResult>,
    #[serde(default)]
    pub last_rsync_result: Option<TransportPrefetchLastResult>,
    pub tal_id: String,
    pub rir_id: String,
    pub priority: u8,
    pub requesters: Vec<TransportPrefetchRequester>,
}

impl TransportPrefetchRequest {
    pub fn from_registered_request(
        identity: &RepoIdentity,
        requester: &RepoRequester,
        priority: u8,
        rsync_scope_uri: String,
        rsync_failure_scope_uri: Option<String>,
        sync_preference: SyncPreference,
    ) -> Self {
        let (dedup_key, mode) = if sync_preference == SyncPreference::RrdpThenRsync {
            if let Some(notification_uri) = identity.notification_uri.clone() {
                (
                    TransportPrefetchDedupKey::RrdpNotify { notification_uri },
                    TransportPrefetchMode::Rrdp,
                )
            } else {
                (
                    TransportPrefetchDedupKey::RsyncScope {
                        rsync_scope_uri: rsync_scope_uri.clone(),
                    },
                    TransportPrefetchMode::Rsync,
                )
            }
        } else {
            (
                TransportPrefetchDedupKey::RsyncScope {
                    rsync_scope_uri: rsync_scope_uri.clone(),
                },
                TransportPrefetchMode::Rsync,
            )
        };

        Self {
            dedup_key,
            rsync_scope_uri,
            rsync_failure_scope_uri,
            repo_identity: TransportPrefetchRepoIdentity::from_identity(identity),
            mode,
            last_result: None,
            last_rsync_result: None,
            tal_id: requester.tal_id.clone(),
            rir_id: requester.rir_id.clone(),
            priority,
            requesters: vec![TransportPrefetchRequester::from_requester(requester)],
        }
    }

    pub fn from_task(task: &RepoTransportTask, rsync_scope_uri: String) -> Self {
        Self {
            dedup_key: TransportPrefetchDedupKey::from_repo_key(&task.dedup_key),
            rsync_scope_uri,
            rsync_failure_scope_uri: task.rsync_failure_scope_uri.clone(),
            repo_identity: TransportPrefetchRepoIdentity::from_identity(&task.repo_identity),
            mode: TransportPrefetchMode::from_mode(task.mode),
            last_result: None,
            last_rsync_result: None,
            tal_id: task.tal_id.clone(),
            rir_id: task.rir_id.clone(),
            priority: task.priority,
            requesters: task
                .requesters
                .iter()
                .map(TransportPrefetchRequester::from_requester)
                .collect(),
        }
    }

    pub fn to_identity(&self) -> RepoIdentity {
        self.repo_identity.to_identity()
    }

    pub fn to_requester(&self) -> RepoRequester {
        self.requesters
            .first()
            .map(TransportPrefetchRequester::to_requester)
            .unwrap_or_else(|| RepoRequester {
                tal_id: self.tal_id.clone(),
                rir_id: self.rir_id.clone(),
                parent_node_id: None,
                ca_instance_handle_id: format!(
                    "{}:{}",
                    self.tal_id, self.repo_identity.rsync_base_uri
                ),
                publication_point_rsync_uri: self.repo_identity.rsync_base_uri.clone(),
                manifest_rsync_uri: format!("{}prefetch.mft", self.repo_identity.rsync_base_uri),
            })
    }

    fn recorder_key(&self) -> String {
        self.dedup_key.stable_key()
    }

    pub fn retry_short_timeout(&self) -> bool {
        matches!(
            self.last_result,
            Some(TransportPrefetchLastResult { ok: false })
        )
    }

    pub fn retry_short_rsync_timeout(&self) -> bool {
        matches!(
            self.last_rsync_result,
            Some(TransportPrefetchLastResult { ok: false })
        ) || (self.mode == TransportPrefetchMode::Rsync && self.retry_short_timeout())
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TransportPrefetchDedupKey {
    RrdpNotify { notification_uri: String },
    RsyncScope { rsync_scope_uri: String },
}

impl TransportPrefetchDedupKey {
    fn from_repo_key(key: &RepoDedupKey) -> Self {
        match key {
            RepoDedupKey::RrdpNotify { notification_uri } => Self::RrdpNotify {
                notification_uri: notification_uri.clone(),
            },
            RepoDedupKey::RsyncScope { rsync_scope_uri } => Self::RsyncScope {
                rsync_scope_uri: rsync_scope_uri.clone(),
            },
        }
    }

    fn stable_key(&self) -> String {
        match self {
            Self::RrdpNotify { notification_uri } => format!("rrdp:{notification_uri}"),
            Self::RsyncScope { rsync_scope_uri } => format!("rsync:{rsync_scope_uri}"),
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TransportPrefetchRepoIdentity {
    pub notification_uri: Option<String>,
    pub rsync_base_uri: String,
}

impl TransportPrefetchRepoIdentity {
    fn from_identity(identity: &RepoIdentity) -> Self {
        Self {
            notification_uri: identity.notification_uri.clone(),
            rsync_base_uri: identity.rsync_base_uri.clone(),
        }
    }

    fn to_identity(&self) -> RepoIdentity {
        RepoIdentity::new(self.notification_uri.clone(), self.rsync_base_uri.clone())
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TransportPrefetchMode {
    Rrdp,
    Rsync,
}

impl TransportPrefetchMode {
    fn from_mode(mode: RepoTransportMode) -> Self {
        match mode {
            RepoTransportMode::Rrdp => Self::Rrdp,
            RepoTransportMode::Rsync => Self::Rsync,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TransportPrefetchLastResult {
    pub ok: bool,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TransportPrefetchRequester {
    pub tal_id: String,
    pub rir_id: String,
    pub parent_node_id: Option<u64>,
    pub ca_instance_handle_id: String,
    pub publication_point_rsync_uri: String,
    pub manifest_rsync_uri: String,
}

impl TransportPrefetchRequester {
    fn from_requester(requester: &RepoRequester) -> Self {
        Self {
            tal_id: requester.tal_id.clone(),
            rir_id: requester.rir_id.clone(),
            parent_node_id: requester.parent_node_id,
            ca_instance_handle_id: requester.ca_instance_handle_id.clone(),
            publication_point_rsync_uri: requester.publication_point_rsync_uri.clone(),
            manifest_rsync_uri: requester.manifest_rsync_uri.clone(),
        }
    }

    fn to_requester(&self) -> RepoRequester {
        RepoRequester {
            tal_id: self.tal_id.clone(),
            rir_id: self.rir_id.clone(),
            parent_node_id: self.parent_node_id,
            ca_instance_handle_id: self.ca_instance_handle_id.clone(),
            publication_point_rsync_uri: self.publication_point_rsync_uri.clone(),
            manifest_rsync_uri: self.manifest_rsync_uri.clone(),
        }
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TransportPrefetchRecorder {
    requests_by_key: BTreeMap<String, TransportPrefetchRequest>,
    results_by_key: BTreeMap<String, TransportPrefetchLastResult>,
    rsync_results_by_scope: BTreeMap<String, TransportPrefetchLastResult>,
    request_order: Vec<String>,
}

impl TransportPrefetchRecorder {
    pub fn record_registered_request(
        &mut self,
        identity: &RepoIdentity,
        requester: &RepoRequester,
        priority: u8,
        rsync_scope_uri: String,
        rsync_failure_scope_uri: Option<String>,
        sync_preference: SyncPreference,
    ) {
        let request = TransportPrefetchRequest::from_registered_request(
            identity,
            requester,
            priority,
            rsync_scope_uri,
            rsync_failure_scope_uri,
            sync_preference,
        );
        self.record_request(request);
    }

    pub fn record_task(&mut self, task: &RepoTransportTask, rsync_scope_uri: String) {
        let request = TransportPrefetchRequest::from_task(task, rsync_scope_uri);
        self.record_request(request);
    }

    pub fn record_result(&mut self, result: &crate::parallel::types::RepoTransportResultEnvelope) {
        let key = TransportPrefetchDedupKey::from_repo_key(&result.dedup_key).stable_key();
        let last_result = TransportPrefetchLastResult {
            ok: matches!(
                result.result,
                crate::parallel::types::RepoTransportResultKind::Success { .. }
            ),
        };
        self.results_by_key.insert(key.clone(), last_result);
        if let Some(request) = self.requests_by_key.get_mut(&key) {
            request.last_result = Some(last_result);
        }
        if let crate::parallel::types::RepoDedupKey::RsyncScope { rsync_scope_uri } =
            &result.dedup_key
        {
            self.rsync_results_by_scope
                .insert(rsync_scope_uri.clone(), last_result);
            for request in self.requests_by_key.values_mut() {
                if request.rsync_scope_uri == *rsync_scope_uri {
                    request.last_rsync_result = Some(last_result);
                }
            }
        }
    }

    pub fn snapshot(&self, sync_preference: SyncPreference) -> TransportPrefetchSnapshot {
        TransportPrefetchSnapshot::new(
            sync_preference,
            self.request_order
                .iter()
                .filter_map(|key| self.requests_by_key.get(key))
                .cloned()
                .collect(),
        )
    }

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

    fn record_request(&mut self, request: TransportPrefetchRequest) {
        let key = request.recorder_key();
        if !self.requests_by_key.contains_key(&key) {
            self.request_order.push(key.clone());
            let mut request = request;
            if request.last_result.is_none() {
                request.last_result = self.results_by_key.get(&key).copied();
            }
            if request.last_rsync_result.is_none() {
                request.last_rsync_result = self
                    .rsync_results_by_scope
                    .get(&request.rsync_scope_uri)
                    .copied();
            }
            self.requests_by_key.insert(key, request);
        }
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TransportPrefetchDispatchStats {
    pub loaded_requests: u64,
    pub enqueued_tasks: u64,
    pub waiting_requests: u64,
    pub reused_results: u64,
    pub skipped_incompatible: u64,
}

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

    fn requester(uri: &str) -> RepoRequester {
        RepoRequester {
            tal_id: "apnic".to_string(),
            rir_id: "apnic".to_string(),
            parent_node_id: None,
            ca_instance_handle_id: format!("apnic:{uri}"),
            publication_point_rsync_uri: "rsync://example.test/repo/".to_string(),
            manifest_rsync_uri: uri.to_string(),
        }
    }

    fn task(notification_uri: &str, manifest_uri: &str) -> RepoTransportTask {
        RepoTransportTask {
            dedup_key: RepoDedupKey::RrdpNotify {
                notification_uri: notification_uri.to_string(),
            },
            rsync_failure_scope_uri: Some("rsync://example.test/".to_string()),
            repo_identity: RepoIdentity::new(
                Some(notification_uri.to_string()),
                "rsync://example.test/repo/",
            ),
            mode: RepoTransportMode::Rrdp,
            retry_short_timeout: false,
            tal_id: "apnic".to_string(),
            rir_id: "apnic".to_string(),
            validation_time: time::OffsetDateTime::UNIX_EPOCH,
            priority: 0,
            requesters: vec![requester(manifest_uri)],
        }
    }

    #[test]
    fn recorder_deduplicates_by_transport_key() {
        let mut recorder = TransportPrefetchRecorder::default();
        recorder.record_task(
            &task(
                "https://example.test/notification.xml",
                "rsync://example.test/repo/a.mft",
            ),
            "rsync://example.test/repo/".to_string(),
        );
        recorder.record_task(
            &task(
                "https://example.test/notification.xml",
                "rsync://example.test/repo/b.mft",
            ),
            "rsync://example.test/repo/".to_string(),
        );

        let snapshot = recorder.snapshot(SyncPreference::RrdpThenRsync);
        assert_eq!(snapshot.requests.len(), 1);
        assert_eq!(
            snapshot.requests[0].requesters[0].manifest_rsync_uri,
            "rsync://example.test/repo/a.mft"
        );
    }

    #[test]
    fn recorder_preserves_first_discovery_order() {
        let mut recorder = TransportPrefetchRecorder::default();
        recorder.record_task(
            &task(
                "https://z.example.test/notification.xml",
                "rsync://z.example.test/repo/root.mft",
            ),
            "rsync://z.example.test/repo/".to_string(),
        );
        recorder.record_task(
            &task(
                "https://a.example.test/notification.xml",
                "rsync://a.example.test/repo/root.mft",
            ),
            "rsync://a.example.test/repo/".to_string(),
        );

        let snapshot = recorder.snapshot(SyncPreference::RrdpThenRsync);
        assert_eq!(
            snapshot.requests[0]
                .repo_identity
                .notification_uri
                .as_deref(),
            Some("https://z.example.test/notification.xml")
        );
        assert_eq!(
            snapshot.requests[1]
                .repo_identity
                .notification_uri
                .as_deref(),
            Some("https://a.example.test/notification.xml")
        );
    }

    #[test]
    fn snapshot_checks_schema_and_sync_preference() {
        let snapshot = TransportPrefetchSnapshot::new(SyncPreference::RrdpThenRsync, Vec::new());
        assert!(snapshot.is_compatible_with(SyncPreference::RrdpThenRsync));
        assert!(!snapshot.is_compatible_with(SyncPreference::RsyncOnly));
    }

    #[test]
    fn recorder_marks_failed_rrdp_and_rsync_for_next_run_short_timeouts() {
        let mut recorder = TransportPrefetchRecorder::default();
        let rrdp_task = task(
            "https://example.test/notification.xml",
            "rsync://example.test/repo/a.mft",
        );
        recorder.record_task(&rrdp_task, "rsync://example.test/repo/".to_string());
        recorder.record_result(&crate::parallel::types::RepoTransportResultEnvelope {
            dedup_key: RepoDedupKey::RrdpNotify {
                notification_uri: "https://example.test/notification.xml".to_string(),
            },
            rsync_failure_scope_uri: None,
            repo_identity: rrdp_task.repo_identity.clone(),
            mode: RepoTransportMode::Rrdp,
            tal_id: "apnic".to_string(),
            rir_id: "apnic".to_string(),
            timing_ms: 1,
            result: crate::parallel::types::RepoTransportResultKind::Failed {
                detail: "timeout".to_string(),
                warnings: Vec::new(),
                error_class: crate::parallel::types::RepoTransportErrorClass::Unknown,
            },
        });
        recorder.record_result(&crate::parallel::types::RepoTransportResultEnvelope {
            dedup_key: RepoDedupKey::RsyncScope {
                rsync_scope_uri: "rsync://example.test/repo/".to_string(),
            },
            rsync_failure_scope_uri: None,
            repo_identity: rrdp_task.repo_identity.clone(),
            mode: RepoTransportMode::Rsync,
            tal_id: "apnic".to_string(),
            rir_id: "apnic".to_string(),
            timing_ms: 1,
            result: crate::parallel::types::RepoTransportResultKind::Failed {
                detail: "timeout".to_string(),
                warnings: Vec::new(),
                error_class: crate::parallel::types::RepoTransportErrorClass::Unknown,
            },
        });

        let snapshot = recorder.snapshot(SyncPreference::RrdpThenRsync);
        assert_eq!(snapshot.requests.len(), 1);
        assert!(snapshot.requests[0].retry_short_timeout());
        assert!(snapshot.requests[0].retry_short_rsync_timeout());
    }
}
