use crate::decode::{self, DecodeError, DecodedImage};
use crate::model::{ImageId, ImageRecord};
use std::collections::HashMap;
use std::sync::{Arc, Condvar, Mutex, MutexGuard};

/// RAII handle to a decoded image held by the pool. While alive it pins the
/// image against eviction; dropping it frees the pin and wakes waiters.
pub struct PixelGuard {
    img: Option<Arc<DecodedImage>>,
    pool: Arc<PoolInner>,
}

impl std::ops::Deref for PixelGuard {
    type Target = DecodedImage;
    fn deref(&self) -> &DecodedImage {
        self.img.as_ref().expect("guard used after drop")
    }
}

impl Drop for PixelGuard {
    fn drop(&mut self) {
        // Drop our Arc first (lowering strong_count so the entry may become
        // idle/evictable), then wake any get() blocked waiting for room.
        self.img = None;
        let _guard = self.pool.state.lock().unwrap();
        self.pool.cv.notify_all();
    }
}

/// Access to decoded pixels. Holding a returned guard pins the image in memory.
pub trait PixelStore: Sync {
    fn get(&self, record: &ImageRecord) -> Result<PixelGuard, DecodeError>;
    fn get_pair(
        &self,
        a: &ImageRecord,
        b: &ImageRecord,
    ) -> Result<(PixelGuard, PixelGuard), DecodeError>;
}

struct PoolState {
    used: u64,
    entries: HashMap<ImageId, Arc<DecodedImage>>,
    lru: Vec<ImageId>, // front = least-recently-used
}

struct PoolInner {
    budget: u64,
    state: Mutex<PoolState>,
    cv: Condvar,
    // Serializes the reservation phase of `get_pair` so two matchers can never
    // both hold one pinned image while waiting for a second (partial-hold
    // deadlock). Always locked BEFORE `state`; guard-drop only touches `state`,
    // so there is no lock-order inversion.
    acquire: Mutex<()>,
}

/// A byte-bounded decoded-image pool. Resident images count against `budget`;
/// idle entries (held only by the pool) are evicted LRU to make room, and a
/// fetch blocks if all resident bytes are pinned by active guards.
pub struct PooledPixelStore {
    inner: Arc<PoolInner>,
}

impl PooledPixelStore {
    pub fn new(budget: u64) -> Self {
        PooledPixelStore {
            inner: Arc::new(PoolInner {
                budget: budget.max(1),
                state: Mutex::new(PoolState { used: 0, entries: HashMap::new(), lru: Vec::new() }),
                cv: Condvar::new(),
                acquire: Mutex::new(()),
            }),
        }
    }

    /// Current resident bytes (test/observability helper).
    pub fn used_bytes(&self) -> u64 {
        self.inner.state.lock().unwrap().used
    }

    fn guard(&self, arc: Arc<DecodedImage>) -> PixelGuard {
        PixelGuard { img: Some(arc), pool: self.inner.clone() }
    }
}

fn size_of(record: &ImageRecord) -> u64 {
    record.width as u64 * record.height as u64 * 4
}

impl PoolInner {
    /// Move `id` to the most-recently-used position.
    fn touch(st: &mut PoolState, id: ImageId) {
        if let Some(pos) = st.lru.iter().position(|&x| x == id) {
            st.lru.remove(pos);
            st.lru.push(id);
        }
    }

    /// Block until at least `need` free bytes exist, evicting idle (unpinned,
    /// non-excluded) LRU entries first. Returns the still-held guard.
    fn make_room<'a>(
        self: &Arc<Self>,
        mut st: MutexGuard<'a, PoolState>,
        need: u64,
        exclude: &[ImageId],
    ) -> MutexGuard<'a, PoolState> {
        loop {
            if st.used + need <= self.budget {
                return st;
            }
            let victim = st.lru.iter().position(|id| {
                !exclude.contains(id) && Arc::strong_count(&st.entries[id]) == 1
            });
            match victim {
                Some(pos) => {
                    let id = st.lru.remove(pos);
                    let arc = st.entries.remove(&id).unwrap();
                    st.used -= arc.rgba.len() as u64;
                    drop(arc);
                }
                None => {
                    st = self.cv.wait(st).unwrap();
                }
            }
        }
    }

    /// Decode `path` and install it under `id`, consuming the `size` bytes that
    /// the caller already reserved. Handles a decode race (another thread
    /// installed the same id meanwhile) by refunding the reservation.
    fn install(
        self: &Arc<Self>,
        id: ImageId,
        path: &std::path::Path,
        size: u64,
    ) -> Result<Arc<DecodedImage>, DecodeError> {
        let decoded = decode::load(path);
        let mut st = self.state.lock().unwrap();
        match decoded {
            Ok(img) => {
                if let Some(arc) = st.entries.get(&id).cloned() {
                    st.used -= size; // lost the race; refund
                    Self::touch(&mut st, id);
                    drop(st);
                    self.cv.notify_all();
                    Ok(arc)
                } else {
                    let arc = Arc::new(img);
                    st.entries.insert(id, arc.clone());
                    st.lru.push(id);
                    Ok(arc) // `used` already includes `size` (reserved by caller)
                }
            }
            Err(e) => {
                st.used -= size;
                drop(st);
                self.cv.notify_all();
                Err(e)
            }
        }
    }

    fn acquire(self: &Arc<Self>, record: &ImageRecord) -> Result<Arc<DecodedImage>, DecodeError> {
        let id = record.id;
        let size = size_of(record);
        {
            let mut st = self.state.lock().unwrap();
            if let Some(arc) = st.entries.get(&id).cloned() {
                Self::touch(&mut st, id);
                return Ok(arc);
            }
            st = self.make_room(st, size, &[]);
            st.used += size;
        }
        self.install(id, &record.path, size)
    }

    fn acquire_two(
        self: &Arc<Self>,
        a: &ImageRecord,
        b: &ImageRecord,
    ) -> Result<(Arc<DecodedImage>, Arc<DecodedImage>), DecodeError> {
        let (ida, idb) = (a.id, b.id);
        let (sa, sb) = (size_of(a), size_of(b));

        // Phase 1: serialized by `acquire` so only one matcher reserves a
        // working set at a time (deadlock-free). Pin any cached members (cloning
        // prevents eviction while we decode the others) and reserve room for the
        // misses. We loop inside the acquire guard so that after a cv.wait we
        // re-read the cache — an image installed by the *previous* holder of
        // `acquire` may now be present, reducing `need` to 0 and avoiding the
        // scenario where a just-installed idle image is in `exclude` but counts
        // against `budget`, causing permanent stall on a tight budget.
        let (cached_a, cached_b) = {
            let _acq = self.acquire.lock().unwrap();
            let mut st = self.state.lock().unwrap();
            loop {
                let ca = st.entries.get(&ida).cloned();
                let cb = st.entries.get(&idb).cloned();
                let need = (ca.is_none() as u64) * sa + (cb.is_none() as u64) * sb;
                // Only exclude already-pinned (cloned) entries from eviction.
                let mut excl: Vec<ImageId> = Vec::with_capacity(2);
                if ca.is_some() {
                    excl.push(ida);
                }
                if cb.is_some() {
                    excl.push(idb);
                }
                if st.used + need <= self.budget {
                    st.used += need;
                    if ca.is_some() {
                        Self::touch(&mut st, ida);
                    }
                    if cb.is_some() {
                        Self::touch(&mut st, idb);
                    }
                    break (ca, cb);
                }
                // Try to evict an idle entry that isn't one we've already pinned.
                let victim = st.lru.iter().position(|id| {
                    !excl.contains(id) && Arc::strong_count(&st.entries[id]) == 1
                });
                match victim {
                    Some(pos) => {
                        let id = st.lru.remove(pos);
                        let arc = st.entries.remove(&id).unwrap();
                        st.used -= arc.rgba.len() as u64;
                        drop(arc);
                    }
                    None => {
                        st = self.cv.wait(st).unwrap();
                    }
                }
            }
        };

        // Phase 2: decode the misses (room already reserved). On error the
        // reservation is refunded inside `install`.
        let arc_a = match cached_a {
            Some(x) => x,
            None => self.install(ida, &a.path, sa)?,
        };
        let arc_b = match cached_b {
            Some(x) => x,
            None => self.install(idb, &b.path, sb)?,
        };
        Ok((arc_a, arc_b))
    }
}

impl PixelStore for PooledPixelStore {
    fn get(&self, record: &ImageRecord) -> Result<PixelGuard, DecodeError> {
        let arc = self.inner.acquire(record)?;
        Ok(self.guard(arc))
    }

    fn get_pair(
        &self,
        a: &ImageRecord,
        b: &ImageRecord,
    ) -> Result<(PixelGuard, PixelGuard), DecodeError> {
        let (arc_a, arc_b) = self.inner.acquire_two(a, b)?;
        Ok((self.guard(arc_a), self.guard(arc_b)))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::hash::Hash64;
    use crate::model::ImageRecord;
    use image::{Rgba, RgbaImage};
    use std::sync::Arc;
    use std::thread;

    fn write_png(dir: &std::path::Path, name: &str, c: u8) -> std::path::PathBuf {
        let mut img = RgbaImage::new(2, 2); // 2*2*4 = 16 bytes decoded
        for p in img.pixels_mut() {
            *p = Rgba([c, c, c, 255]);
        }
        let path = dir.join(name);
        img.save(&path).unwrap();
        path
    }

    fn record(id: usize, path: std::path::PathBuf) -> ImageRecord {
        ImageRecord { id, path, hash: Hash64(0), width: 2, height: 2 }
    }

    #[test]
    fn get_returns_pixels_and_caches() {
        let dir = tempfile::tempdir().unwrap();
        let rec = record(0, write_png(dir.path(), "a.png", 7));
        let store = PooledPixelStore::new(1 << 20);

        let g = store.get(&rec).unwrap();
        assert_eq!(g.width, 2);
        assert_eq!(&g.rgba[0..4], &[7, 7, 7, 255]);
        assert_eq!(store.used_bytes(), 16);

        drop(g);
        std::fs::remove_file(&rec.path).unwrap();
        let g2 = store.get(&rec).unwrap();
        assert_eq!(&g2.rgba[0..4], &[7, 7, 7, 255]);
    }

    #[test]
    fn evicts_idle_entry_to_stay_within_budget() {
        let dir = tempfile::tempdir().unwrap();
        let a = record(0, write_png(dir.path(), "a.png", 1));
        let b = record(1, write_png(dir.path(), "b.png", 2));
        let store = PooledPixelStore::new(16); // room for exactly one 16-byte image

        {
            let _g = store.get(&a).unwrap();
            assert_eq!(store.used_bytes(), 16);
        }

        let g = store.get(&b).unwrap();
        assert_eq!(&g.rgba[0..4], &[2, 2, 2, 255]);
        assert_eq!(store.used_bytes(), 16);
    }

    #[test]
    fn get_pair_returns_both() {
        let dir = tempfile::tempdir().unwrap();
        let a = record(0, write_png(dir.path(), "a.png", 1));
        let b = record(1, write_png(dir.path(), "b.png", 2));
        let store = PooledPixelStore::new(1 << 20);

        let (ga, gb) = store.get_pair(&a, &b).unwrap();
        assert_eq!(&ga.rgba[0..4], &[1, 1, 1, 255]);
        assert_eq!(&gb.rgba[0..4], &[2, 2, 2, 255]);
        assert_eq!(store.used_bytes(), 32);
    }

    #[test]
    fn concurrent_get_pair_under_tight_budget_does_not_deadlock() {
        let dir = tempfile::tempdir().unwrap();
        let a = record(0, write_png(dir.path(), "a.png", 1));
        let b = record(1, write_png(dir.path(), "b.png", 2));
        let store = Arc::new(PooledPixelStore::new(32)); // exactly one pair

        let mut handles = Vec::new();
        for _ in 0..2 {
            let store = store.clone();
            let (a, b) = (a.clone(), b.clone());
            handles.push(thread::spawn(move || {
                for _ in 0..50 {
                    let (ga, gb) = store.get_pair(&a, &b).unwrap();
                    assert_eq!(ga.width, 2);
                    assert_eq!(gb.width, 2);
                }
            }));
        }
        for h in handles {
            h.join().unwrap();
        }
        assert!(store.used_bytes() <= 32);
    }

    #[test]
    fn concurrent_get_pair_distinct_pairs_tight_budget_does_not_deadlock() {
        // Four distinct images, budget = one pair. Two threads request DIFFERENT
        // pairs concurrently — the classic partial-hold deadlock scenario that
        // the `acquire` mutex must prevent.
        let dir = tempfile::tempdir().unwrap();
        let a = record(0, write_png(dir.path(), "a.png", 1));
        let b = record(1, write_png(dir.path(), "b.png", 2));
        let c = record(2, write_png(dir.path(), "c.png", 3));
        let d = record(3, write_png(dir.path(), "d.png", 4));
        let store = Arc::new(PooledPixelStore::new(32)); // one pair fits at a time

        let s1 = store.clone();
        let (a1, b1) = (a.clone(), b.clone());
        let h1 = thread::spawn(move || {
            for _ in 0..50 {
                let _ = s1.get_pair(&a1, &b1).unwrap();
            }
        });
        let s2 = store.clone();
        let h2 = thread::spawn(move || {
            for _ in 0..50 {
                let _ = s2.get_pair(&c, &d).unwrap();
            }
        });
        h1.join().unwrap();
        h2.join().unwrap();
        assert!(store.used_bytes() <= 32);
    }

    #[test]
    fn mixed_get_and_get_pair_does_not_deadlock_or_overflow() {
        // Interleave single get() (no `acquire` mutex) with get_pair() under a
        // one-pair budget — the path most likely to corrupt accounting.
        let dir = tempfile::tempdir().unwrap();
        let a = record(0, write_png(dir.path(), "a.png", 1));
        let b = record(1, write_png(dir.path(), "b.png", 2));
        let store = Arc::new(PooledPixelStore::new(32));

        let s1 = store.clone();
        let (a1, b1) = (a.clone(), b.clone());
        let h1 = thread::spawn(move || {
            for _ in 0..100 {
                let _ = s1.get_pair(&a1, &b1).unwrap();
            }
        });
        let s2 = store.clone();
        let a2 = a.clone();
        let h2 = thread::spawn(move || {
            for _ in 0..100 {
                let g = s2.get(&a2).unwrap();
                assert_eq!(g.width, 2);
            }
        });
        h1.join().unwrap();
        h2.join().unwrap();
        assert!(store.used_bytes() <= 32);
    }

    #[test]
    fn cache_hit_while_pinned_does_not_double_count() {
        let dir = tempfile::tempdir().unwrap();
        let a = record(0, write_png(dir.path(), "a.png", 9));
        let store = PooledPixelStore::new(64);

        let g1 = store.get(&a).unwrap();
        assert_eq!(store.used_bytes(), 16);
        // A second get of the same (still-pinned) image reuses the cached entry;
        // resident bytes must not grow.
        let g2 = store.get(&a).unwrap();
        assert_eq!(store.used_bytes(), 16);
        assert_eq!(g1.rgba, g2.rgba);
    }
}
