cortiq-gateway/src/media.rs

//! Media generation: run `cortiq imagine` (text→image, Lumina) and
//! `cortiq animate` (text→video+audio, MiniMax-H3) over local `.cmf` files as
//! tracked background jobs, the same shape as import jobs so the admin UI can
//! poll progress and fetch the result.
//!
//! These models are NOT served by `cortiq serve` — each generation is one CLI
//! run. The gateway owns the process, parses its progress, and keeps the output
//! under `<data>/media/<job>/`.

use serde::Serialize;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::{Arc, Mutex};
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::sync::Notify;

fn now() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

/// Random job id — media jobs are never deduplicated (every run is a new work).
fn gen_media_id() -> String {
    let mut b = [0u8; 8];
    let _ = getrandom::getrandom(&mut b);
    b.iter().map(|x| format!("{x:02x}")).collect()
}

/// What a `.cmf` under models_dir can do, decided by `cortiq info`.
#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize)]
#[serde(rename_all = "lowercase")]
pub enum Modality {
    Chat,
    Image,
    Video,
}

/// `cortiq info` output → modality. The arch line names the family; the media
/// packs also carry their pipeline components (DiT/VAE/vocoder) in the text.
fn classify_info(info: &str) -> Modality {
    let lower = info.to_ascii_lowercase();
    let arch = lower
        .lines()
        .find_map(|l| l.trim().strip_prefix("arch:"))
        .unwrap_or("")
        .trim()
        .to_string();
    // generic arch suffixes first (e.g. "lumina2-image", "mmh3-video"), then
    // known family names as a fallback for packs without the suffix
    if arch.contains("video") || arch.contains("mmh3") || arch.contains("minimax") {
        return Modality::Video;
    }
    if arch.contains("image") || arch.contains("lumina") || arch.contains("diffusion") {
        return Modality::Image;
    }
    if lower.contains("text-to-video") {
        return Modality::Video;
    }
    if lower.contains("text-to-image") {
        return Modality::Image;
    }
    Modality::Chat
}

/// One local `.cmf` with its detected modality.
#[derive(Clone, Serialize)]
pub struct MediaModel {
    pub name: String, // file stem
    pub path: String,
    pub size_bytes: u64,
    pub modality: Modality,
}

#[derive(Default)]
struct InfoCache(Mutex<HashMap<String, (u64, u64, Modality)>>); // path → (mtime, size, modality)

/// Modality of one `.cmf`, via `cortiq info` — cached by (path, mtime, size),
/// so the process runs once per file version.
pub async fn classify_file(path: &Path, cortiq_bin: &str) -> Modality {
    static CACHE: std::sync::OnceLock<InfoCache> = std::sync::OnceLock::new();
    let cache = CACHE.get_or_init(InfoCache::default);

    let size = crate::import::path_size(path);
    let mtime = std::fs::metadata(path)
        .ok()
        .and_then(|m| m.modified().ok())
        .and_then(|t| t.duration_since(UNIX_EPOCH).ok())
        .map(|d| d.as_secs())
        .unwrap_or(0);
    let key = path.to_string_lossy().to_string();
    if let Some(k) = cache
        .0
        .lock()
        .unwrap()
        .get(&key)
        .filter(|(m, s, _)| *m == mtime && *s == size)
        .map(|(_, _, k)| *k)
    {
        return k;
    }
    let output = tokio::process::Command::new(cortiq_bin)
        .arg("info")
        .arg(path)
        .output()
        .await;
    let text = output
        .map(|o| {
            format!(
                "{}{}",
                String::from_utf8_lossy(&o.stdout),
                String::from_utf8_lossy(&o.stderr)
            )
        })
        .unwrap_or_default();
    let k = classify_info(&text);
    cache.0.lock().unwrap().insert(key, (mtime, size, k));
    k
}

/// Scan models_dir and classify every `.cmf` (file or sharded dir).
pub async fn list_media_models(models_dir: &str, cortiq_bin: &str) -> Vec<MediaModel> {
    let mut out = Vec::new();
    let Ok(rd) = std::fs::read_dir(models_dir) else {
        return out;
    };
    for entry in rd.filter_map(|e| e.ok()) {
        let path = entry.path();
        let name = path
            .file_name()
            .unwrap_or_default()
            .to_string_lossy()
            .to_string();
        if !name.to_ascii_lowercase().ends_with(".cmf") {
            continue;
        }
        let size = crate::import::path_size(&path);
        if size == 0 {
            continue; // empty or still being written
        }
        let modality = classify_file(&path, cortiq_bin).await;
        out.push(MediaModel {
            name: name
                .trim_end_matches(".cmf")
                .trim_end_matches(".CMF")
                .to_string(),
            path: path.to_string_lossy().to_string(),
            size_bytes: size,
            modality,
        });
    }
    out.sort_by(|a, b| a.name.cmp(&b.name));
    out
}

/// Parameters for one generation, straight from the playground form.
#[derive(Clone, Debug, serde::Deserialize)]
pub struct MediaParams {
    pub model: String, // file stem under models_dir
    pub kind: String,  // image | video
    pub prompt: String,
    #[serde(default)]
    pub width: Option<u32>,
    #[serde(default)]
    pub height: Option<u32>,
    #[serde(default)]
    pub steps: Option<u32>,
    #[serde(default)]
    pub cfg: Option<f32>, // image only
    #[serde(default)]
    pub seed: Option<u64>,
    #[serde(default)]
    pub frames: Option<u32>, // video only
    #[serde(default)]
    pub quality: Option<u32>, // video only: AVI JPEG quality
    /// Optional first/last keyframe for video, base64-encoded binary P6 PPM.
    #[serde(default)]
    pub first_frame_b64: Option<String>,
    #[serde(default)]
    pub last_frame_b64: Option<String>,
}

/// Live state of one generation.
#[derive(Clone, Serialize, serde::Deserialize)]
pub struct MediaJob {
    pub id: String,
    pub kind: String, // image | video
    pub model: String,
    pub prompt: String,
    pub width: u32,
    pub height: u32,
    pub steps: u32,
    pub seed: u64,
    pub frames: Option<u32>,
    pub state: String, // queued | running | done | error | cancelled
    pub progress: Option<f32>,
    pub phase: Option<String>,
    pub log: Vec<String>,
    pub started: u64,
    pub finished: Option<u64>,
    /// Files this job produced: name → relative kind ("image", "video", "audio", "mp4").
    pub outputs: Vec<String>,
    pub out_bytes: Option<u64>,
}

/// Everything needed to actually launch a queued job when its turn comes.
struct Prepared {
    bin: String,
    args: Vec<String>,
    gpu: bool,
    dir: PathBuf,
    out_path: PathBuf,
    kind: String,
}

#[derive(Default)]
pub struct MediaStore {
    jobs: Mutex<HashMap<String, MediaJob>>,
    cancels: Mutex<HashMap<String, Arc<Notify>>>,
    /// One generation runs at a time — these models take the whole GPU/RAM
    /// budget — but new requests queue up instead of being rejected.
    queue: Mutex<std::collections::VecDeque<String>>,
    current: Mutex<Option<String>>,
    prepared: Mutex<HashMap<String, Prepared>>,
    persist: Mutex<Option<PathBuf>>,
}

impl MediaStore {
    pub fn new() -> Arc<Self> {
        Arc::new(Self::default())
    }
    /// Load saved generations (their files are still under media_dir) and start
    /// persisting. Jobs that were queued/running when the gateway died settle
    /// as errors — the process is gone.
    pub fn attach_persistence(&self, path: PathBuf) {
        if let Ok(text) = std::fs::read_to_string(&path) {
            if let Ok(mut jobs) = serde_json::from_str::<Vec<MediaJob>>(&text) {
                let mut g = self.jobs.lock().unwrap();
                for mut j in jobs.drain(..) {
                    if j.state == "running" || j.state == "queued" {
                        j.state = "error".into();
                        j.log.push("✗ interrupted by a gateway restart".into());
                        j.finished = Some(now());
                    }
                    g.insert(j.id.clone(), j);
                }
            }
        }
        *self.persist.lock().unwrap() = Some(path);
    }
    fn save(&self) {
        let Some(path) = self.persist.lock().unwrap().clone() else {
            return;
        };
        let jobs: Vec<MediaJob> = self.jobs.lock().unwrap().values().cloned().collect();
        if let Ok(json) = serde_json::to_vec(&jobs) {
            let tmp = path.with_extension("json.tmp");
            if std::fs::write(&tmp, json).is_ok() {
                let _ = std::fs::rename(&tmp, &path);
            }
        }
    }
    /// How many jobs sit in the queue ahead of `id` (0 = next up).
    pub fn queue_position(&self, id: &str) -> Option<usize> {
        self.queue.lock().unwrap().iter().position(|x| x == id)
    }
    pub fn list(&self) -> Vec<MediaJob> {
        let mut v: Vec<MediaJob> = self.jobs.lock().unwrap().values().cloned().collect();
        v.sort_by_key(|j| std::cmp::Reverse(j.started));
        v.truncate(20);
        v
    }
    pub fn get(&self, id: &str) -> Option<MediaJob> {
        self.jobs.lock().unwrap().get(id).cloned()
    }
    pub fn cancel(self: &Arc<Self>, id: &str, media_dir: &Path) -> bool {
        let state = self
            .jobs
            .lock()
            .unwrap()
            .get(id)
            .map(|j| j.state.clone())
            .unwrap_or_default();
        match state.as_str() {
            "queued" => {
                self.queue.lock().unwrap().retain(|x| x != id);
                self.prepared.lock().unwrap().remove(id);
                let _ = std::fs::remove_dir_all(media_dir.join(id));
                self.set_cancelled(id);
                true
            }
            "running" => {
                if let Some(n) = self.cancels.lock().unwrap().get(id) {
                    n.notify_one();
                    true
                } else {
                    false
                }
            }
            _ => false,
        }
    }
    /// Remove a settled job and its output directory. Err = still running.
    pub fn delete(&self, id: &str, media_dir: &Path) -> Result<bool, String> {
        let job = self.jobs.lock().unwrap().get(id).cloned();
        let Some(job) = job else { return Ok(false) };
        if job.state == "running" || job.state == "queued" {
            return Err("cancel the running generation first".into());
        }
        let _ = std::fs::remove_dir_all(media_dir.join(id));
        self.jobs.lock().unwrap().remove(id);
        self.save();
        Ok(true)
    }
    fn push_line(&self, id: &str, line: String) {
        if let Some(j) = self.jobs.lock().unwrap().get_mut(id) {
            j.log.push(line);
            let n = j.log.len();
            if n > 120 {
                j.log.drain(0..n - 120);
            }
        }
    }
    fn set_progress(&self, id: &str, frac: f32, phase: Option<String>) {
        if let Some(j) = self.jobs.lock().unwrap().get_mut(id) {
            let f = frac.clamp(0.0, 1.0);
            if j.progress.is_none() || f > j.progress.unwrap() {
                j.progress = Some(f);
            }
            if let Some(p) = phase {
                j.phase = Some(p);
            }
        }
    }
    fn finish(self: &Arc<Self>, id: &str, ok: bool, outputs: Vec<String>, out_bytes: u64) {
        let mut g = self.jobs.lock().unwrap();
        if let Some(j) = g.get_mut(id) {
            if j.state == "cancelled" {
                drop(g);
                self.after_slot_freed(id);
                return;
            }
            j.finished = Some(now());
            j.state = if ok { "done" } else { "error" }.into();
            if ok {
                j.progress = Some(1.0);
                j.outputs = outputs;
                j.out_bytes = Some(out_bytes);
            }
        }
        drop(g);
        self.after_slot_freed(id);
    }
    fn set_cancelled(self: &Arc<Self>, id: &str) {
        if let Some(j) = self.jobs.lock().unwrap().get_mut(id) {
            j.state = "cancelled".into();
            j.finished = Some(now());
        }
        self.after_slot_freed(id);
    }
    fn after_slot_freed(self: &Arc<Self>, id: &str) {
        {
            let mut cur = self.current.lock().unwrap();
            if cur.as_deref() == Some(id) {
                *cur = None;
            }
        }
        self.cancels.lock().unwrap().remove(id);
        self.save();
        promote_next(self.clone());
    }
}

/// Start the next queued generation, if the slot is free.
fn promote_next(store: Arc<MediaStore>) {
    let next = {
        let mut cur = store.current.lock().unwrap();
        if cur.is_some() {
            return;
        }
        let Some(id) = store.queue.lock().unwrap().pop_front() else {
            return;
        };
        *cur = Some(id.clone());
        id
    };
    let Some(p) = store.prepared.lock().unwrap().remove(&next) else {
        // job was cancelled while queued and slipped through — free the slot
        *store.current.lock().unwrap() = None;
        return;
    };
    if let Some(j) = store.jobs.lock().unwrap().get_mut(&next) {
        j.state = "running".into();
        j.phase = Some("starting".into());
        j.started = now();
    }
    store.save();
    let cancel = Arc::new(Notify::new());
    store
        .cancels
        .lock()
        .unwrap()
        .insert(next.clone(), cancel.clone());
    let s = store.clone();
    tokio::spawn(async move {
        run_generation(
            s, next, p.bin, p.args, p.gpu, p.dir, p.out_path, p.kind, cancel,
        )
        .await;
    });
}

/// Progress markers in `cortiq imagine`/`animate` stdout. The exact wording has
/// varied between releases, so accept anything that carries an X/Y pair near a
/// step/frame word, a bare percentage, or the converter-style `@PROGRESS`.
fn parse_media_progress(line: &str) -> Option<(f32, Option<String>)> {
    let l = line.trim();
    if let Some(rest) = l.strip_prefix("@PROGRESS ") {
        let mut it = rest.splitn(2, ' ');
        let frac: f32 = it.next()?.trim().parse().ok()?;
        let phase = it
            .next()
            .map(|s| s.trim().to_string())
            .filter(|s| !s.is_empty());
        return Some((frac.clamp(0.0, 1.0), phase));
    }
    let lower = l.to_ascii_lowercase();
    // "step 12/30", "denoise 3/4", "frame 7/39" — any word followed by x/y
    for word in ["step", "frame", "sigma", "denois"] {
        if let Some(pos) = lower.find(word) {
            let tail = &lower[pos..];
            if let Some((x, y)) = extract_ratio(tail) {
                return Some((x / y, Some(l.to_string())));
            }
        }
    }
    // bare "42%" (avoid matching sizes like "42%faster")
    if let Some(idx) = lower.find('%') {
        let head: String = lower[..idx]
            .chars()
            .rev()
            .take_while(|c| c.is_ascii_digit() || *c == '.')
            .collect();
        let num: String = head.chars().rev().collect();
        if let Ok(p) = num.parse::<f32>() {
            if (0.0..=100.0).contains(&p) {
                return Some((p / 100.0, None));
            }
        }
    }
    None
}

/// First "N/M" pair in the string with M > 0.
fn extract_ratio(s: &str) -> Option<(f32, f32)> {
    let bytes = s.as_bytes();
    let mut i = 0;
    while i < bytes.len() {
        if bytes[i].is_ascii_digit() {
            let start = i;
            while i < bytes.len() && bytes[i].is_ascii_digit() {
                i += 1;
            }
            if i < bytes.len() && bytes[i] == b'/' {
                let x: f32 = s[start..i].parse().ok()?;
                let ystart = i + 1;
                let mut j = ystart;
                while j < bytes.len() && bytes[j].is_ascii_digit() {
                    j += 1;
                }
                if j > ystart {
                    if let Ok(y) = s[ystart..j].parse::<f32>() {
                        if y > 0.0 && x <= y {
                            return Some((x, y));
                        }
                    }
                }
            }
        } else {
            i += 1;
        }
    }
    None
}

/// Root for generated files: `<data>/media`.
pub fn media_dir() -> PathBuf {
    crate::config::data_dir().join("media")
}

const B64: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
/// Minimal base64 decode (standard alphabet, padding optional) — the keyframes
/// are the only base64 input the gateway takes; not worth a dependency.
fn b64_decode(s: &str) -> Option<Vec<u8>> {
    let mut lut = [255u8; 256];
    for (i, c) in B64.iter().enumerate() {
        lut[*c as usize] = i as u8;
    }
    let mut out = Vec::with_capacity(s.len() * 3 / 4);
    let mut buf = 0u32;
    let mut bits = 0u32;
    for c in s.bytes() {
        if c == b'=' || c == b'\n' || c == b'\r' {
            continue;
        }
        let v = lut[c as usize];
        if v == 255 {
            return None;
        }
        buf = (buf << 6) | v as u32;
        bits += 6;
        if bits >= 8 {
            bits -= 8;
            out.push((buf >> bits) as u8);
        }
    }
    Some(out)
}

/// Kick off a generation. Returns the job id, or an error when another one is
/// already running / the model file is missing.
pub fn start_generation(
    store: Arc<MediaStore>,
    cmf: &crate::config::CmfCfg,
    p: MediaParams,
) -> Result<String, String> {
    if p.prompt.trim().is_empty() {
        return Err("empty prompt".into());
    }
    let kind = p.kind.as_str();
    if kind != "image" && kind != "video" {
        return Err(format!("unknown kind '{kind}'"));
    }
    let model_path = Path::new(&cmf.models_dir).join(format!("{}.cmf", p.model));
    if !model_path.exists() {
        return Err(format!("model file not found: {}", model_path.display()));
    }

    let id = gen_media_id();
    let dir = media_dir().join(&id);
    if let Err(e) = std::fs::create_dir_all(&dir) {
        return Err(format!("cannot create media dir {}: {e}", dir.display()));
    }

    // Optional keyframes (video): decode to PPM files next to the output.
    let mut first_frame: Option<PathBuf> = None;
    let mut last_frame: Option<PathBuf> = None;
    for (b64, slot, name) in [
        (&p.first_frame_b64, &mut first_frame, "first.ppm"),
        (&p.last_frame_b64, &mut last_frame, "last.ppm"),
    ] {
        if let Some(data) = b64.as_deref().filter(|s| !s.is_empty()) {
            let bytes = b64_decode(data).ok_or("invalid keyframe base64")?;
            if !bytes.starts_with(b"P6") {
                return Err("keyframe must be a binary P6 PPM".into());
            }
            let path = dir.join(name);
            std::fs::write(&path, bytes).map_err(|e| format!("write keyframe: {e}"))?;
            *slot = Some(path);
        }
    }

    let width = p.width.unwrap_or(512).clamp(64, 2048);
    let height = p
        .height
        .unwrap_or(if kind == "video" { 288 } else { 512 })
        .clamp(64, 2048);
    let steps = p
        .steps
        .unwrap_or(if kind == "video" { 4 } else { 30 })
        .clamp(1, 200);
    let seed = p.seed.unwrap_or(42);
    let frames = p.frames.unwrap_or(39).clamp(1, 1000);

    let out_name = if kind == "image" {
        "out.ppm"
    } else {
        "out.avi"
    };
    let out_path = dir.join(out_name);

    let mut args: Vec<String> = if kind == "image" {
        vec![
            "imagine".into(),
            model_path.to_string_lossy().to_string(),
            "--prompt".into(),
            p.prompt.clone(),
            "--width".into(),
            width.to_string(),
            "--height".into(),
            height.to_string(),
            "--steps".into(),
            steps.to_string(),
            "--cfg".into(),
            format!("{}", p.cfg.unwrap_or(4.0)),
            "--seed".into(),
            seed.to_string(),
            "--out".into(),
            out_path.to_string_lossy().to_string(),
        ]
    } else {
        let mut a = vec![
            "animate".into(),
            model_path.to_string_lossy().to_string(),
            "--prompt".into(),
            p.prompt.clone(),
            "--width".into(),
            width.to_string(),
            "--height".into(),
            height.to_string(),
            "--frames".into(),
            frames.to_string(),
            "--steps".into(),
            steps.to_string(),
            "--seed".into(),
            seed.to_string(),
            "--quality".into(),
            p.quality.unwrap_or(92).clamp(10, 100).to_string(),
            "--out".into(),
            out_path.to_string_lossy().to_string(),
        ];
        if let Some(f) = &first_frame {
            a.push("--first-frame".into());
            a.push(f.to_string_lossy().to_string());
        }
        if let Some(f) = &last_frame {
            a.push("--last-frame".into());
            a.push(f.to_string_lossy().to_string());
        }
        a
    };
    // keep the model path first for readability in logs
    let _ = &mut args;

    store.jobs.lock().unwrap().insert(
        id.clone(),
        MediaJob {
            id: id.clone(),
            kind: kind.to_string(),
            model: p.model.clone(),
            prompt: p.prompt.clone(),
            width,
            height,
            steps,
            seed,
            frames: (kind == "video").then_some(frames),
            state: "queued".into(),
            progress: None,
            phase: None,
            log: vec![format!("→ cortiq {}", args.join(" "))],
            started: now(),
            finished: None,
            outputs: Vec::new(),
            out_bytes: None,
        },
    );
    store.prepared.lock().unwrap().insert(
        id.clone(),
        Prepared {
            bin: cmf.cortiq_bin.clone(),
            args,
            gpu: cmf.gpu,
            dir,
            out_path,
            kind: kind.to_string(),
        },
    );
    store.queue.lock().unwrap().push_back(id.clone());
    store.save();
    promote_next(store.clone());
    Ok(id)
}

#[allow(clippy::too_many_arguments)]
async fn run_generation(
    store: Arc<MediaStore>,
    id: String,
    bin: String,
    args: Vec<String>,
    gpu: bool,
    dir: PathBuf,
    out_path: PathBuf,
    kind: String,
    cancel: Arc<Notify>,
) {
    let mut cmd = tokio::process::Command::new(&bin);
    cmd.args(&args)
        .env("CMF_GPU", if gpu { "1" } else { "0" })
        .stdout(Stdio::piped())
        .stderr(Stdio::piped());
    let mut child = match cmd.spawn() {
        Ok(c) => c,
        Err(e) => {
            store.push_line(&id, format!("✗ failed to start cortiq: {e}"));
            store.finish(&id, false, Vec::new(), 0);
            return;
        }
    };
    let mut readers = Vec::new();
    if let Some(o) = child.stdout.take() {
        readers.push(tokio::spawn(read_lines(store.clone(), id.clone(), o)));
    }
    if let Some(e) = child.stderr.take() {
        readers.push(tokio::spawn(read_lines(store.clone(), id.clone(), e)));
    }

    let status = tokio::select! {
        s = child.wait() => s,
        _ = cancel.notified() => {
            let _ = child.start_kill();
            let _ = child.wait().await;
            for r in readers { r.abort(); }
            let _ = std::fs::remove_dir_all(&dir);
            store.push_line(&id, "✗ cancelled".into());
            store.set_cancelled(&id);
            return;
        }
    };
    for r in readers {
        let _ = r.await;
    }

    let ok = status.map(|s| s.success()).unwrap_or(false);
    if !ok || !out_path.exists() {
        store.push_line(&id, "✗ generation failed — see the log above".into());
        store.finish(&id, false, Vec::new(), 0);
        return;
    }

    let mut outputs = vec![if kind == "image" { "image" } else { "video" }.to_string()];
    if kind == "video" {
        // `animate` writes a .wav next to the .avi
        if dir.join("out.wav").exists() {
            outputs.push("audio".into());
        }
        // opportunistic browser-friendly remux when ffmpeg is present
        if let Ok(p) = which_ffmpeg() {
            store.push_line(&id, "→ remuxing to mp4 (ffmpeg found)".into());
            let mp4 = dir.join("out.mp4");
            let mut f = tokio::process::Command::new(p);
            f.arg("-y").arg("-i").arg(dir.join("out.avi"));
            if dir.join("out.wav").exists() {
                f.arg("-i")
                    .arg(dir.join("out.wav"))
                    .args(["-map", "0:v", "-map", "1:a", "-c:a", "aac"]);
            }
            f.args(["-c:v", "libx264", "-pix_fmt", "yuv420p"]).arg(&mp4);
            f.stdout(Stdio::null()).stderr(Stdio::null());
            if matches!(f.status().await, Ok(s) if s.success()) && mp4.exists() {
                outputs.push("mp4".into());
            }
        }
    }
    let bytes = crate::import::path_size(&out_path);
    store.push_line(
        &id,
        format!("✓ done — {}", crate::import::format_bytes(bytes)),
    );
    store.finish(&id, true, outputs, bytes);
}

/// Strip ANSI color escapes that `tracing` writes even into a pipe.
fn strip_ansi(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut chars = s.chars();
    while let Some(c) = chars.next() {
        if c == '\u{1b}' {
            for d in chars.by_ref() {
                if d.is_ascii_alphabetic() {
                    break;
                }
            }
        } else {
            out.push(c);
        }
    }
    out
}

async fn read_lines(store: Arc<MediaStore>, id: String, stream: impl tokio::io::AsyncRead + Unpin) {
    let mut lines = BufReader::new(stream).lines();
    while let Ok(Some(l)) = lines.next_line().await {
        let l = strip_ansi(l.trim());
        if l.is_empty() {
            continue;
        }
        // the Metal backend logs one line per pipeline — dozens per run, all identical
        if l.contains("Metal GPU path") {
            continue;
        }
        if let Some((frac, phase)) = parse_media_progress(&l) {
            store.set_progress(&id, frac, phase);
        } else {
            store.push_line(&id, l);
        }
    }
}

fn which_ffmpeg() -> Result<String, ()> {
    for p in [
        "ffmpeg",
        "/usr/local/bin/ffmpeg",
        "/opt/homebrew/bin/ffmpeg",
    ] {
        if std::process::Command::new(p)
            .arg("-version")
            .stdout(Stdio::null())
            .stderr(Stdio::null())
            .status()
            .map(|s| s.success())
            .unwrap_or(false)
        {
            return Ok(p.to_string());
        }
    }
    Err(())
}

/// (relative file, content-type) for an output kind of a finished job.
pub fn output_file(kind: &str) -> Option<(&'static str, &'static str)> {
    match kind {
        "image" => Some(("out.ppm", "image/x-portable-pixmap")),
        "video" => Some(("out.avi", "video/x-msvideo")),
        "audio" => Some(("out.wav", "audio/wav")),
        "mp4" => Some(("out.mp4", "video/mp4")),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn progress_parses_common_shapes() {
        assert_eq!(parse_media_progress("step 15/30").unwrap().0, 0.5);
        assert_eq!(
            parse_media_progress("Denoising step 3/4 ...").unwrap().0,
            0.75
        );
        assert_eq!(parse_media_progress("frame 39/39 decoded").unwrap().0, 1.0);
        assert_eq!(parse_media_progress("@PROGRESS 0.25 vae").unwrap().0, 0.25);
        assert!(parse_media_progress("loaded 2361 tensors").is_none());
    }

    #[test]
    fn classify_by_arch() {
        assert_eq!(classify_info("  Arch:        lumina2\n"), Modality::Image);
        assert_eq!(classify_info("  Arch:        mmh3\n"), Modality::Video);
        assert_eq!(classify_info("  Arch:        qwen3\n"), Modality::Chat);
    }

    #[test]
    fn b64_roundtrip() {
        assert_eq!(b64_decode("UDYgd2g=").unwrap(), b"P6 wh");
        assert!(b64_decode("!!").is_none());
    }
}