momentry_core/src/main.rs

use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use std::path::Path;

use momentry_core::{Database, PostgresDb, VideoRecord};

#[derive(Parser)]
#[command(name = "momentry")]
#[command(about = "Digital asset management system with video analysis and RAG")]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    /// Register a video file
    Register {
        /// Video file path or URL
        path: String,
    },
    /// Process video (generate all JSON files)
    Process {
        /// UUID or path
        target: String,
    },
    /// Generate chunks and store in database
    Chunk {
        /// UUID
        uuid: String,
    },
    /// Vectorize chunks
    Vectorize {
        /// UUID (or 'all' for all)
        uuid: String,
    },
    /// Play video with overlays
    Play {
        /// Video path or UUID
        target: String,
    },
    /// Start watching directories
    Watch {
        /// Directories to watch (comma separated)
        directories: Option<String>,
    },
    /// Start API server
    Server {
        /// Host
        #[arg(long, default_value = "127.0.0.1")]
        host: String,
        /// Port
        #[arg(long, default_value = "3000")]
        port: u16,
    },
    /// Query using RAG
    Query {
        /// Query text
        query: String,
    },
    /// Lookup UUID from path
    Lookup {
        /// File path
        path: String,
    },
    /// Resolve path from UUID
    Resolve {
        /// UUID
        uuid: String,
    },
    /// Generate thumbnails for videos
    Thumbnails {
        /// UUID (optional, generates for all if not specified)
        uuid: Option<String>,
        /// Number of thumbnails per video
        #[arg(short, long, default_value = "6")]
        count: u32,
    },
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt::init();

    let cli = Cli::parse();

    match cli.command {
        Commands::Register { path } => {
            println!("Registering: {}", path);

            // Compute UUID
            let uuid = momentry_core::uuid::compute_uuid_from_path(&path);
            println!("UUID: {}", uuid);

            // Run ffprobe
            let probe_result = momentry_core::core::probe::probe_video(&path)?;

            println!("\nVideo probe results:");
            let duration = probe_result
                .format
                .duration
                .as_ref()
                .and_then(|s| s.parse::<f64>().ok())
                .unwrap_or(0.0);
            println!("  Duration: {}s", duration);
            if let Some(size) = &probe_result.format.size {
                println!("  Size: {}", size);
            }

            let mut width = 0u32;
            let mut height = 0u32;
            let mut fps = 0.0;

            for stream in &probe_result.streams {
                if stream.codec_type.as_deref() == Some("video") {
                    width = stream.width.unwrap_or(0);
                    height = stream.height.unwrap_or(0);
                    if let Some(fps_str) = &stream.r_frame_rate {
                        if let Some((num, den)) = fps_str.split_once('/') {
                            if let (Ok(n), Ok(d)) = (num.parse::<f64>(), den.parse::<f64>()) {
                                if d > 0.0 {
                                    fps = n / d;
                                }
                            }
                        }
                    }
                    println!("  Video: {}x{}", width, height);
                    if let Some(fps) = &stream.r_frame_rate {
                        println!("  FPS: {}", fps);
                    }
                }
                if stream.codec_type.as_deref() == Some("audio") {
                    println!("  Audio: {} channels", stream.channels.unwrap_or(0));
                    if let Some(sr) = &stream.sample_rate {
                        println!("  Sample Rate: {}", sr);
                    }
                }
            }

            // Save probe JSON to file
            let file_manager = momentry_core::FileManager::new(std::path::PathBuf::from("."));
            let json_str = serde_json::to_string_pretty(&probe_result)?;
            let json_path = file_manager.save_json(&uuid, "probe", &json_str)?;
            println!("\nProbe JSON saved to: {:?}", json_path);

            // Store in PostgreSQL
            println!("\nStoring in database...");
            let db = PostgresDb::init().await?;
            let file_path = Path::new(&path)
                .canonicalize()
                .map(|p| p.to_string_lossy().to_string())
                .unwrap_or_else(|_| path.clone());
            let file_name = Path::new(&path)
                .file_name()
                .map(|n| n.to_string_lossy().to_string())
                .unwrap_or_default();

            let record = VideoRecord {
                id: 0,
                uuid: uuid.clone(),
                file_path,
                file_name,
                duration,
                width,
                height,
                fps,
                probe_json: Some(json_str),
                created_at: String::new(),
            };

            let video_id = db.register_video(&record).await?;
            println!("Video registered with ID: {}", video_id);

            Ok(())
        }
        Commands::Process { target } => {
            println!("Processing: {}", target);

            // Compute UUID if path is given
            let uuid = if target.len() == 16 && !target.contains('/') {
                target.clone()
            } else {
                momentry_core::uuid::compute_uuid_from_path(&target)
            };

            // Get video from database
            let db = PostgresDb::init().await?;
            let video = db
                .get_video_by_uuid(&uuid)
                .await?
                .ok_or_else(|| anyhow::anyhow!("Video not found: {}", uuid))?;

            let video_path = &video.file_path;
            let file_manager = momentry_core::FileManager::new(std::path::PathBuf::from("."));

            // Process ASR
            println!("\nRunning ASR...");
            let asr_path = format!("{}.asr.json", uuid);
            let asr_result =
                momentry_core::core::processor::process_asr(video_path, &asr_path).await?;
            let asr_json = serde_json::to_string_pretty(&asr_result)?;
            std::fs::write(&asr_path, &asr_json)?;
            println!("ASR saved to: {}", asr_path);
            println!("  {} segments found", asr_result.segments.len());

            // TODO: Process OCR, YOLO, Face, Pose, ASRx
            println!("\nOther processors not yet implemented.");

            Ok(())
        }
        Commands::Chunk { uuid } => {
            println!("Chunking: {}", uuid);

            let db = PostgresDb::init().await?;
            let video = db
                .get_video_by_uuid(&uuid)
                .await?
                .ok_or_else(|| anyhow::anyhow!("Video not found: {}", uuid))?;

            // Read ASR JSON
            let asr_path = format!("{}.asr.json", uuid);
            let asr_json = std::fs::read_to_string(&asr_path)
                .context("ASR file not found. Run 'process' first.")?;

            let asr_result: momentry_core::core::processor::asr::AsrResult =
                serde_json::from_str(&asr_json)?;

            println!("Processing {} ASR segments...", asr_result.segments.len());

            // Split into sentence chunks
            let mut sentence_chunks = Vec::new();
            for (i, seg) in asr_result.segments.iter().enumerate() {
                let chunk = momentry_core::Chunk::new(
                    uuid.clone(),
                    i as u32,
                    momentry_core::ChunkType::Sentence,
                    seg.start,
                    seg.end,
                    serde_json::json!({
                        "text": seg.text,
                    }),
                );
                sentence_chunks.push(chunk);
            }

            // Split into time-based chunks (10 seconds)
            let splitter = momentry_core::core::chunk::ChunkSplitter::new(10.0);
            let time_chunks = splitter.split_time_based(&uuid, video.duration);

            // Store in database
            println!("Storing {} sentence chunks...", sentence_chunks.len());
            for chunk in &sentence_chunks {
                db.store_chunk(chunk).await?;
            }

            println!("Storing {} time-based chunks...", time_chunks.len());
            for chunk in &time_chunks {
                db.store_chunk(chunk).await?;
            }

            println!(
                "Done! {} total chunks stored.",
                sentence_chunks.len() + time_chunks.len()
            );

            Ok(())
        }
        Commands::Vectorize { uuid } => {
            println!("Vectorizing: {}", uuid);
            // TODO: Implement vectorize
            Ok(())
        }
        Commands::Play { target } => {
            println!("Playing: {}", target);
            // TODO: Implement play
            Ok(())
        }
        Commands::Watch { directories } => {
            println!("Starting watcher: {:?}", directories);
            // TODO: Implement watch
            Ok(())
        }
        Commands::Server { host, port } => {
            println!("Starting API server at {}:{}", host, port);
            // TODO: Implement server
            Ok(())
        }
        Commands::Query { query } => {
            println!("Query: {}", query);
            // TODO: Implement query
            Ok(())
        }
        Commands::Lookup { path } => {
            let uuid = momentry_core::uuid::compute_uuid_from_path(&path);
            println!("Path: {}", path);
            println!("UUID: {}", uuid);
            Ok(())
        }
        Commands::Resolve { uuid } => {
            println!("Resolving UUID: {}", uuid);
            // TODO: Look up path from UUID in database
            println!("(Database lookup not implemented yet)");
            Ok(())
        }
        Commands::Thumbnails { uuid, count } => {
            let db = PostgresDb::init().await?;

            let videos = if let Some(ref uuid) = uuid {
                vec![db
                    .get_video_by_uuid(uuid)
                    .await?
                    .ok_or_else(|| anyhow::anyhow!("Video not found: {}", uuid))?]
            } else {
                db.list_videos().await?
            };

            let output_dir = std::path::PathBuf::from("thumbnails");
            let extractor = momentry_core::ThumbnailExtractor::new(output_dir, count);

            for video in videos {
                println!(
                    "\nGenerating thumbnails for: {} ({})",
                    video.file_name, video.uuid
                );

                match extractor.get_or_create(&video.file_path, &video.uuid) {
                    Ok(result) => {
                        println!("  Generated {} thumbnails", result.count);
                    }
                    Err(e) => {
                        println!("  Error: {}", e);
                    }
                }
            }

            println!("\nThumbnails generated successfully!");
            Ok(())
        }
    }
}