use std::{
    sync::Arc,
    time::{Duration, Instant},
};

use anyhow::Result;
use quick_cache::{
    Weighter,
    sync::{Cache, GuardResult},
};
use tantivy::{DocAddress, TantivyDocument};
use tici_system::readable::ReadableSize;

use crate::utils::metrics::{
    inc_doc_cache_coalesced_wait_count, inc_doc_cache_load_miss_count,
    inc_doc_cache_resident_hit_count, observe_doc_cache_load_latency,
};

// Heuristics for `quick_cache` byte limits (`capacity_bytes`) and internal
// sizing.
//
// `ESTIMATED_DOC_BYTES`: assumed average serialized doc size used only by
// `estimated_items_capacity` to derive an entry-count hint for
// `Cache::with_weighter`. It does not affect eviction weight.
//
// `DOC_FIXED_OVERHEAD` / `DOC_FIELD_OVERHEAD`: fudge terms in `doc_weight` so
// total weight tracks memory pressure better than `node_data.capacity()`
// alone—fixed shell cost per `Arc<TantivyDocument>` plus a linear term per
// field (`doc.len()`).
//
// All values are approximate; adjust with profiling if tighter accounting is
// needed.
const ESTIMATED_DOC_BYTES: u64 = 4 * 1024;
const DOC_FIXED_OVERHEAD: u64 = 64;
const DOC_FIELD_OVERHEAD: u64 = 16;

type InnerCache = Cache<DocCacheKey, Arc<TantivyDocument>, DocWeighter>;

#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct DocCacheKey {
    frag_id: u64,
    doc_addr: DocAddress,
}

impl DocCacheKey {
    pub fn new(frag_id: u64, doc_addr: DocAddress) -> Self {
        Self { frag_id, doc_addr }
    }

    fn bytes(&self) -> u64 {
        std::mem::size_of_val(self) as u64
    }
}

#[derive(Clone, Debug, Default)]
struct DocWeighter;

impl Weighter<DocCacheKey, Arc<TantivyDocument>> for DocWeighter {
    fn weight(&self, key: &DocCacheKey, doc: &Arc<TantivyDocument>) -> u64 {
        key.bytes() + doc_weight(doc)
    }
}

pub struct DocCache {
    cache: InnerCache,
}

impl DocCache {
    pub fn new(capacity: ReadableSize) -> Self {
        Self {
            cache: build_cache(capacity.0),
        }
    }

    pub fn get_or_insert_with(
        &self,
        key: DocCacheKey,
        load_doc: impl FnOnce() -> Result<TantivyDocument>,
    ) -> Result<Arc<TantivyDocument>> {
        // `quick_cache` serializes concurrent loads with placeholder guards.
        // `Duration::ZERO` makes this probe non-blocking: Value = resident hit;
        // Guard = miss and this thread loads; Timeout = a peer holds the insert
        // guard, so fall back to `get_or_insert_with` to wait on the in-flight
        // load. Metrics: resident_hit / load_miss / coalesced_wait.
        match self.cache.get_value_or_guard(&key, Some(Duration::ZERO)) {
            GuardResult::Value(doc) => {
                inc_doc_cache_resident_hit_count();
                Ok(doc)
            }
            GuardResult::Guard(guard) => {
                inc_doc_cache_load_miss_count();
                let result = load_doc_for_cache(load_doc);
                if let Ok(doc) = &result {
                    // `PlaceholderGuard::insert` returns `Err(value)` only when the
                    // placeholder was removed (e.g. `remove` or a non-guard
                    // `insert` on the same key). DocCache never does either, so
                    // this should not fail in practice; ignore and still return
                    // the loaded doc (uncached until the next access).
                    let _ = guard.insert(doc.clone());
                }
                result
            }
            GuardResult::Timeout => {
                let mut loaded = false;
                let result = self.cache.get_or_insert_with(&key, || {
                    loaded = true;
                    load_doc_for_cache(load_doc)
                });
                if loaded {
                    inc_doc_cache_load_miss_count();
                } else {
                    inc_doc_cache_coalesced_wait_count();
                }
                result
            }
        }
    }
}

fn load_doc_for_cache(
    load_doc: impl FnOnce() -> Result<TantivyDocument>,
) -> Result<Arc<TantivyDocument>> {
    let start = Instant::now();
    let result = load_doc().map(|mut doc| {
        doc.shrink_to_fit();
        Arc::new(doc)
    });
    let result_label = if result.is_ok() { "ok" } else { "error" };
    observe_doc_cache_load_latency(result_label, start.elapsed());
    result
}

fn build_cache(capacity_bytes: u64) -> InnerCache {
    Cache::with_weighter(
        estimated_items_capacity(capacity_bytes),
        capacity_bytes,
        DocWeighter,
    )
}

fn estimated_items_capacity(capacity_bytes: u64) -> usize {
    let estimated = capacity_bytes / ESTIMATED_DOC_BYTES;
    estimated.clamp(1, usize::MAX as u64) as usize
}

fn doc_weight(doc: &TantivyDocument) -> u64 {
    DOC_FIXED_OVERHEAD + doc.node_data.capacity() as u64 + doc.len() as u64 * DOC_FIELD_OVERHEAD
}