use crate::video::VideoInfo;
use std::collections::{HashMap, VecDeque};
use std::io::{BufRead, BufReader};
use std::path::PathBuf;
use std::process::{Child, ChildStderr, ChildStdout, Command, Stdio};
use std::sync::{Arc, Condvar, Mutex};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

pub const DEFAULT_CACHE_CAPACITY: usize = 160;
pub const MIN_CACHE_CAPACITY: usize = 50;
pub const MAX_CACHE_CAPACITY: usize = 800;
const READ_BUF_BYTES: usize = 256 * 1024;
/// If the requested frame is this far ahead of the worker, kill + reseek.
const SEEK_AHEAD_THRESHOLD: u64 = 60;
const DECODE_TIMEOUT: Duration = Duration::from_secs(15);
/// Keep the last N stderr lines from ffmpeg so error messages can include
/// the real reason (e.g. "vaapi: failed to load driver") instead of just
/// "pipe closed before frame completed".
const STDERR_TAIL_LINES: usize = 32;

struct FrameCache {
    map: HashMap<u64, Vec<u8>>,
    order: VecDeque<u64>,
    capacity: usize,
}

impl FrameCache {
    fn new(capacity: usize) -> Self {
        Self {
            map: HashMap::new(),
            order: VecDeque::new(),
            capacity,
        }
    }

    fn contains(&self, idx: u64) -> bool {
        self.map.contains_key(&idx)
    }

    fn get(&mut self, idx: u64) -> Option<Vec<u8>> {
        let bytes = self.map.get(&idx)?.clone();
        if let Some(pos) = self.order.iter().position(|&i| i == idx) {
            self.order.remove(pos);
        }
        self.order.push_back(idx);
        Some(bytes)
    }

    fn insert(&mut self, idx: u64, bytes: Vec<u8>) {
        if self.map.contains_key(&idx) {
            if let Some(pos) = self.order.iter().position(|&i| i == idx) {
                self.order.remove(pos);
            }
        } else if self.map.len() >= self.capacity {
            if let Some(oldest) = self.order.pop_front() {
                self.map.remove(&oldest);
            }
        }
        self.map.insert(idx, bytes);
        self.order.push_back(idx);
    }

    fn set_capacity(&mut self, new_cap: usize) {
        self.capacity = new_cap.clamp(MIN_CACHE_CAPACITY, MAX_CACHE_CAPACITY);
        while self.map.len() > self.capacity {
            if let Some(oldest) = self.order.pop_front() {
                self.map.remove(&oldest);
            } else {
                break;
            }
        }
    }
}

struct SharedState {
    /// Owned by shared state (not the worker closure) so stop_worker can
    /// kill ffmpeg directly and unblock a pipe read that's waiting for the
    /// first post-seek frame.
    child: Option<Child>,
    cache: FrameCache,
    /// Frame the worker's pipe will emit next.
    next_idx: u64,
    /// Most recent frame the frontend requested — used to bound lookahead.
    reader_idx: u64,
    /// Worker reads this many frames past reader_idx before parking.
    lookahead: u64,
    eof: bool,
    error: Option<String>,
    shutdown: bool,
}

struct Shared {
    state: Mutex<SharedState>,
    cond: Condvar,
    /// Ring buffer of recent ffmpeg stderr lines, filled by a sibling drain
    /// thread. Cleared at each new seek.
    stderr_tail: Mutex<VecDeque<String>>,
}

pub struct DecoderSession {
    path: PathBuf,
    info: VideoInfo,
    /// ffmpeg `-hwaccel` argument. None / empty = software decode. Cleared
    /// on any decode failure so the next attempt uses software — we'd
    /// rather degrade than keep retrying a flaky GPU path.
    hwaccel: Option<String>,
    shared: Arc<Shared>,
    worker: Option<JoinHandle<()>>,
    /// Sibling thread draining ffmpeg's stderr into `shared.stderr_tail`.
    /// Joined in `stop_worker` so the tail doesn't get clobbered by a
    /// straggling drain across seeks.
    stderr_drain: Option<JoinHandle<()>>,
}

impl DecoderSession {
    pub fn open(path: PathBuf, info: VideoInfo, hwaccel: Option<String>) -> Self {
        Self {
            path,
            info,
            hwaccel: hwaccel.filter(|s| !s.is_empty()),
            shared: Arc::new(Shared {
                state: Mutex::new(SharedState {
                    child: None,
                    cache: FrameCache::new(DEFAULT_CACHE_CAPACITY),
                    next_idx: 0,
                    reader_idx: 0,
                    lookahead: (DEFAULT_CACHE_CAPACITY as u64) / 2,
                    eof: false,
                    error: None,
                    shutdown: false,
                }),
                cond: Condvar::new(),
                stderr_tail: Mutex::new(VecDeque::new()),
            }),
            worker: None,
            stderr_drain: None,
        }
    }

    pub fn info(&self) -> &VideoInfo {
        &self.info
    }

    pub fn set_cache_capacity(&self, capacity: usize) {
        let clamped = capacity.clamp(MIN_CACHE_CAPACITY, MAX_CACHE_CAPACITY);
        let mut st = self.shared.state.lock().unwrap();
        st.cache.set_capacity(clamped);
        st.lookahead = (clamped as u64) / 2;
    }

    pub fn decode_at(&mut self, idx: u64) -> Result<Vec<u8>, String> {
        match self.decode_at_inner(idx) {
            Ok(bytes) => Ok(bytes),
            Err(e) => {
                // Auto-fallback: any decode failure while hwaccel is on
                // retries once with software. Covers both first-frame init
                // failures and later seek/codec-edge failures (e.g. a GPU
                // path that worked on h264 but chokes on an HEVC segment).
                // Once cleared, subsequent failures propagate.
                if let Some(prev) = self.hwaccel.take() {
                    eprintln!(
                        "[decoder] hwaccel '{prev}' failed ({e}); falling back to software"
                    );
                    self.stop_worker();
                    self.decode_at_inner(idx)
                } else {
                    Err(e)
                }
            }
        }
    }

    fn decode_at_inner(&mut self, idx: u64) -> Result<Vec<u8>, String> {
        if idx >= self.info.total_frames {
            return Err(format!(
                "frame {} out of range (total {})",
                idx, self.info.total_frames
            ));
        }

        // Update reader cursor, check cache, note whether a seek is needed.
        let need_seek = {
            let mut st = self.shared.state.lock().unwrap();
            st.reader_idx = idx;
            if let Some(bytes) = st.cache.get(idx) {
                self.shared.cond.notify_all();
                return Ok(bytes);
            }
            self.worker.is_none()
                || idx < st.next_idx
                || idx > st.next_idx.saturating_add(SEEK_AHEAD_THRESHOLD)
        };

        if need_seek {
            self.seek_and_start_worker(idx)?;
        } else {
            self.shared.cond.notify_all();
        }

        // Wait for the worker to produce idx.
        let wait_started = Instant::now();
        let mut st = self.shared.state.lock().unwrap();
        loop {
            if st.cache.contains(idx) {
                return st
                    .cache
                    .get(idx)
                    .ok_or_else(|| "frame vanished from cache".to_string());
            }
            if let Some(e) = &st.error {
                return Err(e.clone());
            }
            if st.eof {
                return Err(format!("frame {} not produced (EOF)", idx));
            }
            if st.shutdown {
                return Err("decoder shut down".into());
            }
            let (next_st, timeout) = self
                .shared
                .cond
                .wait_timeout(st, Duration::from_millis(250))
                .unwrap();
            st = next_st;
            if timeout.timed_out() && wait_started.elapsed() >= DECODE_TIMEOUT {
                drop(st);
                self.stop_worker();
                return Err(format!("timed out decoding frame {idx}"));
            }
        }
    }

    fn seek_and_start_worker(&mut self, idx: u64) -> Result<(), String> {
        self.stop_worker();

        let (child, stdout, stderr) =
            spawn_ffmpeg(&self.path, &self.info, idx, self.hwaccel.as_deref())?;

        {
            let mut st = self.shared.state.lock().unwrap();
            st.child = Some(child);
            st.next_idx = idx;
            st.reader_idx = idx;
            st.eof = false;
            st.error = None;
            st.shutdown = false;
        }
        self.shared.stderr_tail.lock().unwrap().clear();
        self.shared.cond.notify_all();

        // Drain ffmpeg's stderr into a ring buffer so error messages can
        // include the real reason. We must drain even on success — leaving
        // stderr unread would let the pipe buffer fill and stall ffmpeg.
        let drain_shared = self.shared.clone();
        self.stderr_drain = Some(thread::spawn(move || {
            let reader = BufReader::new(stderr);
            for line in reader.lines().map_while(Result::ok) {
                let mut tail = drain_shared.stderr_tail.lock().unwrap();
                if tail.len() >= STDERR_TAIL_LINES {
                    tail.pop_front();
                }
                tail.push_back(line);
            }
        }));

        let shared = self.shared.clone();
        let handle = thread::spawn(move || {
            worker_loop(stdout, shared);
        });
        self.worker = Some(handle);
        Ok(())
    }

    fn stop_worker(&mut self) {
        // Take ownership of the child out of shared state and kill it FIRST so
        // a pipe read in the worker unblocks with EOF. Then join the worker.
        let child_opt = {
            let mut st = self.shared.state.lock().unwrap();
            st.shutdown = true;
            st.child.take()
        };
        if let Some(mut child) = child_opt {
            let _ = child.kill();
            let _ = child.wait();
        }
        self.shared.cond.notify_all();

        if let Some(handle) = self.worker.take() {
            let _ = handle.join();
        }
        // Killing the child closed stderr, so the drain thread sees EOF and
        // exits — joining it here prevents a straggling write into the new
        // session's stderr_tail after the next seek clears it.
        if let Some(handle) = self.stderr_drain.take() {
            let _ = handle.join();
        }

        let mut st = self.shared.state.lock().unwrap();
        st.shutdown = false;
        st.eof = false;
        st.error = None;
    }
}

impl Drop for DecoderSession {
    fn drop(&mut self) {
        self.stop_worker();
    }
}

/// Pick a VAAPI render node. `SAM_TOOL_VAAPI_DEVICE` overrides; otherwise we
/// probe `/dev/dri/renderD128..199` for the first that exists. Multi-GPU
/// systems can end up with the integrated card on D128 and a discrete card
/// on D129+ — the env var is the escape hatch when the auto pick is wrong.
fn vaapi_device() -> String {
    if let Ok(p) = std::env::var("SAM_TOOL_VAAPI_DEVICE") {
        if !p.is_empty() {
            return p;
        }
    }
    for n in 128..200u32 {
        let p = format!("/dev/dri/renderD{n}");
        if std::path::Path::new(&p).exists() {
            return p;
        }
    }
    "/dev/dri/renderD128".to_string()
}

fn spawn_ffmpeg(
    path: &PathBuf,
    info: &VideoInfo,
    idx: u64,
    hwaccel: Option<&str>,
) -> Result<(Child, BufReader<ChildStdout>, ChildStderr), String> {
    let path_str = path.to_str().ok_or_else(|| "non-utf8 path".to_string())?;
    let t = idx as f64 / info.fps;
    let mut cmd = Command::new("ffmpeg");
    cmd.args(["-v", "error", "-threads", "0"]);
    // -hwaccel must come BEFORE -i. Output goes to CPU memory by default
    // (no -hwaccel_output_format), so the downstream mjpeg encoder works
    // unchanged. If the GPU init fails, ffmpeg will exit nonzero and the
    // worker thread surfaces the error via shared.error.
    match hwaccel {
        Some("auto") => {
            cmd.args(["-hwaccel", "auto"]);
        }
        Some("vaapi") => {
            let dev = vaapi_device();
            cmd.args(["-hwaccel", "vaapi", "-vaapi_device", &dev]);
        }
        Some("cuda") => {
            cmd.args(["-hwaccel", "cuda"]);
        }
        Some("vdpau") => {
            cmd.args(["-hwaccel", "vdpau"]);
        }
        _ => {}
    }
    if t > 1.0 {
        cmd.args(["-ss", &format!("{:.3}", t - 1.0)]);
        cmd.args(["-i", path_str]);
        cmd.args(["-ss", "1.0"]);
    } else {
        cmd.args(["-i", path_str]);
        cmd.args(["-ss", &format!("{:.3}", t)]);
    }
    cmd.args([
        "-an",
        "-q:v",
        "3",
        "-f",
        "image2pipe",
        "-vcodec",
        "mjpeg",
        "pipe:1",
    ]);
    cmd.stdout(Stdio::piped()).stderr(Stdio::piped());
    let mut child = cmd.spawn().map_err(|e| format!("ffmpeg spawn: {e}"))?;
    let stdout = child
        .stdout
        .take()
        .ok_or_else(|| "ffmpeg stdout unavailable".to_string())?;
    let stderr = child
        .stderr
        .take()
        .ok_or_else(|| "ffmpeg stderr unavailable".to_string())?;
    Ok((
        child,
        BufReader::with_capacity(READ_BUF_BYTES, stdout),
        stderr,
    ))
}

fn worker_loop(mut stdout: BufReader<ChildStdout>, shared: Arc<Shared>) {
    loop {
        // Backpressure: wait if we're already far enough ahead.
        {
            let st = shared.state.lock().unwrap();
            if st.shutdown {
                return;
            }
            let ahead = st.next_idx.saturating_sub(st.reader_idx);
            if ahead >= st.lookahead {
                let (st2, _) = shared
                    .cond
                    .wait_timeout(st, Duration::from_millis(100))
                    .unwrap();
                if st2.shutdown {
                    return;
                }
                continue;
            }
        }

        match read_jpeg_bytes(&mut stdout) {
            Ok(bytes) => {
                let mut st = shared.state.lock().unwrap();
                if st.shutdown {
                    return;
                }
                let idx = st.next_idx;
                st.cache.insert(idx, bytes);
                st.next_idx = idx + 1;
                shared.cond.notify_all();
            }
            Err(e) => {
                // ffmpeg closes stderr right alongside stdout; on real
                // failures the drain thread may not have flushed yet, so
                // we wait briefly. On a clean exit there's nothing to
                // flush — collect_stderr_tail short-circuits via the
                // child's exit code so end-of-stream EOF isn't penalised.
                let tail = collect_stderr_tail(&shared, Duration::from_millis(150));
                let mut st = shared.state.lock().unwrap();
                if !tail.is_empty() {
                    // With `-v error`, anything on stderr means ffmpeg
                    // actually failed (vs a clean exit) — surface it as a
                    // real error so decode_at can fall back / report.
                    st.error = Some(format!("{e}: ffmpeg: {tail}"));
                } else if e.contains("pipe closed") || e.contains("EOF") {
                    st.eof = true;
                } else {
                    st.error = Some(e);
                }
                shared.cond.notify_all();
                return;
            }
        }
    }
}

fn read_tail_now(shared: &Arc<Shared>) -> Option<String> {
    let tail = shared.stderr_tail.lock().unwrap();
    if tail.is_empty() {
        None
    } else {
        Some(tail.iter().cloned().collect::<Vec<_>>().join(" | "))
    }
}

fn collect_stderr_tail(shared: &Arc<Shared>, max_wait: Duration) -> String {
    // Fast path: drain already flushed — common when ffmpeg writes its
    // error before exiting and stdout EOF arrives after the drain.
    if let Some(tail) = read_tail_now(shared) {
        return tail;
    }
    // Skip the poll on clean ffmpeg exit (exit code 0): natural end-of-
    // stream has nothing on stderr (we run with `-v error`), and 150ms ×
    // every EOF seek near the final frames adds up.
    let exited_cleanly = {
        let mut st = shared.state.lock().unwrap();
        st.child
            .as_mut()
            .and_then(|c| c.try_wait().ok().flatten())
            .map(|status| status.success())
            .unwrap_or(false)
    };
    if exited_cleanly {
        return String::new();
    }
    // Slow path: ffmpeg exited non-zero (or hasn't exited yet) and the
    // drain thread hasn't flushed. Poll briefly so the error message
    // carries the real reason.
    let deadline = Instant::now() + max_wait;
    loop {
        if let Some(tail) = read_tail_now(shared) {
            return tail;
        }
        if Instant::now() >= deadline {
            return String::new();
        }
        thread::sleep(Duration::from_millis(10));
    }
}

fn read_jpeg_bytes(reader: &mut BufReader<ChildStdout>) -> Result<Vec<u8>, String> {
    let mut out = Vec::with_capacity(64 * 1024);
    let mut last = 0u8;
    let mut in_jpeg = false;
    loop {
        let buf = reader.fill_buf().map_err(|e| format!("pipe read: {e}"))?;
        if buf.is_empty() {
            return Err("ffmpeg pipe closed before frame completed".into());
        }
        let mut consumed = 0usize;
        for &b in buf {
            consumed += 1;
            if !in_jpeg {
                if last == 0xFF && b == 0xD8 {
                    in_jpeg = true;
                    out.push(0xFF);
                    out.push(0xD8);
                }
                last = b;
            } else {
                out.push(b);
                if last == 0xFF && b == 0xD9 {
                    reader.consume(consumed);
                    return Ok(out);
                }
                last = b;
            }
        }
        reader.consume(consumed);
    }
}