momentry_core/src/api/identities.rs

use axum::{
    body::Body,
    extract::{Path, Query, State},
    http::{header, StatusCode},
    response::{IntoResponse, Json},
    routing::{get, post},
    Router,
};
use serde::{Deserialize, Serialize};
use std::process::Command;

use crate::core::db::{Database, PostgresDb};

#[derive(Debug, Deserialize)]
pub struct CreateIdentityRequest {
    pub face_json_path: String,
    pub identity_name: String,
    pub schema: Option<String>,
}

#[derive(Debug, Serialize)]
pub struct CreateIdentityResponse {
    pub success: bool,
    pub message: String,
    pub identity_uuid: Option<String>,
    pub identity_name: String,
    pub total_vectors: Option<i32>,
    pub angle_coverage: Option<Vec<String>>,
    pub quality_avg: Option<f64>,
}

pub fn identity_routes() -> Router<crate::api::types::AppState> {
    Router::new()
        .route("/api/v1/identities", get(list_identities))
        .route("/api/v1/identity", post(create_identity))
        .route("/api/v1/faces/candidates", get(list_face_candidates))
        .route("/api/v1/traces/unassigned", get(list_unassigned_traces))
}

/// Register a Global Identity from face.json with multi-angle reference vectors.
/// Calls select_face_reference_vectors_v2.py for automatic reference selection.
async fn create_identity(
    State(_state): State<crate::api::types::AppState>,
    Json(req): Json<CreateIdentityRequest>,
) -> Result<Json<CreateIdentityResponse>, (StatusCode, String)> {
    let schema = req.schema.unwrap_or("dev".to_string());
    let python_path =
        std::env::var("MOMENTRY_PYTHON_PATH").unwrap_or("/opt/homebrew/bin/python3.11".to_string());

    let scripts_dir = std::env::var("MOMENTRY_SCRIPTS_DIR").unwrap_or_else(|_| {
        let mut path = std::env::current_dir().unwrap_or_default();
        path.push("scripts");
        path.to_string_lossy().to_string()
    });

    let script_path = format!("{}/select_face_reference_vectors_v2.py", scripts_dir);

    tracing::info!(
        "Registering identity '{}' from face.json: {}",
        req.identity_name,
        req.face_json_path
    );

    let output = Command::new(&python_path)
        .arg(&script_path)
        .arg("--face-json")
        .arg(&req.face_json_path)
        .arg("--identity-name")
        .arg(&req.identity_name)
        .arg("--register")
        .arg("--schema")
        .arg(&schema)
        .output()
        .map_err(|e| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Failed to execute script: {}", e),
            )
        })?;

    if !output.status.success() {
        let stderr = String::from_utf8_lossy(&output.stderr);
        return Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Script failed: {}", stderr),
        ));
    }

    let db = PostgresDb::init().await.map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("DB error: {}", e),
        )
    })?;

    let id_table = crate::core::db::schema::table_name("identities");
    let query = format!(
        "SELECT uuid, reference_data->'total_references' as total,
               reference_data->'angles_covered' as angles,
               reference_data->'quality_avg' as quality
        FROM {}
        WHERE name = $1
        ORDER BY created_at DESC
        LIMIT 1",
        id_table
    );

    let row: Option<(String, Option<i32>, Option<Vec<String>>, Option<f64>)> =
        sqlx::query_as(&query)
            .bind(&req.identity_name)
            .fetch_optional(db.pool())
            .await
            .map_err(|e| {
                (
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Query error: {}", e),
                )
            })?;

    match row {
        Some((uuid, total, angles, quality)) => Ok(Json(CreateIdentityResponse {
            success: true,
            message: format!(
                "Successfully registered identity '{}' with {} reference vectors",
                req.identity_name,
                total.unwrap_or(0)
            ),
            identity_uuid: Some(uuid),
            identity_name: req.identity_name,
            total_vectors: total,
            angle_coverage: angles,
            quality_avg: quality,
        })),
        None => Ok(Json(CreateIdentityResponse {
            success: true,
            message: format!(
                "Identity '{}' registered, but details not found",
                req.identity_name
            ),
            identity_uuid: None,
            identity_name: req.identity_name,
            total_vectors: None,
            angle_coverage: None,
            quality_avg: None,
        })),
    }
}

/// List all global identities
async fn list_identities(
    State(_state): State<crate::api::types::AppState>,
    Query(query): Query<ListIdentitiesQuery>,
) -> Result<Json<IdentityListResponse>, (StatusCode, String)> {
    let db = match PostgresDb::init().await {
        Ok(db) => db,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("DB error: {}", e),
            ))
        }
    };

    let page = query.page.unwrap_or(1);
    let page_size = query.page_size.unwrap_or(20);
    let offset = ((page - 1) as i64) * (page_size as i64);

    let id_table = crate::core::db::schema::table_name("identities");

    let total: i64 = sqlx::query_scalar(&format!(
        "SELECT COUNT(*) FROM {} WHERE status IS NULL OR status != 'merged'",
        id_table
    ))
    .fetch_one(db.pool())
    .await
    .map_err(|e| {
        (
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Count error: {}", e),
        )
    })?;

    let sql = format!(
        r#"SELECT i.id::int, i.uuid, i.name, i.metadata, i.status, i.starred,
           COALESCE(
               jsonb_agg(jsonb_build_object(
                   'file_uuid', fi.file_uuid,
                   'confidence', fi.confidence,
                   'source', fi.metadata->>'source'
               ) ORDER BY fi.created_at DESC),
               '[]'::jsonb
           ) as file_bindings
           FROM {} i
           LEFT JOIN {} fi ON i.id = fi.identity_id
           WHERE i.status IS NULL OR i.status != 'merged'
           GROUP BY i.id, i.uuid, i.name, i.metadata, i.status, i.starred
           ORDER BY i.id DESC LIMIT $1 OFFSET $2"#,
        id_table,
        crate::core::db::schema::table_name("file_identities")
    );

    let rows: Vec<(
        i32,
        uuid::Uuid,
        String,
        Option<serde_json::Value>,
        Option<String>,
        Option<bool>,
        serde_json::Value,
    )> = match sqlx::query_as(&sql)
        .bind(page_size as i64)
        .bind(offset)
        .fetch_all(db.pool())
        .await
    {
        Ok(rows) => rows,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Query error: {}", e),
            ))
        }
    };

    let identities: Vec<IdentityResponse> = rows
        .into_iter()
        .map(|r| {
            let file_bindings: Vec<FileBinding> =
                r.6.as_array()
                    .map(|arr| {
                        arr.iter()
                            .filter_map(|v| serde_json::from_value(v.clone()).ok())
                            .collect()
                    })
                    .unwrap_or_default();
            let file_uuids: Vec<String> = file_bindings
                .iter()
                .map(|fb| fb.file_uuid.clone())
                .collect();
            IdentityResponse {
                id: r.0,
                identity_uuid: r.1.to_string().replace('-', ""),
                name: r.2,
                metadata: r.3,
                status: r.4,
                starred: r.5.unwrap_or(false),
                file_uuids,
                file_bindings: Some(file_bindings),
            }
        })
        .collect();

    let identities_table = crate::core::db::schema::table_name("identities");
    let total_identities: i64 =
        sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {}", identities_table))
            .fetch_one(db.pool())
            .await
            .unwrap_or(0);
    let tmdb_identities: i64 = sqlx::query_scalar(&format!(
        "SELECT COUNT(*) FROM {} WHERE source = 'tmdb'",
        identities_table
    ))
    .fetch_one(db.pool())
    .await
    .unwrap_or(0);
    let auto_identities: i64 = sqlx::query_scalar(&format!(
        "SELECT COUNT(*) FROM {} WHERE file_uuid IS NOT NULL",
        crate::core::db::schema::table_name("strangers")
    ))
    .fetch_one(db.pool())
    .await
    .unwrap_or(0);

    Ok(Json(IdentityListResponse {
        identities,
        count: total,
        page,
        page_size,
        total_identities,
        tmdb_identities,
        auto_identities,
    }))
}

#[derive(Debug, Deserialize)]
pub struct ListIdentitiesQuery {
    pub page: Option<usize>,
    pub page_size: Option<usize>,
}

#[derive(Debug, Deserialize)]
pub struct FaceCandidatesQuery {
    pub file_uuid: Option<String>,
    pub min_confidence: Option<f64>,
    pub page: Option<usize>,
    pub page_size: Option<usize>,
    pub limit: Option<usize>,
}

#[derive(Debug, Serialize)]
pub struct FaceCandidate {
    pub id: i32,
    pub face_id: Option<String>,
    pub file_uuid: String,
    pub frame_number: i64,
    pub confidence: f32,
    pub bbox: Option<serde_json::Value>,
    pub attributes: Option<serde_json::Value>,
}

#[derive(Debug, Serialize)]
pub struct FaceCandidatesResponse {
    pub candidates: Vec<FaceCandidate>,
    pub total: i64,
    pub page: usize,
    pub page_size: usize,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct FileBinding {
    pub file_uuid: String,
    pub confidence: f64,
    pub source: Option<String>,
}

#[derive(Debug, Serialize)]
pub struct IdentityResponse {
    pub id: i32,
    pub identity_uuid: String,
    pub name: String,
    pub metadata: Option<serde_json::Value>,
    pub status: Option<String>,
    pub starred: bool,
    pub file_uuids: Vec<String>,
    pub file_bindings: Option<Vec<FileBinding>>,
}

#[derive(Debug, Serialize)]
pub struct IdentityListResponse {
    pub identities: Vec<IdentityResponse>,
    pub count: i64,
    pub page: usize,
    pub page_size: usize,
    pub total_identities: i64,
    pub tmdb_identities: i64,
    pub auto_identities: i64,
}

async fn list_face_candidates(
    Query(query): Query<FaceCandidatesQuery>,
) -> Result<Json<FaceCandidatesResponse>, (StatusCode, String)> {
    let page = query.page.unwrap_or(1);
    let page_size = std::cmp::min(query.page_size.unwrap_or(15), 100);
    let offset = (page - 1) * page_size;
    let min_confidence = query.min_confidence.unwrap_or(0.5);

    // Query Qdrant _faces for unbound faces (identity_id IS NULL)
    let qdrant = crate::core::db::qdrant_db::QdrantDb::new();
    let mut filter_must = vec![
        serde_json::json!({"is_null": {"key": "identity_id"}}),
        serde_json::json!({"key": "confidence", "range": {"gte": min_confidence}}),
    ];
    if let Some(ref file_uuid) = query.file_uuid {
        filter_must.push(serde_json::json!({"key": "file_uuid", "match": {"value": file_uuid}}));
    }
    let scroll_filter = serde_json::json!({"must": filter_must});

    let all_points = qdrant
        .scroll_all_points("_faces", scroll_filter, 1000)
        .await
        .map_err(|e| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Qdrant scroll failed: {}", e),
            )
        })?;

    let total = all_points.len() as i64;

    // Sort by confidence DESC then paginate
    let mut sorted: Vec<&serde_json::Value> = all_points.iter().collect();
    sorted.sort_by(|a, b| {
        let ca = a["payload"]["confidence"].as_f64().unwrap_or(0.0);
        let cb = b["payload"]["confidence"].as_f64().unwrap_or(0.0);
        cb.partial_cmp(&ca).unwrap_or(std::cmp::Ordering::Equal)
    });
    let paginated: Vec<&&serde_json::Value> = sorted.iter().skip(offset).take(page_size).collect();

    let candidates: Vec<FaceCandidate> = paginated
        .into_iter()
        .map(|p| {
            let payload = &p["payload"];
            let point_id = p["id"].as_u64().unwrap_or(0);
            FaceCandidate {
                id: point_id as i32,
                face_id: Some(format!("{:x}", point_id)),
                file_uuid: payload["file_uuid"].as_str().unwrap_or("").to_string(),
                frame_number: payload["frame"].as_i64().unwrap_or(0),
                confidence: payload["confidence"].as_f64().unwrap_or(0.0) as f32,
                bbox: payload.get("bbox").cloned(),
                attributes: None,
            }
        })
        .collect();

    Ok(Json(FaceCandidatesResponse {
        candidates,
        total,
        page,
        page_size,
    }))
}

#[derive(Debug, Deserialize)]
pub struct UnassignedTracesQuery {
    pub file_uuid: Option<String>,
    pub page: Option<usize>,
    pub page_size: Option<usize>,
}

#[derive(Debug, Serialize)]
pub struct UnassignedTrace {
    pub trace_id: i32,
    pub file_uuid: String,
    pub frame_count: i64,
    pub start_frame: i64,
    pub end_frame: i64,
    pub best_face_id: i32,
    pub best_face_frame: i64,
    pub best_face_confidence: f64,
    pub best_face_bbox: Option<serde_json::Value>,
}

#[derive(Debug, Serialize)]
pub struct UnassignedTracesResponse {
    pub traces: Vec<UnassignedTrace>,
    pub total: i64,
    pub page: usize,
    pub page_size: usize,
}

/// List unassigned traces (identity_id IS NULL, grouped by trace_id)
async fn list_unassigned_traces(
    Query(query): Query<UnassignedTracesQuery>,
) -> Result<Json<UnassignedTracesResponse>, (StatusCode, String)> {
    let page = query.page.unwrap_or(1);
    let page_size = std::cmp::min(query.page_size.unwrap_or(20), 100);
    let offset = (page - 1) * page_size;

    // Query Qdrant _faces for unbound traces (identity_id IS NULL, trace_id > 0)
    let qdrant = crate::core::db::qdrant_db::QdrantDb::new();
    let mut filter_must: Vec<serde_json::Value> = vec![
        serde_json::json!({"is_null": {"key": "identity_id"}}),
        serde_json::json!({"key": "trace_id", "range": {"gt": 0}}),
    ];
    if let Some(ref file_uuid) = query.file_uuid {
        filter_must.push(serde_json::json!({"key": "file_uuid", "match": {"value": file_uuid}}));
    }
    let scroll_filter = serde_json::json!({"must": filter_must});

    let all_points = qdrant
        .scroll_all_points("_faces", scroll_filter, 1000)
        .await
        .map_err(|e| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Qdrant scroll failed: {}", e),
            )
        })?;

    // Group by (file_uuid, trace_id) and aggregate
    use std::collections::BTreeMap;
    #[derive(Default)]
    struct TraceAgg {
        frame_count: i64,
        start_frame: i64,
        end_frame: i64,
        best_confidence: f64,
        best_point_id: i64,
        best_frame: i64,
        best_bbox: Option<serde_json::Value>,
    }

    let mut trace_map: BTreeMap<(String, i32), TraceAgg> = BTreeMap::new();
    for point in &all_points {
        let payload = &point["payload"];
        let file_uuid = match payload["file_uuid"].as_str() {
            Some(f) => f.to_string(),
            None => continue,
        };
        let trace_id = payload["trace_id"].as_i64().unwrap_or(0) as i32;
        if trace_id <= 0 {
            continue;
        }
        let frame = payload["frame"].as_i64().unwrap_or(0);
        let confidence = payload["confidence"].as_f64().unwrap_or(0.0);
        let point_id = point["id"].as_i64().unwrap_or(0);

        let entry = trace_map.entry((file_uuid, trace_id)).or_default();
        entry.frame_count += 1;
        if frame < entry.start_frame || entry.start_frame == 0 {
            entry.start_frame = frame;
        }
        if frame > entry.end_frame {
            entry.end_frame = frame;
        }
        if confidence > entry.best_confidence {
            entry.best_confidence = confidence;
            entry.best_point_id = point_id;
            entry.best_frame = frame;
            entry.best_bbox = payload.get("bbox").cloned();
        }
    }

    let total = trace_map.len() as i64;

    // Sort by frame_count DESC, paginate
    let mut sorted_traces: Vec<((String, i32), TraceAgg)> = trace_map.into_iter().collect();
    sorted_traces.sort_by(|a, b| b.1.frame_count.cmp(&a.1.frame_count));
    let paginated: Vec<_> = sorted_traces
        .into_iter()
        .skip(offset)
        .take(page_size)
        .collect();

    let traces: Vec<UnassignedTrace> = paginated
        .into_iter()
        .map(|((file_uuid, trace_id), agg)| UnassignedTrace {
            trace_id,
            file_uuid,
            frame_count: agg.frame_count,
            start_frame: agg.start_frame,
            end_frame: agg.end_frame,
            best_face_id: agg.best_point_id as i32,
            best_face_frame: agg.best_frame,
            best_face_confidence: agg.best_confidence,
            best_face_bbox: agg.best_bbox,
        })
        .collect();

    Ok(Json(UnassignedTracesResponse {
        traces,
        total,
        page,
        page_size,
    }))
}