momentry_core/src/api/trace_agent_api.rs

use axum::{
    extract::{Path, Query, State},
    http::StatusCode,
    response::Json,
    routing::{get, post},
    Router,
};
use serde::{Deserialize, Serialize};

use crate::core::db::PostgresDb;

pub fn trace_agent_routes() -> Router<crate::api::server::AppState> {
    Router::new()
        .route(
            "/api/v1/file/:file_uuid/traces",
            post(list_traces_sorted),
        )
        .route(
            "/api/v1/file/:file_uuid/trace/:trace_id/faces",
            get(list_trace_faces),
        )
}

#[derive(Debug, Deserialize)]
struct TracesRequest {
    min_faces: Option<i64>,
    sort_by: Option<String>,
    page: Option<i64>,
    page_size: Option<i64>,
    limit: Option<i64>,
    min_confidence: Option<f64>,
    max_confidence: Option<f64>,
}

#[derive(Debug, Serialize)]
struct TraceInfo {
    trace_id: i32,
    face_count: i64,
    start_frame: i32,
    end_frame: i32,
    start_time: f64,
    end_time: f64,
    duration_sec: f64,
    avg_confidence: f64,
    sample_face_id: Option<String>,
}

#[derive(Debug, Serialize)]
struct TracesResponse {
    success: bool,
    file_uuid: String,
    fps: f64,
    total_traces: i64,
    total_faces: i64,
    page: i64,
    page_size: i64,
    traces: Vec<TraceInfo>,
}

async fn list_traces_sorted(
    State(state): State<crate::api::server::AppState>,
    Path(file_uuid): Path<String>,
    Json(req): Json<TracesRequest>,
) -> Result<Json<TracesResponse>, (StatusCode, String)> {
    let min_faces = req.min_faces.unwrap_or(1);
    let sort = req.sort_by.as_deref().unwrap_or("first_appearance");
    let page = req.page.unwrap_or(1).max(1);
    let page_size = req.page_size.unwrap_or(50).max(1).min(500);
    let hard_limit = req.limit.unwrap_or(500);
    let effective_limit = hard_limit.min(page_size);
    let db_offset = (page - 1) * page_size;
    let min_confidence = req.min_confidence.unwrap_or(0.0);
    let max_confidence = req.max_confidence.unwrap_or(1.0);

    let order_clause = match sort {
        "face_count" => "face_count DESC",
        "duration" => "duration_sec DESC",
        _ => "start_frame ASC",
    };

    let fps: f64 =
        sqlx::query_scalar(&format!("SELECT COALESCE(fps, 24.0) FROM {} WHERE file_uuid = $1",
            crate::core::db::schema::table_name("videos")))
            .bind(&file_uuid)
            .fetch_optional(state.db.pool())
            .await
            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
            .unwrap_or(24.0);

    let query = format!(
        "SELECT tt.*, fd.id AS sample_face_id FROM (
            SELECT trace_id::int AS trace_id,
                   COUNT(*) AS face_count,
                   MIN(frame_number)::int AS start_frame,
                   MAX(frame_number)::int AS end_frame,
                   (MAX(frame_number) - MIN(frame_number))::float8 AS duration_sec,
                   AVG(confidence)::float8 AS avg_confidence
            FROM {}
            WHERE file_uuid = $1 AND trace_id IS NOT NULL
              AND confidence >= $5 AND confidence <= $6
            GROUP BY trace_id
            HAVING COUNT(*) >= $2
            ORDER BY {}
            LIMIT $3 OFFSET $4
        ) tt
         LEFT JOIN LATERAL (
             SELECT id FROM {}
             WHERE trace_id = tt.trace_id AND file_uuid = $1
             ORDER BY confidence DESC LIMIT 1
         ) fd ON true",
        crate::core::db::schema::table_name("face_detections"),
        order_clause,
        crate::core::db::schema::table_name("face_detections"),
    );

    let rows: Vec<(i32, i64, i32, i32, f64, f64, Option<i32>)> =
        sqlx::query_as(&query)
            .bind(&file_uuid)
            .bind(min_faces)
            .bind(effective_limit)
            .bind(db_offset)
            .bind(min_confidence)
            .bind(max_confidence)
            .fetch_all(state.db.pool())
            .await
            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    let traces: Vec<TraceInfo> = rows
        .into_iter()
        .map(|(tid, fc, sf, ef, dur, conf, fid)| TraceInfo {
            trace_id: tid,
            face_count: fc,
            start_frame: sf,
            end_frame: ef,
            start_time: sf as f64 / fps,
            end_time: ef as f64 / fps,
            duration_sec: dur / fps,
            avg_confidence: conf,
            sample_face_id: fid.map(|v| v.to_string()),
        })
        .collect();

    let (total_traces, total_faces): (i64, i64) = sqlx::query_as(
        &format!("SELECT COUNT(DISTINCT trace_id), COUNT(*) FROM {} WHERE file_uuid = $1 AND trace_id IS NOT NULL",
            crate::core::db::schema::table_name("face_detections"))
    )
    .bind(&file_uuid)
    .fetch_one(state.db.pool())
    .await
    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    Ok(Json(TracesResponse {
        success: true,
        file_uuid,
        fps,
        total_traces,
        total_faces,
        page,
        page_size,
        traces,
    }))
}

// ── Individual face detections for a trace ──

#[derive(Debug, Deserialize)]
struct TraceFacesQuery {
    page: Option<i64>,
    page_size: Option<i64>,
    limit: Option<i64>,
    offset: Option<i64>,
    interpolate: Option<bool>,
}

#[derive(Debug, Serialize)]
struct TraceFaceItem {
    id: i32,
    start_frame: i32,
    end_frame: i32,
    start_time: f64,
    end_time: f64,
    x: Option<i32>,
    y: Option<i32>,
    width: Option<i32>,
    height: Option<i32>,
    confidence: f64,
    interpolated: bool,
}

#[derive(Debug, Serialize)]
struct TraceFacesResponse {
    success: bool,
    file_uuid: String,
    trace_id: i32,
    fps: f64,
    total: i64,
    faces: Vec<TraceFaceItem>,
}

fn lerp_i32(a: Option<i32>, b: Option<i32>, t: f64) -> Option<i32> {
    match (a, b) {
        (Some(av), Some(bv)) => Some((av as f64 + (bv - av) as f64 * t).round() as i32),
        _ => None,
    }
}

async fn list_trace_faces(
    State(state): State<crate::api::server::AppState>,
    Path((file_uuid, trace_id)): Path<(String, i32)>,
    Query(q): Query<TraceFacesQuery>,
) -> Result<Json<TraceFacesResponse>, (StatusCode, String)> {
    let limit = q.limit.unwrap_or(200).min(1000);
    // Support both page/page_size and offset; page/page_size takes precedence
    let offset = if q.page.is_some() || q.page_size.is_some() {
        let p = q.page.unwrap_or(1).max(1);
        let ps = q.page_size.unwrap_or(200).max(1).min(1000);
        (p - 1) * ps
    } else {
        q.offset.unwrap_or(0)
    };
    let interpolate = q.interpolate.unwrap_or(false);

    let fps: f64 =
        sqlx::query_scalar(&format!("SELECT COALESCE(fps, 24.0) FROM {} WHERE file_uuid = $1",
            crate::core::db::schema::table_name("videos")))
            .bind(&file_uuid)
            .fetch_optional(state.db.pool())
            .await
            .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?
            .unwrap_or(24.0);

    let total_detected: i64 = sqlx::query_scalar(
        &format!("SELECT COUNT(*) FROM {} WHERE file_uuid = $1 AND trace_id = $2",
            crate::core::db::schema::table_name("face_detections"))
    )
    .bind(&file_uuid)
    .bind(trace_id)
    .fetch_one(state.db.pool())
    .await
    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    let rows: Vec<(
        i32,
        i32,
        Option<i32>,
        Option<i32>,
        Option<i32>,
        Option<i32>,
        f32,
    )> = sqlx::query_as(
        &format!("SELECT id, frame_number::int, x, y, width, height, confidence::float4 \
                   FROM {} WHERE file_uuid = $1 AND trace_id = $2 \
                   ORDER BY frame_number ASC LIMIT $3 OFFSET $4",
            crate::core::db::schema::table_name("face_detections"))
    )
    .bind(&file_uuid)
    .bind(trace_id)
    .bind(limit)
    .bind(offset)
    .fetch_all(state.db.pool())
    .await
    .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;

    let mut faces: Vec<TraceFaceItem> = Vec::new();

    for (i, (id, frame, x, y, w, h, conf)) in rows.iter().enumerate() {
        let cur = (x, y, w, h);

        // Add interpolated frames between previous and current detection
        if interpolate && i > 0 {
            let prev = &rows[i - 1];
            let prev_frame = prev.1;
            let gap = frame - prev_frame;
            if gap > 1 {
                for mid in 1..gap {
                    let t = mid as f64 / gap as f64;
                    let mid_x = lerp_i32(prev.2, *x, t);
                    let mid_y = lerp_i32(prev.3, *y, t);
                    let mid_w = lerp_i32(prev.4, *w, t);
                    let mid_h = lerp_i32(prev.5, *h, t);
                    let mid_frame = prev_frame + mid;
                    let mt = (mid_frame as f64 / fps * 10.0).round() / 10.0;
                    faces.push(TraceFaceItem {
                        id: 0,
                        start_frame: mid_frame,
                        end_frame: mid_frame,
                        start_time: mt,
                        end_time: mt,
                        x: mid_x,
                        y: mid_y,
                        width: mid_w,
                        height: mid_h,
                        confidence: 0.0,
                        interpolated: true,
                    });
                }
            }
        }

        // Add the real detection
        let frame_val = *frame;
        let ft = (frame_val as f64 / fps * 10.0).round() / 10.0;
        faces.push(TraceFaceItem {
            id: *id,
            start_frame: frame_val,
            end_frame: frame_val,
            start_time: ft,
            end_time: ft,
            x: *x,
            y: *y,
            width: *w,
            height: *h,
            confidence: *conf as f64,
            interpolated: false,
        });
    }

    let total = if interpolate && faces.len() as i64 > total_detected {
        faces.len() as i64
    } else {
        total_detected
    };

    Ok(Json(TraceFacesResponse {
        success: true,
        file_uuid,
        trace_id,
        fps,
        total,
        faces,
    }))
}