use std::{
    collections::hash_map::DefaultHasher,
    fs,
    hash::{Hash, Hasher},
    path::PathBuf,
};

use rocksdb::{DB, IteratorMode, Options};
use serde::{Deserialize, Serialize};
use serde_json;

use crate::error::{EventError, Result};

pub struct TenantAssignmentStore {
    db: DB,
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct TenantRecord {
    #[serde(default)]
    pub shard: Option<String>,
    #[serde(default, rename = "storage_quota_mb", alias = "aggregate_quota")]
    pub storage_quota_mb: Option<u64>,
    #[serde(default, rename = "storage_usage_bytes", alias = "aggregate_count")]
    pub storage_usage_bytes: Option<u64>,
}

impl TenantRecord {
    fn cleanup(&mut self) {
        if let Some(shard) = &self.shard {
            let trimmed = shard.trim();
            if trimmed.is_empty() {
                self.shard = None;
            } else if trimmed != shard {
                self.shard = Some(trimmed.to_string());
            }
        }
    }

    fn is_empty(&self) -> bool {
        self.shard.is_none()
            && self.storage_quota_mb.is_none()
            && self.storage_usage_bytes.is_none()
    }
}

pub const BYTES_PER_MEGABYTE: u64 = 1024 * 1024;
pub const SHARD_PREFIX: &str = "shard-";

impl TenantAssignmentStore {
    pub fn open(path: PathBuf) -> Result<Self> {
        fs::create_dir_all(&path)?;
        let mut opts = Options::default();
        opts.create_if_missing(true);
        let db = DB::open(&opts, path).map_err(map_db_error)?;
        Ok(Self { db })
    }

    pub fn open_read_only(path: PathBuf) -> Result<Self> {
        if path.exists() {
            let mut opts = Options::default();
            opts.create_if_missing(false);
            let db = DB::open_for_read_only(&opts, path, false).map_err(map_db_error)?;
            Ok(Self { db })
        } else {
            Self::open(path)
        }
    }

    fn key(tenant: &str) -> Vec<u8> {
        tenant.as_bytes().to_vec()
    }

    fn read_record(&self, tenant: &str) -> Result<Option<TenantRecord>> {
        let key = Self::key(tenant);
        let value = match self.db.get(&key).map_err(map_db_error)? {
            Some(bytes) => bytes,
            None => return Ok(None),
        };
        let mut record = decode_record(value)?;
        record.cleanup();
        if record.is_empty() {
            Ok(None)
        } else {
            Ok(Some(record))
        }
    }

    fn write_record(&self, tenant: &str, record: TenantRecord) -> Result<()> {
        let key = Self::key(tenant);
        if record.is_empty() {
            self.db.delete(key).map_err(map_db_error)?;
        } else {
            self.db
                .put(
                    key,
                    serde_json::to_vec(&record).map_err(|err| {
                        EventError::Serialization(format!("invalid tenant record: {err}"))
                    })?,
                )
                .map_err(map_db_error)?;
        }
        Ok(())
    }

    pub fn assign(&self, tenant: &str, shard: &str) -> Result<bool> {
        let mut record = self.read_record(tenant)?.unwrap_or_default();
        let shard_string = Some(shard.to_string());
        let changed = record.shard != shard_string;
        record.shard = shard_string;
        self.write_record(tenant, record)?;
        Ok(changed)
    }

    pub fn unassign(&self, tenant: &str) -> Result<bool> {
        let Some(mut record) = self.read_record(tenant)? else {
            return Ok(false);
        };
        if record.shard.is_none() {
            return Ok(false);
        }
        record.shard = None;
        self.write_record(tenant, record)?;
        Ok(true)
    }

    pub fn shard_for(&self, tenant: &str) -> Result<Option<String>> {
        Ok(self.read_record(tenant)?.and_then(|record| record.shard))
    }

    pub fn quota_for(&self, tenant: &str) -> Result<Option<u64>> {
        Ok(self
            .read_record(tenant)?
            .and_then(|record| record.storage_quota_mb))
    }

    pub fn set_quota(&self, tenant: &str, quota: Option<u64>) -> Result<bool> {
        let mut record = self.read_record(tenant)?.unwrap_or_default();
        let changed = record.storage_quota_mb != quota;
        record.storage_quota_mb = quota;
        self.write_record(tenant, record)?;
        Ok(changed)
    }

    pub fn storage_usage_bytes(&self, tenant: &str) -> Result<Option<u64>> {
        Ok(self
            .read_record(tenant)?
            .and_then(|record| record.storage_usage_bytes))
    }

    pub fn ensure_storage_usage_bytes<F>(&self, tenant: &str, initializer: F) -> Result<u64>
    where
        F: FnOnce() -> Result<u64>,
    {
        let mut record = self.read_record(tenant)?.unwrap_or_default();
        if let Some(bytes) = record.storage_usage_bytes {
            return Ok(bytes);
        }
        let bytes = initializer()?;
        record.storage_usage_bytes = Some(bytes);
        self.write_record(tenant, record)?;
        Ok(bytes)
    }

    pub fn update_storage_usage_bytes(&self, tenant: &str, usage: u64) -> Result<u64> {
        let mut record = self.read_record(tenant)?.unwrap_or_default();
        record.storage_usage_bytes = Some(usage);
        self.write_record(tenant, record)?;
        Ok(usage)
    }

    pub fn record_for(&self, tenant: &str) -> Result<Option<TenantRecord>> {
        self.read_record(tenant)
    }

    pub fn list(&self) -> Result<Vec<(String, TenantRecord)>> {
        let mut entries = Vec::new();
        for entry in self.db.iterator(IteratorMode::Start) {
            let (key, value) = entry.map_err(map_db_error)?;
            let tenant = String::from_utf8(key.to_vec())
                .map_err(|err| EventError::Serialization(format!("invalid tenant entry: {err}")))?;
            let mut record = decode_record(value.to_vec())?;
            record.cleanup();
            if record.is_empty() {
                continue;
            }
            entries.push((tenant, record));
        }
        Ok(entries)
    }
}

fn decode_record(bytes: Vec<u8>) -> Result<TenantRecord> {
    match serde_json::from_slice::<TenantRecord>(&bytes) {
        Ok(record) => Ok(record),
        Err(_) => {
            let shard = String::from_utf8(bytes)
                .map_err(|err| EventError::Serialization(format!("invalid tenant entry: {err}")))?;
            let shard = shard.trim();
            if shard.is_empty() {
                Ok(TenantRecord::default())
            } else {
                Ok(TenantRecord {
                    shard: Some(shard.to_string()),
                    storage_quota_mb: None,
                    storage_usage_bytes: None,
                })
            }
        }
    }
}

pub fn compute_default_shard(tenant: &str, shard_count: u16) -> String {
    let mut hasher = DefaultHasher::new();
    tenant.hash(&mut hasher);
    let max = shard_count.max(1) as u64;
    let index = (hasher.finish() % max) as u16;
    format_shard_id(index)
}

pub fn format_shard_id(index: u16) -> String {
    format!("{SHARD_PREFIX}{index:04}")
}

fn map_db_error(err: rocksdb::Error) -> EventError {
    EventError::Storage(format!("tenant metadata error: {err}"))
}
