use std::collections::BTreeMap;

use clusterflux_core::{ProjectId, TenantId, UserId};
use serde::{Deserialize, Serialize};
use thiserror::Error;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ArtifactRelayConfiguration {
    pub enabled: bool,
    pub max_artifact_bytes: u64,
    pub max_active_per_project: usize,
    pub max_active_per_tenant: usize,
    pub max_active_per_account: usize,
    pub max_active_global: usize,
    pub max_period_bytes_per_project: u64,
    pub max_period_bytes_per_tenant: u64,
    pub max_period_bytes_per_account: u64,
    pub period_seconds: u64,
    pub framing_overhead_bytes: u64,
    pub max_tracked_scopes: usize,
}

impl ArtifactRelayConfiguration {
    pub fn unlimited() -> Self {
        Self {
            enabled: true,
            max_artifact_bytes: u64::MAX,
            max_active_per_project: usize::MAX,
            max_active_per_tenant: usize::MAX,
            max_active_per_account: usize::MAX,
            max_active_global: usize::MAX,
            max_period_bytes_per_project: u64::MAX,
            max_period_bytes_per_tenant: u64::MAX,
            max_period_bytes_per_account: u64::MAX,
            period_seconds: u64::MAX,
            framing_overhead_bytes: 512,
            max_tracked_scopes: usize::MAX,
        }
    }

    pub fn validate(&self) -> Result<(), ArtifactRelayError> {
        if self.period_seconds == 0
            || self.max_active_per_project == 0
            || self.max_active_per_tenant == 0
            || self.max_active_per_account == 0
            || self.max_active_global == 0
            || self.max_tracked_scopes == 0
        {
            return Err(ArtifactRelayError::InvalidConfiguration);
        }
        Ok(())
    }
}

impl Default for ArtifactRelayConfiguration {
    fn default() -> Self {
        Self::unlimited()
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ArtifactRelayUsage {
    pub active_transfers: usize,
    pub ingress_bytes: u64,
    pub egress_bytes: u64,
    pub abandoned_or_failed_bytes: u64,
    pub reserved_bytes: u64,
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct ArtifactRelayDurableState {
    pub reservations: BTreeMap<String, ArtifactRelayReservationState>,
    pub period: u64,
    pub project_used: Vec<(TenantId, ProjectId, u64)>,
    pub tenant_used: Vec<(TenantId, u64)>,
    pub account_used: Vec<(TenantId, UserId, u64)>,
    pub ingress_used: u64,
    pub egress_used: u64,
    pub abandoned_or_failed_used: u64,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ArtifactRelayReservationState {
    pub tenant: TenantId,
    pub project: ProjectId,
    pub account: UserId,
    pub remaining_reserved_bytes: u64,
    pub ingress_bytes: u64,
    pub egress_bytes: u64,
    pub expires_at_epoch_seconds: u64,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum RelayFinishReason {
    Completed,
    Failed,
    Cancelled,
    Expired,
}

#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum ArtifactRelayError {
    #[error("artifact relay is disabled by hosted policy")]
    Disabled,
    #[error("artifact exceeds the configured relay size limit")]
    ArtifactTooLarge,
    #[error("artifact relay concurrency limit reached for {0}")]
    Concurrency(&'static str),
    #[error("artifact relay is temporarily at process capacity")]
    Capacity,
    #[error("artifact relay period byte budget exhausted for {0}")]
    PeriodBudget(&'static str),
    #[error("artifact relay retained scope state limit reached")]
    StateLimit,
    #[error("artifact relay reservation is missing or expired")]
    MissingReservation,
    #[error("artifact relay configuration contains a zero bound")]
    InvalidConfiguration,
}

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
struct RelayScope {
    tenant: TenantId,
    project: ProjectId,
    account: UserId,
}

#[derive(Clone, Debug)]
struct RelayReservation {
    scope: RelayScope,
    remaining_reserved_bytes: u64,
    ingress_bytes: u64,
    egress_bytes: u64,
    expires_at_epoch_seconds: u64,
}

#[derive(Clone, Debug)]
pub(super) struct ArtifactRelayLedger {
    configuration: ArtifactRelayConfiguration,
    reservations: BTreeMap<String, RelayReservation>,
    period: u64,
    project_used: BTreeMap<(TenantId, ProjectId), u64>,
    tenant_used: BTreeMap<TenantId, u64>,
    account_used: BTreeMap<(TenantId, UserId), u64>,
    ingress_used: u64,
    egress_used: u64,
    abandoned_or_failed_used: u64,
}

impl Default for ArtifactRelayLedger {
    fn default() -> Self {
        Self::new(ArtifactRelayConfiguration::default())
    }
}

impl ArtifactRelayLedger {
    pub(super) fn new(configuration: ArtifactRelayConfiguration) -> Self {
        Self {
            configuration,
            reservations: BTreeMap::new(),
            period: 0,
            project_used: BTreeMap::new(),
            tenant_used: BTreeMap::new(),
            account_used: BTreeMap::new(),
            ingress_used: 0,
            egress_used: 0,
            abandoned_or_failed_used: 0,
        }
    }

    pub(super) fn configure(
        &mut self,
        configuration: ArtifactRelayConfiguration,
    ) -> Result<(), ArtifactRelayError> {
        configuration.validate()?;
        self.configuration = configuration;
        Ok(())
    }

    pub(super) fn from_durable(
        configuration: ArtifactRelayConfiguration,
        state: ArtifactRelayDurableState,
    ) -> Self {
        Self {
            configuration,
            reservations: state
                .reservations
                .into_iter()
                .map(|(key, reservation)| {
                    (
                        key,
                        RelayReservation {
                            scope: RelayScope {
                                tenant: reservation.tenant,
                                project: reservation.project,
                                account: reservation.account,
                            },
                            remaining_reserved_bytes: reservation.remaining_reserved_bytes,
                            ingress_bytes: reservation.ingress_bytes,
                            egress_bytes: reservation.egress_bytes,
                            expires_at_epoch_seconds: reservation.expires_at_epoch_seconds,
                        },
                    )
                })
                .collect(),
            period: state.period,
            project_used: state
                .project_used
                .into_iter()
                .map(|(tenant, project, bytes)| ((tenant, project), bytes))
                .collect(),
            tenant_used: state.tenant_used.into_iter().collect(),
            account_used: state
                .account_used
                .into_iter()
                .map(|(tenant, account, bytes)| ((tenant, account), bytes))
                .collect(),
            ingress_used: state.ingress_used,
            egress_used: state.egress_used,
            abandoned_or_failed_used: state.abandoned_or_failed_used,
        }
    }

    pub(super) fn durable_state(&self) -> ArtifactRelayDurableState {
        ArtifactRelayDurableState {
            reservations: self
                .reservations
                .iter()
                .map(|(key, reservation)| {
                    (
                        key.clone(),
                        ArtifactRelayReservationState {
                            tenant: reservation.scope.tenant.clone(),
                            project: reservation.scope.project.clone(),
                            account: reservation.scope.account.clone(),
                            remaining_reserved_bytes: reservation.remaining_reserved_bytes,
                            ingress_bytes: reservation.ingress_bytes,
                            egress_bytes: reservation.egress_bytes,
                            expires_at_epoch_seconds: reservation.expires_at_epoch_seconds,
                        },
                    )
                })
                .collect(),
            period: self.period,
            project_used: self
                .project_used
                .iter()
                .map(|((tenant, project), bytes)| (tenant.clone(), project.clone(), *bytes))
                .collect(),
            tenant_used: self
                .tenant_used
                .iter()
                .map(|(tenant, bytes)| (tenant.clone(), *bytes))
                .collect(),
            account_used: self
                .account_used
                .iter()
                .map(|((tenant, account), bytes)| (tenant.clone(), account.clone(), *bytes))
                .collect(),
            ingress_used: self.ingress_used,
            egress_used: self.egress_used,
            abandoned_or_failed_used: self.abandoned_or_failed_used,
        }
    }

    pub(super) fn reconcile_after_restart(&mut self, now_epoch_seconds: u64) {
        let reservations = self.reservations.keys().cloned().collect::<Vec<_>>();
        for key in reservations {
            self.finish(&key, RelayFinishReason::Failed);
        }
        self.prepare_period(now_epoch_seconds);
    }

    fn prepare_period(&mut self, now_epoch_seconds: u64) {
        let period = now_epoch_seconds / self.configuration.period_seconds.max(1);
        if self.period != period {
            self.period = period;
            self.project_used.clear();
            self.tenant_used.clear();
            self.account_used.clear();
            self.ingress_used = 0;
            self.egress_used = 0;
            self.abandoned_or_failed_used = 0;
        }
    }

    pub(super) fn estimated_wire_bytes(&self, artifact_bytes: u64, chunk_bytes: u64) -> u64 {
        let encoded = artifact_bytes
            .saturating_add(2)
            .saturating_div(3)
            .saturating_mul(4);
        let chunks = artifact_bytes
            .saturating_add(chunk_bytes.saturating_sub(1))
            .saturating_div(chunk_bytes.max(1))
            .max(1);
        encoded.saturating_mul(2).saturating_add(
            chunks
                .saturating_mul(2)
                .saturating_mul(self.configuration.framing_overhead_bytes),
        )
    }

    pub(super) fn framing_overhead_bytes(&self) -> u64 {
        self.configuration.framing_overhead_bytes
    }

    pub(super) fn reserve(
        &mut self,
        key: String,
        tenant: TenantId,
        project: ProjectId,
        account: UserId,
        artifact_bytes: u64,
        chunk_bytes: u64,
        expires_at_epoch_seconds: u64,
        now_epoch_seconds: u64,
    ) -> Result<(), ArtifactRelayError> {
        self.expire(now_epoch_seconds);
        self.prepare_period(now_epoch_seconds);
        if !self.configuration.enabled {
            return Err(ArtifactRelayError::Disabled);
        }
        if artifact_bytes > self.configuration.max_artifact_bytes {
            return Err(ArtifactRelayError::ArtifactTooLarge);
        }
        let scope = RelayScope {
            tenant,
            project,
            account,
        };
        self.check_concurrency(&scope)?;
        self.check_tracked_scope_capacity(&scope)?;
        let reserved = self.estimated_wire_bytes(artifact_bytes, chunk_bytes);
        self.check_budget(&scope, reserved)?;
        self.reservations.insert(
            key,
            RelayReservation {
                scope,
                remaining_reserved_bytes: reserved,
                ingress_bytes: 0,
                egress_bytes: 0,
                expires_at_epoch_seconds,
            },
        );
        Ok(())
    }

    pub(super) fn rekey(&mut self, from: &str, to: String) -> Result<(), ArtifactRelayError> {
        let reservation = self
            .reservations
            .remove(from)
            .ok_or(ArtifactRelayError::MissingReservation)?;
        self.reservations.insert(to, reservation);
        Ok(())
    }

    fn check_concurrency(&self, scope: &RelayScope) -> Result<(), ArtifactRelayError> {
        if self.reservations.len() >= self.configuration.max_active_global {
            return Err(ArtifactRelayError::Capacity);
        }
        let project = self
            .reservations
            .values()
            .filter(|reservation| {
                reservation.scope.tenant == scope.tenant
                    && reservation.scope.project == scope.project
            })
            .count();
        if project >= self.configuration.max_active_per_project {
            return Err(ArtifactRelayError::Concurrency("project"));
        }
        let tenant = self
            .reservations
            .values()
            .filter(|reservation| reservation.scope.tenant == scope.tenant)
            .count();
        if tenant >= self.configuration.max_active_per_tenant {
            return Err(ArtifactRelayError::Concurrency("tenant"));
        }
        let account = self
            .reservations
            .values()
            .filter(|reservation| {
                reservation.scope.tenant == scope.tenant
                    && reservation.scope.account == scope.account
            })
            .count();
        if account >= self.configuration.max_active_per_account {
            return Err(ArtifactRelayError::Concurrency("account"));
        }
        Ok(())
    }

    fn check_tracked_scope_capacity(&self, scope: &RelayScope) -> Result<(), ArtifactRelayError> {
        let project_key = (scope.tenant.clone(), scope.project.clone());
        let account_key = (scope.tenant.clone(), scope.account.clone());
        let new_scopes = usize::from(!self.project_used.contains_key(&project_key))
            + usize::from(!self.tenant_used.contains_key(&scope.tenant))
            + usize::from(!self.account_used.contains_key(&account_key));
        let tracked = self.project_used.len() + self.tenant_used.len() + self.account_used.len();
        if tracked.saturating_add(new_scopes) > self.configuration.max_tracked_scopes {
            return Err(ArtifactRelayError::StateLimit);
        }
        Ok(())
    }

    fn reserved_for(&self, scope: &RelayScope) -> (u64, u64, u64, u64) {
        let mut project = 0_u64;
        let mut tenant = 0_u64;
        let mut account = 0_u64;
        let mut global = 0_u64;
        for reservation in self.reservations.values() {
            let bytes = reservation.remaining_reserved_bytes;
            global = global.saturating_add(bytes);
            if reservation.scope.tenant == scope.tenant {
                tenant = tenant.saturating_add(bytes);
                if reservation.scope.project == scope.project {
                    project = project.saturating_add(bytes);
                }
                if reservation.scope.account == scope.account {
                    account = account.saturating_add(bytes);
                }
            }
        }
        (project, tenant, account, global)
    }

    fn check_budget(&self, scope: &RelayScope, additional: u64) -> Result<(), ArtifactRelayError> {
        let (project_reserved, tenant_reserved, account_reserved, _) = self.reserved_for(scope);
        let project_used = self
            .project_used
            .get(&(scope.tenant.clone(), scope.project.clone()))
            .copied()
            .unwrap_or(0);
        let tenant_used = self.tenant_used.get(&scope.tenant).copied().unwrap_or(0);
        let account_used = self
            .account_used
            .get(&(scope.tenant.clone(), scope.account.clone()))
            .copied()
            .unwrap_or(0);
        for (used, reserved, limit, label) in [
            (
                project_used,
                project_reserved,
                self.configuration.max_period_bytes_per_project,
                "project",
            ),
            (
                tenant_used,
                tenant_reserved,
                self.configuration.max_period_bytes_per_tenant,
                "tenant",
            ),
            (
                account_used,
                account_reserved,
                self.configuration.max_period_bytes_per_account,
                "account",
            ),
        ] {
            if used.saturating_add(reserved).saturating_add(additional) > limit {
                return Err(ArtifactRelayError::PeriodBudget(label));
            }
        }
        Ok(())
    }

    fn charge(
        &mut self,
        key: &str,
        bytes: u64,
        ingress: bool,
        now_epoch_seconds: u64,
    ) -> Result<(), ArtifactRelayError> {
        self.prepare_period(now_epoch_seconds);
        let reservation = self
            .reservations
            .get(key)
            .cloned()
            .ok_or(ArtifactRelayError::MissingReservation)?;
        let additional = bytes.saturating_sub(reservation.remaining_reserved_bytes);
        if additional > 0 {
            self.check_budget(&reservation.scope, additional)?;
        }
        let reservation = self
            .reservations
            .get_mut(key)
            .ok_or(ArtifactRelayError::MissingReservation)?;
        reservation.remaining_reserved_bytes =
            reservation.remaining_reserved_bytes.saturating_sub(bytes);
        if ingress {
            reservation.ingress_bytes = reservation.ingress_bytes.saturating_add(bytes);
            self.ingress_used = self.ingress_used.saturating_add(bytes);
        } else {
            reservation.egress_bytes = reservation.egress_bytes.saturating_add(bytes);
            self.egress_used = self.egress_used.saturating_add(bytes);
        }
        let project_key = (
            reservation.scope.tenant.clone(),
            reservation.scope.project.clone(),
        );
        let account_key = (
            reservation.scope.tenant.clone(),
            reservation.scope.account.clone(),
        );
        *self.project_used.entry(project_key).or_default() = self
            .project_used
            .get(&(
                reservation.scope.tenant.clone(),
                reservation.scope.project.clone(),
            ))
            .copied()
            .unwrap_or(0)
            .saturating_add(bytes);
        *self
            .tenant_used
            .entry(reservation.scope.tenant.clone())
            .or_default() = self
            .tenant_used
            .get(&reservation.scope.tenant)
            .copied()
            .unwrap_or(0)
            .saturating_add(bytes);
        *self.account_used.entry(account_key).or_default() = self
            .account_used
            .get(&(
                reservation.scope.tenant.clone(),
                reservation.scope.account.clone(),
            ))
            .copied()
            .unwrap_or(0)
            .saturating_add(bytes);
        Ok(())
    }

    pub(super) fn charge_ingress(
        &mut self,
        key: &str,
        bytes: u64,
        now_epoch_seconds: u64,
    ) -> Result<(), ArtifactRelayError> {
        self.charge(key, bytes, true, now_epoch_seconds)
    }

    pub(super) fn charge_egress(
        &mut self,
        key: &str,
        bytes: u64,
        now_epoch_seconds: u64,
    ) -> Result<(), ArtifactRelayError> {
        self.charge(key, bytes, false, now_epoch_seconds)
    }

    pub(super) fn finish(&mut self, key: &str, reason: RelayFinishReason) {
        if let Some(reservation) = self.reservations.remove(key) {
            if reason != RelayFinishReason::Completed {
                self.abandoned_or_failed_used = self
                    .abandoned_or_failed_used
                    .saturating_add(reservation.ingress_bytes)
                    .saturating_add(reservation.egress_bytes);
            }
        }
    }

    pub(super) fn expire(&mut self, now_epoch_seconds: u64) {
        let expired = self
            .reservations
            .iter()
            .filter(|(_, reservation)| reservation.expires_at_epoch_seconds < now_epoch_seconds)
            .map(|(key, _)| key.clone())
            .collect::<Vec<_>>();
        for key in expired {
            self.finish(&key, RelayFinishReason::Expired);
        }
    }

    pub(super) fn usage(&self) -> ArtifactRelayUsage {
        ArtifactRelayUsage {
            active_transfers: self.reservations.len(),
            ingress_bytes: self.ingress_used,
            egress_bytes: self.egress_used,
            abandoned_or_failed_bytes: self.abandoned_or_failed_used,
            reserved_bytes: self
                .reservations
                .values()
                .map(|reservation| reservation.remaining_reserved_bytes)
                .fold(0_u64, u64::saturating_add),
        }
    }
}