PersonalWebApplication/Server/src/media.rs

//! Special-folder media storage: images plus videos with async ffmpeg transcode.

use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use tokio::process::Command;
use tokio::sync::Semaphore;

use crate::folders::MAX_ITEMS;
use crate::photos::{self, MAX_BYTES};

pub const MAX_VIDEO_BYTES: u64 = 500 * 1024 * 1024;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum MediaType {
    Image,
    Video,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum MediaStatus {
    Ready,
    Processing,
    Failed,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct MediaItem {
    pub id: String,
    #[serde(rename = "type")]
    pub media_type: MediaType,
    pub status: MediaStatus,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UploadKind {
    Image(&'static str),
    Video(&'static str),
}

#[derive(Debug, PartialEq, Eq)]
pub enum MediaError {
    Full,
    TooLarge,
    BadType,
    BadFolder,
    Io,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct VideoMeta {
    id: String,
    #[serde(rename = "type")]
    media_type: MediaType,
    status: MediaStatus,
    original_name: String,
    created_at: String,
    output_file: String,
    error: Option<String>,
}

#[derive(Clone)]
pub struct FolderMediaStore {
    base: PathBuf,
    transcodes: Arc<Semaphore>,
}

#[derive(Clone)]
pub struct VideoJob {
    pub folder: String,
    pub id: String,
    pub input_path: PathBuf,
    pub output_path: PathBuf,
    pub meta_path: PathBuf,
}

impl FolderMediaStore {
    pub fn from_env() -> Self {
        let base = std::env::var("SPECIAL_DIR")
            .unwrap_or_else(|_| "special-folders".to_string());
        Self::new(base)
    }

    pub fn new(base: impl Into<PathBuf>) -> Self {
        Self {
            base: base.into(),
            transcodes: Arc::new(Semaphore::new(1)),
        }
    }

    fn folder_dir(&self, folder: &str) -> Option<PathBuf> {
        let safe = photos::sanitize_user(folder)?;
        Some(self.base.join(safe))
    }

    fn media_dir(&self, folder: &str) -> Option<PathBuf> {
        Some(self.folder_dir(folder)?.join("media"))
    }

    fn tmp_dir(&self, folder: &str) -> Option<PathBuf> {
        Some(self.folder_dir(folder)?.join(".tmp"))
    }

    pub fn list(&self, folder: &str) -> (Vec<MediaItem>, u64) {
        let Some(dir) = self.folder_dir(folder) else {
            return (Vec::new(), 0);
        };
        let mut items = Vec::new();
        let mut total_bytes = 0u64;

        if let Ok(rd) = std::fs::read_dir(&dir) {
            for entry in rd.flatten() {
                let path = entry.path();
                if path.is_file() {
                    if let Ok(meta) = entry.metadata() {
                        total_bytes += meta.len();
                    }
                    if let Some(name) = entry.file_name().to_str() {
                        if photos::valid_id(name) {
                            items.push(MediaItem {
                                id: name.to_string(),
                                media_type: MediaType::Image,
                                status: MediaStatus::Ready,
                            });
                        }
                    }
                }
            }
        }

        if let Some(media_dir) = self.media_dir(folder) {
            if let Ok(rd) = std::fs::read_dir(media_dir) {
                for entry in rd.flatten() {
                    let path = entry.path();
                    if path.extension().and_then(|s| s.to_str()) != Some("json") {
                        continue;
                    }
                    if let Ok(bytes) = std::fs::read(&path) {
                        if let Ok(meta) = serde_json::from_slice::<VideoMeta>(&bytes) {
                            let item = MediaItem {
                                id: meta.id.clone(),
                                media_type: MediaType::Video,
                                status: meta.status,
                            };
                            if let Some(parent) = path.parent().and_then(|p| p.parent()) {
                                let output = parent.join(&meta.output_file);
                                if let Ok(meta) = std::fs::metadata(output) {
                                    total_bytes += meta.len();
                                }
                            }
                            items.push(item);
                        }
                    }
                }
            }
        }

        items.sort_by(|a, b| a.id.cmp(&b.id));
        (items, total_bytes)
    }

    pub fn count_items(&self, folder: &str, decoration_count: usize) -> usize {
        let _ = MAX_ITEMS;
        self.list(folder).0.len() + decoration_count
    }

    pub fn save_image_bytes(&self, folder: &str, ext: &str, bytes: &[u8]) -> Result<MediaItem, MediaError> {
        if bytes.len() > MAX_BYTES {
            return Err(MediaError::TooLarge);
        }
        if !matches!(ext, "jpg" | "png" | "gif" | "webp") {
            return Err(MediaError::BadType);
        }
        let dir = self.folder_dir(folder).ok_or(MediaError::BadFolder)?;
        std::fs::create_dir_all(&dir).map_err(|_| MediaError::Io)?;
        let id = format!("{}.{}", new_stem(), ext);
        std::fs::write(dir.join(&id), bytes).map_err(|_| MediaError::Io)?;
        Ok(MediaItem {
            id,
            media_type: MediaType::Image,
            status: MediaStatus::Ready,
        })
    }

    pub fn begin_video_upload(&self, folder: &str, ext: &str, original_name: &str) -> Result<VideoJob, MediaError> {
        if !matches!(ext, "mp4" | "mov" | "m4v" | "webm" | "avi" | "mkv") {
            return Err(MediaError::BadType);
        }
        let folder_dir = self.folder_dir(folder).ok_or(MediaError::BadFolder)?;
        let media_dir = self.media_dir(folder).ok_or(MediaError::BadFolder)?;
        let tmp_dir = self.tmp_dir(folder).ok_or(MediaError::BadFolder)?;
        std::fs::create_dir_all(&folder_dir).map_err(|_| MediaError::Io)?;
        std::fs::create_dir_all(&media_dir).map_err(|_| MediaError::Io)?;
        std::fs::create_dir_all(&tmp_dir).map_err(|_| MediaError::Io)?;

        let id = new_stem();
        let input_path = tmp_dir.join(format!("{id}.{ext}"));
        let output_path = folder_dir.join(format!("{id}.mp4"));
        let meta_path = media_dir.join(format!("{id}.json"));
        let meta = VideoMeta {
            id: id.clone(),
            media_type: MediaType::Video,
            status: MediaStatus::Processing,
            original_name: original_name.to_string(),
            created_at: chrono::Utc::now().to_rfc3339(),
            output_file: format!("{id}.mp4"),
            error: None,
        };
        write_meta(&meta_path, &meta)?;
        Ok(VideoJob {
            folder: folder.to_string(),
            id,
            input_path,
            output_path,
            meta_path,
        })
    }

    pub fn finish_video_upload(&self, job: VideoJob) -> Result<MediaItem, MediaError> {
        self.spawn_transcode(job.clone());
        Ok(MediaItem {
            id: job.id,
            media_type: MediaType::Video,
            status: MediaStatus::Processing,
        })
    }

    pub fn spawn_transcode(&self, job: VideoJob) {
        let semaphore = self.transcodes.clone();
        tokio::spawn(async move {
            let _permit = match semaphore.acquire_owned().await {
                Ok(permit) => permit,
                Err(_) => return,
            };
            let args = ffmpeg_args(&job.input_path, &job.output_path);
            let status = Command::new("ffmpeg").args(&args).status().await;
            match status {
                Ok(code) if code.success() => {
                    update_video_status(&job.meta_path, MediaStatus::Ready, None);
                    let _ = tokio::fs::remove_file(&job.input_path).await;
                }
                Ok(code) => {
                    update_video_status(
                        &job.meta_path,
                        MediaStatus::Failed,
                        Some(format!("ffmpeg exited with status {code}")),
                    );
                    let _ = tokio::fs::remove_file(&job.input_path).await;
                    let _ = tokio::fs::remove_file(&job.output_path).await;
                }
                Err(err) => {
                    update_video_status(&job.meta_path, MediaStatus::Failed, Some(err.to_string()));
                    let _ = tokio::fs::remove_file(&job.input_path).await;
                    let _ = tokio::fs::remove_file(&job.output_path).await;
                }
            }
        });
    }

    pub fn read(&self, folder: &str, id: &str) -> Option<(&'static str, Vec<u8>)> {
        let dir = self.folder_dir(folder)?;
        if photos::valid_id(id) {
            let bytes = std::fs::read(dir.join(id)).ok()?;
            let ext = id.rsplit_once('.').map(|(_, ext)| ext)?;
            return Some((mime_for_image_ext(ext), bytes));
        }
        if valid_video_id(id) {
            let meta = self.video_meta(folder, id)?;
            if meta.status != MediaStatus::Ready {
                return None;
            }
            let bytes = std::fs::read(dir.join(meta.output_file)).ok()?;
            return Some(("video/mp4", bytes));
        }
        None
    }

    pub fn status(&self, folder: &str, id: &str) -> Option<MediaItem> {
        if photos::valid_id(id) {
            let dir = self.folder_dir(folder)?;
            if dir.join(id).is_file() {
                return Some(MediaItem {
                    id: id.to_string(),
                    media_type: MediaType::Image,
                    status: MediaStatus::Ready,
                });
            }
        }
        let meta = self.video_meta(folder, id)?;
        Some(MediaItem {
            id: meta.id,
            media_type: MediaType::Video,
            status: meta.status,
        })
    }

    pub fn delete(&self, folder: &str, id: &str) -> bool {
        let Some(dir) = self.folder_dir(folder) else {
            return false;
        };
        if photos::valid_id(id) {
            return std::fs::remove_file(dir.join(id)).is_ok();
        }
        if !valid_video_id(id) {
            return false;
        }
        let meta = self.video_meta(folder, id);
        let mut removed = false;
        if let Some(meta) = meta {
            removed |= std::fs::remove_file(dir.join(meta.output_file)).is_ok();
        }
        if let Some(tmp_dir) = self.tmp_dir(folder) {
            if let Ok(rd) = std::fs::read_dir(tmp_dir) {
                for entry in rd.flatten() {
                    if let Some(name) = entry.file_name().to_str() {
                        if name.starts_with(id) {
                            removed |= std::fs::remove_file(entry.path()).is_ok();
                        }
                    }
                }
            }
        }
        if let Some(media_dir) = self.media_dir(folder) {
            removed |= std::fs::remove_file(media_dir.join(format!("{id}.json"))).is_ok();
        }
        removed
    }

    fn video_meta(&self, folder: &str, id: &str) -> Option<VideoMeta> {
        if !valid_video_id(id) {
            return None;
        }
        let path = self.media_dir(folder)?.join(format!("{id}.json"));
        serde_json::from_slice(&std::fs::read(path).ok()?).ok()
    }
}

pub fn classify_upload(content_type: Option<&str>, filename: Option<&str>) -> Option<UploadKind> {
    let mime = content_type.unwrap_or("").to_ascii_lowercase();
    match mime.as_str() {
        "image/jpeg" => return Some(UploadKind::Image("jpg")),
        "image/png" => return Some(UploadKind::Image("png")),
        "image/gif" => return Some(UploadKind::Image("gif")),
        "image/webp" => return Some(UploadKind::Image("webp")),
        "video/mp4" => return Some(UploadKind::Video("mp4")),
        "video/quicktime" => return Some(UploadKind::Video("mov")),
        "video/webm" => return Some(UploadKind::Video("webm")),
        "video/x-msvideo" => return Some(UploadKind::Video("avi")),
        "video/x-matroska" => return Some(UploadKind::Video("mkv")),
        _ => {}
    }
    let ext = filename
        .and_then(|name| name.rsplit_once('.').map(|(_, ext)| ext.to_ascii_lowercase()))?;
    match ext.as_str() {
        "jpg" | "jpeg" => Some(UploadKind::Image("jpg")),
        "png" => Some(UploadKind::Image("png")),
        "gif" => Some(UploadKind::Image("gif")),
        "webp" => Some(UploadKind::Image("webp")),
        "mp4" => Some(UploadKind::Video("mp4")),
        "mov" => Some(UploadKind::Video("mov")),
        "m4v" => Some(UploadKind::Video("m4v")),
        "webm" => Some(UploadKind::Video("webm")),
        "avi" => Some(UploadKind::Video("avi")),
        "mkv" => Some(UploadKind::Video("mkv")),
        _ => None,
    }
}

pub fn valid_video_id(id: &str) -> bool {
    id.len() == 16 && id.chars().all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase())
}

pub fn ffmpeg_args(input: &Path, output: &Path) -> Vec<String> {
    vec![
        "-y".to_string(),
        "-i".to_string(),
        input.to_string_lossy().to_string(),
        "-vf".to_string(),
        "scale='min(iw,if(gte(iw,ih),1280,720))':'min(ih,if(gte(iw,ih),720,1280))':force_original_aspect_ratio=decrease:force_divisible_by=2,fps=24".to_string(),
        "-c:v".to_string(),
        "libx264".to_string(),
        "-preset".to_string(),
        "veryfast".to_string(),
        "-crf".to_string(),
        "30".to_string(),
        "-maxrate".to_string(),
        "900k".to_string(),
        "-bufsize".to_string(),
        "1800k".to_string(),
        "-c:a".to_string(),
        "aac".to_string(),
        "-b:a".to_string(),
        "96k".to_string(),
        "-movflags".to_string(),
        "+faststart".to_string(),
        output.to_string_lossy().to_string(),
    ]
}

fn write_meta(path: &Path, meta: &VideoMeta) -> Result<(), MediaError> {
    let bytes = serde_json::to_vec_pretty(meta).map_err(|_| MediaError::Io)?;
    std::fs::write(path, bytes).map_err(|_| MediaError::Io)
}

fn update_video_status(path: &Path, status: MediaStatus, error: Option<String>) {
    let Ok(bytes) = std::fs::read(path) else {
        return;
    };
    let Ok(mut meta) = serde_json::from_slice::<VideoMeta>(&bytes) else {
        return;
    };
    meta.status = status;
    meta.error = error.map(|msg| msg.chars().take(240).collect());
    let _ = write_meta(path, &meta);
}

fn mime_for_image_ext(ext: &str) -> &'static str {
    match ext {
        "jpg" => "image/jpeg",
        "png" => "image/png",
        "gif" => "image/gif",
        "webp" => "image/webp",
        _ => "application/octet-stream",
    }
}

fn new_stem() -> String {
    static COUNTER: AtomicU64 = AtomicU64::new(0);
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|duration| duration.as_nanos())
        .unwrap_or(0);
    let next = COUNTER.fetch_add(1, Ordering::Relaxed);
    let digest = Sha256::digest(format!("{nanos}-{next}").as_bytes());
    digest.iter().take(8).map(|byte| format!("{byte:02x}")).collect()
}

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

    #[test]
    fn classify_upload_accepts_common_images_and_videos() {
        assert_eq!(classify_upload(Some("image/jpeg"), Some("a.jpg")), Some(UploadKind::Image("jpg")));
        assert_eq!(classify_upload(Some("image/png"), Some("a.png")), Some(UploadKind::Image("png")));
        assert_eq!(classify_upload(Some("video/mp4"), Some("clip.mp4")), Some(UploadKind::Video("mp4")));
        assert_eq!(classify_upload(Some("video/quicktime"), Some("clip.mov")), Some(UploadKind::Video("mov")));
        assert_eq!(classify_upload(Some("application/octet-stream"), Some("clip.mkv")), Some(UploadKind::Video("mkv")));
        assert_eq!(classify_upload(Some("text/plain"), Some("note.txt")), None);
    }

    #[test]
    fn valid_video_id_is_lower_hex_stem_only() {
        assert!(valid_video_id("a1b2c3d4e5f60708"));
        assert!(!valid_video_id("a1b2.mp4"));
        assert!(!valid_video_id("../a1b2c3d4e5f60708"));
        assert!(!valid_video_id("A1b2c3d4e5f60708"));
        assert!(!valid_video_id(""));
    }

    #[test]
    fn ffmpeg_args_match_low_bandwidth_target() {
        let args = ffmpeg_args(Path::new("/tmp/in.mov"), Path::new("/tmp/out.mp4"));
        let joined = args.join(" ");
        assert!(joined.contains("-vf"));
        assert!(joined.contains("fps=24"));
        assert!(joined.contains("-maxrate 900k"));
        assert!(joined.contains("-bufsize 1800k"));
        assert!(joined.contains("-b:a 96k"));
        assert!(joined.contains("+faststart"));
        assert!(joined.contains("force_original_aspect_ratio=decrease"));
        assert!(joined.contains("force_divisible_by=2"));
        assert!(joined.contains("/tmp/in.mov"));
        assert!(joined.contains("/tmp/out.mp4"));
    }
}