momentry_core/src/api/server.rs

use axum::{
    extract::{Query, State},
    http::StatusCode,
    response::Json,
    routing::{get, post},
    Router,
};
use serde::{Deserialize, Serialize};
use std::net::SocketAddr;
use std::path::Path;
use std::sync::Arc;

use crate::core::db::{Database, PostgresDb, QdrantDb, RedisClient, VideoRecord};
use crate::{Embedder, FileManager};

#[derive(Clone)]
struct AppState {
    embedder: Arc<Embedder>,
    #[allow(dead_code)]
    embedder_model: String,
}

#[derive(Debug, Deserialize)]
struct RegisterRequest {
    path: String,
}

#[derive(Debug, Serialize)]
struct RegisterResponse {
    uuid: String,
    video_id: i64,
    file_name: String,
    duration: f64,
    width: u32,
    height: u32,
}

#[derive(Debug, Deserialize)]
struct SearchRequest {
    query: String,
    limit: Option<usize>,
    uuid: Option<String>,
}

#[derive(Debug, Serialize)]
struct SearchResult {
    uuid: String,
    chunk_id: String,
    chunk_type: String,
    start_time: f64,
    end_time: f64,
    text: String,
    score: f32,
}

#[derive(Debug, Serialize)]
struct SearchResponse {
    results: Vec<SearchResult>,
    query: String,
}

#[derive(Debug, Deserialize)]
struct LookupQuery {
    path: Option<String>,
    uuid: Option<String>,
}

#[derive(Debug, Serialize)]
struct LookupResponse {
    uuid: String,
    file_path: Option<String>,
    file_name: Option<String>,
    duration: Option<f64>,
}

#[derive(Debug, Serialize)]
struct VideoInfoResponse {
    uuid: String,
    file_path: String,
    file_name: String,
    duration: f64,
    width: u32,
    height: u32,
}

#[derive(Debug, Serialize)]
struct VideosResponse {
    videos: Vec<VideoInfoResponse>,
}

async fn register(
    State(_state): State<AppState>,
    Json(req): Json<RegisterRequest>,
) -> Result<Json<RegisterResponse>, StatusCode> {
    let path = req.path;

    let uuid = crate::uuid::compute_uuid_from_path(&path);

    let probe_result =
        crate::core::probe::probe_video(&path).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let duration = probe_result
        .format
        .duration
        .as_ref()
        .and_then(|s: &String| s.parse::<f64>().ok())
        .unwrap_or(0.0);

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

    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);
        }
    }

    let file_manager = FileManager::new(std::path::PathBuf::from("."));
    let json_str =
        serde_json::to_string(&probe_result).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    let _json_path = file_manager
        .save_json(&uuid, "probe", &json_str)
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let db = PostgresDb::init()
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    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: file_name.clone(),
        duration,
        width,
        height,
        fps: 0.0,
        probe_json: Some(json_str),
        storage: Default::default(),
        created_at: String::new(),
    };

    let video_id = db
        .register_video(&record)
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    Ok(Json(RegisterResponse {
        uuid,
        video_id,
        file_name,
        duration,
        width,
        height,
    }))
}

async fn search(
    State(state): State<AppState>,
    Json(req): Json<SearchRequest>,
) -> Result<Json<SearchResponse>, StatusCode> {
    let limit = req.limit.unwrap_or(10);

    let query_vector = state.embedder.embed_query(&req.query).await.map_err(|e| {
        tracing::error!("Failed to embed query: {}", e);
        StatusCode::INTERNAL_SERVER_ERROR
    })?;

    let qdrant: QdrantDb = QdrantDb::init()
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
    let pg: PostgresDb = PostgresDb::init()
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let search_results = if let Some(uuid) = &req.uuid {
        let query_f64: Vec<f64> = query_vector.iter().map(|&x| x as f64).collect();
        qdrant
            .search_in_uuid(&query_f64, uuid, limit)
            .await
            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
    } else {
        qdrant
            .search(&query_vector, limit)
            .await
            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
    };

    let mut results = Vec::new();
    for r in search_results {
        let chunks = pg.get_chunks_by_uuid(&r.chunk_id).await.unwrap_or_default();

        for chunk in chunks {
            let text = chunk
                .content
                .get("text")
                .and_then(|v: &serde_json::Value| v.as_str())
                .unwrap_or("")
                .to_string();

            results.push(SearchResult {
                uuid: chunk.uuid.clone(),
                chunk_id: chunk.chunk_id.clone(),
                chunk_type: chunk.chunk_type.as_str().to_string(),
                start_time: chunk.start_time,
                end_time: chunk.end_time,
                text,
                score: r.score,
            });
        }
    }

    Ok(Json(SearchResponse {
        results,
        query: req.query,
    }))
}

async fn lookup(Query(query): Query<LookupQuery>) -> Result<Json<LookupResponse>, StatusCode> {
    let db: PostgresDb = PostgresDb::init()
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    if let Some(path) = query.path {
        let uuid = crate::uuid::compute_uuid_from_path(&path);
        return Ok(Json(LookupResponse {
            uuid,
            file_path: None,
            file_name: None,
            duration: None,
        }));
    }

    if let Some(uuid) = query.uuid {
        let video = db
            .get_video_by_uuid(&uuid)
            .await
            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

        if let Some(v) = video {
            return Ok(Json(LookupResponse {
                uuid: v.uuid,
                file_path: Some(v.file_path),
                file_name: Some(v.file_name),
                duration: Some(v.duration),
            }));
        }
    }

    Err(StatusCode::NOT_FOUND)
}

pub async fn start_server(host: &str, port: u16) -> anyhow::Result<()> {
    let embedder = Arc::new(Embedder::new("nomic-embed-text:v1.5".to_string()));

    let state = AppState {
        embedder,
        embedder_model: "nomic-embed-text:v1.5".to_string(),
    };

    let app = Router::new()
        .route("/api/v1/register", post(register))
        .route("/api/v1/search", post(search))
        .route("/api/v1/lookup", get(lookup))
        .route("/api/v1/videos", get(list_videos))
        .route("/api/v1/progress/:uuid", get(get_progress))
        .with_state(state);

    let addr = SocketAddr::new(host.parse().unwrap(), port);
    tracing::info!("Starting API server at http://{}", addr);

    let listener = tokio::net::TcpListener::bind(addr).await?;
    axum::serve(listener, app).await?;

    Ok(())
}

async fn list_videos() -> Result<Json<VideosResponse>, StatusCode> {
    let db: PostgresDb = PostgresDb::init()
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let videos = db
        .list_videos()
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let video_infos: Vec<VideoInfoResponse> = videos
        .into_iter()
        .map(|v| VideoInfoResponse {
            uuid: v.uuid,
            file_path: v.file_path,
            file_name: v.file_name,
            duration: v.duration,
            width: v.width,
            height: v.height,
        })
        .collect();

    Ok(Json(VideosResponse {
        videos: video_infos,
    }))
}

#[derive(Debug, Serialize)]
struct ProgressResponse {
    uuid: String,
    processors: Vec<ProcessorProgressInfo>,
}

#[derive(Debug, Serialize)]
struct ProcessorProgressInfo {
    name: String,
    status: String,
    current: u32,
    total: u32,
    message: String,
}

async fn get_progress(
    axum::extract::Path(uuid): axum::extract::Path<String>,
) -> Result<Json<ProgressResponse>, StatusCode> {
    let redis = RedisClient::new().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let mut conn = redis
        .get_conn()
        .await
        .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;

    let processor_names = ["asr", "cut", "asrx", "yolo", "ocr", "face", "pose"];
    let mut processors = Vec::new();

    for name in processor_names {
        let key = format!("momentry:job:{}:processor:{}", uuid, name);

        let status: String = redis::cmd("HGET")
            .arg(&key)
            .arg("status")
            .query_async(&mut conn)
            .await
            .unwrap_or_else(|_| "pending".to_string());

        let current: u32 = redis::cmd("HGET")
            .arg(&key)
            .arg("current")
            .query_async(&mut conn)
            .await
            .unwrap_or_else(|_| "0".to_string())
            .parse()
            .unwrap_or(0);

        let total: u32 = redis::cmd("HGET")
            .arg(&key)
            .arg("total")
            .query_async(&mut conn)
            .await
            .unwrap_or_else(|_| "0".to_string())
            .parse()
            .unwrap_or(0);

        let message: String = redis::cmd("HGET")
            .arg(&key)
            .arg("message")
            .query_async(&mut conn)
            .await
            .unwrap_or_else(|_| "".to_string());

        processors.push(ProcessorProgressInfo {
            name: name.to_string(),
            status,
            current,
            total,
            message,
        });
    }

    Ok(Json(ProgressResponse { uuid, processors }))
}