use crate::error::SKResult;
use crate::meta::SKFileMeta;
use lru::LruCache;
use niffler::Level;
use niffler::send::compression::Format;
use obikseq::SuperKmer;
use rustix::process::{Resource, getrlimit};
use std::fs::{File, OpenOptions};
use std::io::{BufWriter, Write};
use std::num::NonZeroUsize;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, OnceLock};

/// Hard upper bound on pool size regardless of OS fd limit.
pub const MAX_POOL_SIZE: usize = 512;

/// Default pending buffer threshold (bytes) before draining to the physical fd.
pub const DEFAULT_FLUSH_THRESHOLD: usize = 64 * 1024;

// Convenient alias for the per-entry physical writer slot.
type PhysWriter = Option<Box<dyn Write + Send>>;

// ── WriteEntry ─────────────────────────────────────────────────────────────────

struct WriteEntry {
    path: PathBuf,
    format: Format,
    level: Level,
    /// Per-entry mutex for the physical fd.
    /// Independent of the pool mutex to allow parallel writes to different entries.
    fd: Arc<Mutex<PhysWriter>>,
    logically_closed: bool,
}

// ── SKFilePool ─────────────────────────────────────────────────────────────────

/// LRU pool of open write file descriptors.
///
/// # Locking model
///
/// Two independent locks:
///
/// | Lock | Scope | Held during |
/// |---|---|---|
/// | `Arc<Mutex<SKFilePool>>` (pool lock) | all entries | LRU management only — microseconds |
/// | `Arc<Mutex<PhysWriter>>` (entry lock) | one entry | Writing a pending buffer to the fd |
///
/// **Pool lock is never held while writing data.** This allows parallel writes
/// to different partitions.
///
/// Lock ordering: always pool lock → entry lock, never reverse.
///
/// # Eviction
///
/// `evict_lru()` uses `try_lock()` on entry fd locks.
/// An entry whose fd is currently being written to is skipped; the next LRU
/// candidate is tried instead.  If all open fds are in use, an error is returned.
pub struct SKFilePool {
    max_open: usize,
    /// All registered entries. Index = stable token id. Never shrinks.
    entries: Vec<WriteEntry>,
    /// IDs of entries currently holding an open fd, in LRU order.
    ///
    /// Invariant: `id ∈ open ↔ entries[id].fd.lock().is_some()`
    open: LruCache<usize, ()>,
}

/// Shared reference to a pool; the primary way to create `SKFileWriter`s.
pub type SharedPool = Arc<Mutex<SKFilePool>>;

static GLOBAL_POOL: OnceLock<SharedPool> = OnceLock::new();

fn global_pool() -> &'static SharedPool {
    GLOBAL_POOL.get_or_init(|| Arc::new(Mutex::new(SKFilePool::from_system_limits())))
}

impl SKFilePool {
    /// Create a pool allowing at most `max_open` simultaneously open fds.
    pub fn new(max_open: usize) -> Self {
        let cap = NonZeroUsize::new(max_open.max(1)).unwrap();
        Self {
            max_open,
            entries: Vec::new(),
            open: LruCache::new(cap),
        }
    }

    /// Derive pool size from the OS fd limit (75 %, clamped to `[16, MAX_POOL_SIZE]`).
    pub fn from_system_limits() -> Self {
        let fd_limit = getrlimit(Resource::Nofile).current.unwrap_or(256) as usize;
        let max_open = (fd_limit * 3 / 4).clamp(16, MAX_POOL_SIZE);
        Self::new(max_open)
    }

    pub fn max_open(&self) -> usize {
        self.max_open
    }

    /// Total number of registered entries (open + evicted).
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Number of currently open fds — O(1).
    pub fn open_count(&self) -> usize {
        self.open.len()
    }

    /// Flush and close all registered entries. O(n), one-time.
    ///
    /// Call only after all concurrent drains have completed (e.g., after joining
    /// all write threads). Tokens with unflushed pending must be flushed or closed
    /// before this call, or their data will be lost.
    pub fn close_all(&mut self) -> SKResult<()> {
        for entry in self.entries.iter_mut() {
            entry.logically_closed = true;
            let mut fd_guard = entry.fd.lock().unwrap();
            if let Some(mut w) = fd_guard.take() {
                w.flush()?;
                // drop(w) → Zstd frame footer written
            }
        }
        self.open.clear();
        Ok(())
    }

    // ── private ───────────────────────────────────────────────────────────────

    /// Create file on disk (empty, fd immediately closed), register entry.
    fn register(&mut self, path: PathBuf, format: Format, level: Level) -> SKResult<usize> {
        File::create(&path)?;
        let id = self.entries.len();
        self.entries.push(WriteEntry {
            path,
            format,
            level,
            fd: Arc::new(Mutex::new(None)),
            logically_closed: false,
        });
        Ok(id)
    }

    /// Ensure entry `id` has an open fd. Evicts LRU if at capacity.
    /// Must be called under pool lock.
    fn ensure_open(&mut self, id: usize) -> SKResult<()> {
        if self.open.get(&id).is_some() {
            return Ok(()); // promotes to MRU
        }
        if self.entries[id].logically_closed {
            return Err(std::io::Error::new(
                std::io::ErrorKind::BrokenPipe,
                "write to logically closed entry",
            )
            .into());
        }
        if self.open.len() >= self.max_open {
            self.evict_lru()?;
        }
        let entry = &self.entries[id];
        let file = OpenOptions::new().append(true).open(&entry.path)?;
        let writer =
            niffler::send::get_writer(Box::new(BufWriter::new(file)), entry.format, entry.level)?;
        // Brief fd lock acquisition under pool lock. fd is None here so uncontested.
        *self.entries[id].fd.lock().unwrap() = Some(writer);
        self.open.put(id, ());
        Ok(())
    }

    /// Evict the least recently used *non-busy* entry. O(1) typical; O(open) worst-case.
    ///
    /// Uses `try_lock()` so that entries currently being written to are skipped.
    /// Returns an error if all open fds are in use.
    fn evict_lru(&mut self) -> SKResult<()> {
        // Collect ids before the loop: iterating `self.open` borrows it immutably,
        // which would conflict with `self.open.pop()` inside the loop body.
        // iter() yields MRU→LRU; rev() gives LRU→MRU (try least recently used first).
        let candidates: Vec<usize> = self.open.iter().rev().map(|(&id, _)| id).collect();
        for lru_id in candidates {
            let fd_arc = Arc::clone(&self.entries[lru_id].fd);
            if let Ok(mut fd_guard) = fd_arc.try_lock() {
                if let Some(mut w) = fd_guard.take() {
                    let _ = w.flush(); // best-effort; drop finalizes Zstd frame
                }
                self.open.pop(&lru_id);
                return Ok(());
            }
            // fd locked by a writer thread — skip this candidate
        }
        Err(std::io::Error::new(
            std::io::ErrorKind::ResourceBusy,
            "pool saturated: all open fds are currently in use",
        )
        .into())
    }
}

// ── SKFileWriter ────────────────────────────────────────────────────────────────

/// Handle to a registered pool entry; owns a local pending write buffer.
///
/// # Write path (hot path)
///
/// `write(sk)` serialises into `pending` with **no locking**.
/// When `pending ≥ flush_threshold`, `drain()` is called:
///
/// 1. Pool lock acquired briefly → `ensure_open(id)` → entry fd lock acquired **under** pool lock.
/// 2. Pool lock released. Entry fd lock still held.
/// 3. Entire `pending` buffer written to fd in one `write_all`. Entry fd lock released.
///
/// Acquiring the fd lock while the pool lock is still held prevents a race where
/// another thread could evict the just-opened fd between releasing pool lock and
/// acquiring fd lock.  `evict_lru()` uses `try_lock()` and will skip an entry
/// whose fd lock is already held.
pub struct SKFileWriter {
    id: usize,
    pool: Arc<Mutex<SKFilePool>>,
    path: PathBuf,
    pending: Vec<u8>,
    flush_threshold: usize,
    logically_closed: bool,
    meta: SKFileMeta,
}

/// Create a `SKFileWriter` for a new file (Zstd, level 3).
pub fn create_token(pool: &SharedPool, path: PathBuf) -> SKResult<SKFileWriter> {
    create_token_with(pool, path, Format::Zstd, Level::Three)
}

/// Create a `SKFileWriter` for a new file with explicit format and level.
pub fn create_token_with(
    pool: &SharedPool,
    path: PathBuf,
    format: Format,
    level: Level,
) -> SKResult<SKFileWriter> {
    let id = pool.lock().unwrap().register(path.clone(), format, level)?;
    Ok(SKFileWriter {
        id,
        pool: Arc::clone(pool),
        path,
        pending: Vec::with_capacity(DEFAULT_FLUSH_THRESHOLD + 128),
        flush_threshold: DEFAULT_FLUSH_THRESHOLD,
        logically_closed: false,
        meta: SKFileMeta::zero(),
    })
}

impl SKFileWriter {
    /// Create a standalone file writer (Zstd, level 3).
    /// The pool is created internally and is not accessible to the caller.
    pub fn create<P: AsRef<Path>>(path: P) -> SKResult<Self> {
        Self::create_with(path, Format::Zstd, Level::Three)
    }

    /// Create a standalone file writer with explicit format and level.
    pub fn create_with<P: AsRef<Path>>(path: P, format: Format, level: Level) -> SKResult<Self> {
        create_token_with(global_pool(), path.as_ref().to_owned(), format, level)
    }

    /// `true` if the underlying fd is currently open in the pool.
    pub fn is_physically_open(&self) -> bool {
        self.pool.lock().unwrap().open.contains(&self.id)
    }

    /// Accumulate one SuperKmer. Drains to fd when `pending ≥ flush_threshold`.
    pub fn write(&mut self, sk: &SuperKmer) -> SKResult<()> {
        self.check_not_closed()?;
        self.write_one(sk)
    }

    /// Accumulate a slice of SuperKmers, draining whenever the threshold is exceeded.
    pub fn write_batch(&mut self, sks: &[SuperKmer]) -> SKResult<()> {
        self.check_not_closed()?;
        for sk in sks {
            self.write_one(sk)?;
        }
        Ok(())
    }

    fn write_one(&mut self, sk: &SuperKmer) -> SKResult<()> {
        sk.write_to_binary(&mut self.pending)?;
        self.meta.instances += 1;
        self.meta.count_sum += sk.count() as u64;
        self.meta.length_sum += sk.seql() as u64;
        if self.pending.len() >= self.flush_threshold {
            self.drain()?;
        }
        Ok(())
    }

    /// Drain pending bytes to the fd **and** flush the compressor's internal buffer.
    pub fn flush(&mut self) -> SKResult<()> {
        self.check_not_closed()?;
        if self.pending.is_empty() {
            return Ok(());
        }
        let fd_arc;
        let mut fd_guard;
        {
            let mut pool = self.pool.lock().unwrap();
            pool.ensure_open(self.id)?;
            fd_arc = Arc::clone(&pool.entries[self.id].fd);
            fd_guard = fd_arc.lock().unwrap(); // acquire fd lock under pool lock
            // pool drops here → pool lock released, fd lock still held
        }
        let w = fd_guard.as_mut().expect("fd open after ensure_open");
        w.write_all(&self.pending)?;
        w.flush()?;
        self.pending.clear();
        Ok(())
    }

    /// Flush pending bytes, finalize Zstd frame, permanently seal this token.
    pub fn close(&mut self) -> SKResult<()> {
        if self.logically_closed {
            return Ok(());
        }
        self.logically_closed = true;

        let fd_arc;
        let mut fd_guard;
        {
            let mut pool = self.pool.lock().unwrap();
            let has_pending = !self.pending.is_empty();
            if has_pending {
                pool.ensure_open(self.id)?;
            }
            pool.entries[self.id].logically_closed = true;
            pool.open.pop(&self.id);
            fd_arc = Arc::clone(&pool.entries[self.id].fd);
            fd_guard = fd_arc.lock().unwrap(); // acquire fd lock under pool lock
            // pool drops here → pool lock released
        }

        if !self.pending.is_empty() {
            fd_guard
                .as_mut()
                .expect("fd open after ensure_open")
                .write_all(&self.pending)?;
            self.pending.clear();
        }
        if let Some(mut w) = fd_guard.take() {
            w.flush()?;
            // drop(w) → Zstd frame finalized
        }
        self.meta.write(&self.path)?;
        Ok(())
    }

    /// Adjust the byte threshold at which pending is drained to the fd.
    /// Default: `DEFAULT_FLUSH_THRESHOLD` (8 KiB).
    pub fn set_flush_threshold(&mut self, bytes: usize) {
        self.flush_threshold = bytes;
    }

    /// `true` if this token has not been closed.
    pub fn is_open(&self) -> bool {
        !self.logically_closed
    }

    /// Physical path of the file.
    pub fn path(&self) -> &Path {
        &self.path
    }

    // ── private ───────────────────────────────────────────────────────────────

    fn check_not_closed(&self) -> SKResult<()> {
        if self.logically_closed {
            Err(std::io::Error::new(
                std::io::ErrorKind::BrokenPipe,
                "write to logically closed SKFileWriter",
            )
            .into())
        } else {
            Ok(())
        }
    }

    /// Drain pending bytes to the fd (no compressor flush).
    ///
    /// Two-phase locking:
    /// 1. Pool lock → ensure_open (promotes to MRU) → acquire entry fd lock (under pool lock).
    /// 2. Release pool lock. Write pending under entry fd lock only.
    ///
    /// Holding fd lock while releasing pool lock prevents eviction of our entry
    /// during the write: `evict_lru` uses `try_lock` and will skip us.
    fn drain(&mut self) -> SKResult<()> {
        debug_assert!(!self.pending.is_empty());
        let fd_arc;
        let mut fd_guard;
        {
            let mut pool = self.pool.lock().unwrap();
            pool.ensure_open(self.id)?;
            fd_arc = Arc::clone(&pool.entries[self.id].fd);
            fd_guard = fd_arc.lock().unwrap(); // acquire fd lock under pool lock
            // pool drops here → pool lock released, fd lock still held
        }
        fd_guard
            .as_mut()
            .expect("fd open after ensure_open")
            .write_all(&self.pending)?;
        // fd_guard drops → entry fd lock released
        self.pending.clear();
        Ok(())
    }
}

impl Drop for SKFileWriter {
    fn drop(&mut self) {
        if !self.logically_closed {
            let _ = self.close();
        }
    }
}

// ── tests ──────────────────────────────────────────────────────────────────────

#[cfg(test)]
#[path = "tests/pool.rs"]
mod tests;
