<reponame>lebai-robotics/simple-dns
use std::collections::HashMap;
use std::time::{Duration, Instant};

use simple_dns::ResourceRecord;

#[derive(Debug)]
pub struct ResourceRecordManager<'a> {
    pub resources: HashMap<ResourceRecord<'a>, ResourceRecordType>,
}

impl<'a> ResourceRecordManager<'a> {
    pub fn new() -> Self {
        Self {
            resources: HashMap::new(),
        }
    }

    /// Register a Resource Record
    pub fn add_owned_resource(&mut self, resource: ResourceRecord<'a>) {
        self.resources.insert(resource, ResourceRecordType::Owned);
    }

    pub fn add_expirable_resource(&mut self, resource: ResourceRecord<'a>) {
        let exp_info = ExpirationInfo::new(resource.ttl);
        self.resources
            .insert(resource, ResourceRecordType::Expirable(exp_info));
    }

    pub fn remove_resource_record(&mut self, resource_record: &ResourceRecord<'a>) {
        self.resources.remove(resource_record);
    }

    /// Remove all resource records
    pub fn clear(&mut self) {
        self.resources = HashMap::new();
    }
}

#[derive(Debug, Clone)]
pub enum ResourceRecordType {
    Owned,
    Expirable(ExpirationInfo),
}

impl<'a> ResourceRecordType {
    pub fn is_owned(&self) -> bool {
        matches!(self, &ResourceRecordType::Owned)
    }
    pub fn is_expired(&self) -> bool {
        match self {
            ResourceRecordType::Owned => false,
            ResourceRecordType::Expirable(exp_info) => exp_info.expire_at < Instant::now(),
        }
    }

    pub fn expired_time(&self) -> Option<Duration> {
        match self {
            ResourceRecordType::Owned => None,
            ResourceRecordType::Expirable(exp_info) => {
                let now = Instant::now();
                if exp_info.expire_at < now {
                    Some(now - exp_info.expire_at)
                } else {
                    None
                }
            }
        }
    }

    pub fn should_refresh(&self) -> bool {
        match self {
            ResourceRecordType::Owned => false,
            ResourceRecordType::Expirable(exp_info) => exp_info.refresh_at < Instant::now(),
        }
    }
}

/// Provides known service expiration and refresh times
#[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
pub struct ExpirationInfo {
    refresh_at: Instant,
    expire_at: Instant,
}

impl ExpirationInfo {
    pub fn new(ttl: u32) -> Self {
        let ttl = ttl as u64;
        let added = Instant::now();
        let expire_at = added + Duration::from_secs(ttl);
        let refresh_at = match ttl {
            0 => expire_at,
            ttl if ttl < 60 => added + Duration::from_secs(ttl / 2),
            ttl => added + Duration::from_secs(ttl / 10 * 8),
        };

        Self {
            expire_at,
            refresh_at,
        }
    }
}
