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 RegisterFromPersonRequest {
    pub file_uuid: String,
    pub person_id: String,
    pub identity_name: String,
    pub metadata: Option<serde_json::Value>,
}

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

#[derive(Debug, Serialize)]
pub struct RegisterFromPersonResponse {
    pub success: bool,
    pub message: String,
    pub identity_id: i32,
    pub identity_name: String,
    pub person_id: String,
}

#[derive(Debug, Serialize)]
pub struct RegisterFromFaceResponse {
    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::server::AppState> {
    Router::new()
        .route("/api/v1/identities/from-person", post(register_from_person))
        .route("/api/v1/identities/from-face", post(register_from_face))
        .route("/api/v1/identities", get(list_identities))
        .route("/api/v1/faces/candidates", get(list_face_candidates))
        .route(
            "/api/v1/identities/:identity_id/faces",
            get(get_identity_faces),
        )
        .route("/api/v1/faces/:face_id/thumbnail", get(get_face_thumbnail))
}

/// 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 register_from_face(
    State(_state): State<crate::api::server::AppState>,
    Json(req): Json<RegisterFromFaceRequest>,
) -> Result<Json<RegisterFromFaceResponse>, (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 query = r#"
        SELECT uuid, reference_data->'total_references' as total,
               reference_data->'angles_covered' as angles,
               reference_data->'quality_avg' as quality
        FROM identities
        WHERE name = $1
        ORDER BY created_at DESC
        LIMIT 1
    "#;

    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(RegisterFromFaceResponse {
            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(RegisterFromFaceResponse {
            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,
        })),
    }
}

/// Register a Global Identity from a specific Person in a video.
/// This creates/updates the Identity record, links the Person to the Identity,
/// and updates the Person's name to match the Identity.
async fn register_from_person(
    State(_state): State<crate::api::server::AppState>,
    Json(req): Json<RegisterFromPersonRequest>,
) -> Result<Json<RegisterFromPersonResponse>, (StatusCode, String)> {
    let db = match PostgresDb::init().await {
        Ok(db) => db,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("DB error: {}", e),
            ))
        }
    };

    let mut tx = match db.pool().begin().await {
        Ok(tx) => tx,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Tx error: {}", e),
            ))
        }
    };

    // 1. Check if Person exists
    let person_query =
        "SELECT id, name FROM person_identities WHERE person_id = $1 AND file_uuid = $2";
    let person: Option<(i32, Option<String>)> = match sqlx::query_as(person_query)
        .bind(&req.person_id)
        .bind(&req.file_uuid)
        .fetch_optional(&mut *tx)
        .await
    {
        Ok(p) => p,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Query error: {}", e),
            ))
        }
    };

    let (person_db_id, _old_name) = match person {
        Some(p) => p,
        None => {
            return Err((
                StatusCode::NOT_FOUND,
                format!(
                    "Person '{}' not found in video '{}'",
                    req.person_id, req.file_uuid
                ),
            ))
        }
    };

    // 2. Check if Identity exists
    let identity_query = "SELECT id FROM identities WHERE name = $1";
    let identity_id: Option<i32> = match sqlx::query_scalar(identity_query)
        .bind(&req.identity_name)
        .fetch_optional(&mut *tx)
        .await
    {
        Ok(id) => id,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Query error: {}", e),
            ))
        }
    };

    let final_identity_id = if let Some(id) = identity_id {
        id
    } else {
        // Create new Identity
        let meta_json = req.metadata.clone().unwrap_or(serde_json::json!({}));

        let new_id: i32 = match sqlx::query_scalar(
            r#"
            INSERT INTO identities (name, embedding, metadata)
            VALUES ($1, NULLIF($2, '')::public.vector, $3)
            RETURNING id
        "#,
        )
        .bind(&req.identity_name)
        .bind("".to_string()) // No embedding for now via this API
        .bind(&meta_json)
        .fetch_one(&mut *tx)
        .await
        {
            Ok(id) => id,
            Err(e) => {
                return Err((
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Insert identity error: {}", e),
                ))
            }
        };
        new_id
    };

    // 3. Create Binding
    // Columns: id, identity_id, identity_type, identity_value, confidence, metadata, created_at
    let binding_query = r#"
        INSERT INTO identity_bindings (identity_id, identity_type, identity_value, confidence, metadata)
        VALUES ($1, $2, $3, $4, $5)
        ON CONFLICT DO NOTHING
    "#;

    match sqlx::query(binding_query)
        .bind(final_identity_id)
        .bind("person_id") // identity_type
        .bind(&req.person_id) // identity_value
        .bind(1.0) // confidence
        .bind(serde_json::to_string(&serde_json::json!({"auto_updated": true})).unwrap())
        .execute(&mut *tx)
        .await
    {
        Ok(_) => {}
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Binding error: {}", e),
            ))
        }
    };

    // 4. Update Person Name
    let update_person = "UPDATE person_identities SET name = $1 WHERE id = $2";
    match sqlx::query(update_person)
        .bind(&req.identity_name)
        .bind(person_db_id)
        .execute(&mut *tx)
        .await
    {
        Ok(_) => {}
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Update person error: {}", e),
            ))
        }
    };

    match tx.commit().await {
        Ok(_) => {}
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Commit error: {}", e),
            ))
        }
    };

    Ok(Json(RegisterFromPersonResponse {
        success: true,
        message: format!(
            "Successfully registered identity '{}' and linked to person '{}'",
            req.identity_name, req.person_id
        ),
        identity_id: final_identity_id,
        identity_name: req.identity_name,
        person_id: req.person_id,
    }))
}

/// List all global identities
async fn list_identities(
    State(_state): State<crate::api::server::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 count_sql = "SELECT COUNT(*) FROM identities";
    let total: i64 = match sqlx::query_scalar(count_sql).fetch_one(db.pool()).await {
        Ok(count) => count,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Count error: {}", e),
            ))
        }
    };

    let sql = "SELECT id, name, metadata FROM identities ORDER BY id DESC LIMIT $1 OFFSET $2";

    let rows: Vec<(i32, String, Option<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| IdentityResponse {
            id: r.0,
            name: r.1,
            metadata: r.2,
        })
        .collect();

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

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

#[derive(Debug, Serialize)]
pub struct IdentityResponse {
    pub id: i32,
    pub name: String,
    pub metadata: Option<serde_json::Value>,
}

#[derive(Debug, Serialize)]
pub struct IdentityListResponse {
    pub identities: Vec<IdentityResponse>,
    pub count: i64,
    pub page: usize,
    pub page_size: 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: f64,
    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, Deserialize)]
pub struct IdentityFacesQuery {
    pub page: Option<usize>,
    pub page_size: Option<usize>,
    pub limit: Option<usize>,
}

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

#[derive(Debug, Serialize)]
pub struct IdentityFacesResponse {
    pub identity_id: i32,
    pub faces: Vec<IdentityFace>,
    pub total: i64,
}

async fn list_face_candidates(
    Query(query): Query<FaceCandidatesQuery>,
) -> Result<Json<FaceCandidatesResponse>, (StatusCode, String)> {
    let db = match PostgresDb::init().await {
        Ok(db) => db,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Failed to connect to database: {}", e),
            ))
        }
    };

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

    let table = crate::core::db::schema::table_name("face_detections");

    let total: i64 = if let Some(file_uuid) = &query.file_uuid {
        let count_sql = format!(
            "SELECT COUNT(*) FROM {} WHERE identity_id IS NULL AND confidence >= $1 AND file_uuid = $2",
            table
        );
        match sqlx::query_scalar(&count_sql)
            .bind(min_confidence)
            .bind(file_uuid)
            .fetch_one(db.pool())
            .await
        {
            Ok(count) => count,
            Err(e) => {
                return Err((
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Count error: {}", e),
                ))
            }
        }
    } else {
        let count_sql = format!(
            "SELECT COUNT(*) FROM {} WHERE identity_id IS NULL AND confidence >= $1",
            table
        );
        match sqlx::query_scalar(&count_sql)
            .bind(min_confidence)
            .fetch_one(db.pool())
            .await
        {
            Ok(count) => count,
            Err(e) => {
                return Err((
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Count error: {}", e),
                ))
            }
        }
    };

    let rows = if let Some(file_uuid) = &query.file_uuid {
        let sql = format!(
            "SELECT id, face_id, file_uuid, frame_number, confidence, bbox, attributes 
             FROM {} 
             WHERE identity_id IS NULL AND confidence >= $1 AND file_uuid = $2
             ORDER BY confidence DESC 
             LIMIT $3 OFFSET $4",
            table
        );
        match sqlx::query_as::<
            _,
            (
                i32,
                Option<String>,
                String,
                i64,
                f64,
                Option<serde_json::Value>,
                Option<serde_json::Value>,
            ),
        >(&sql)
        .bind(min_confidence)
        .bind(file_uuid)
        .bind(page_size as i64)
        .bind(offset as i64)
        .fetch_all(db.pool())
        .await
        {
            Ok(rows) => rows,
            Err(e) => {
                return Err((
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Query error: {}", e),
                ))
            }
        }
    } else {
        let sql = format!(
            "SELECT id, face_id, file_uuid, frame_number, confidence, bbox, attributes 
             FROM {} 
             WHERE identity_id IS NULL AND confidence >= $1
             ORDER BY confidence DESC 
             LIMIT $2 OFFSET $3",
            table
        );
        match sqlx::query_as::<
            _,
            (
                i32,
                Option<String>,
                String,
                i64,
                f64,
                Option<serde_json::Value>,
                Option<serde_json::Value>,
            ),
        >(&sql)
        .bind(min_confidence)
        .bind(page_size as i64)
        .bind(offset as i64)
        .fetch_all(db.pool())
        .await
        {
            Ok(rows) => rows,
            Err(e) => {
                return Err((
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Query error: {}", e),
                ))
            }
        }
    };

    let candidates: Vec<FaceCandidate> = rows
        .into_iter()
        .map(|r| FaceCandidate {
            id: r.0,
            face_id: r.1,
            file_uuid: r.2,
            frame_number: r.3,
            confidence: r.4,
            bbox: r.5,
            attributes: r.6,
        })
        .collect();

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

async fn get_identity_faces(
    axum::extract::Path(identity_id): axum::extract::Path<i32>,
    Query(query): Query<IdentityFacesQuery>,
) -> Result<Json<IdentityFacesResponse>, (StatusCode, String)> {
    let db = match PostgresDb::init().await {
        Ok(db) => db,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Failed to connect to database: {}", e),
            ))
        }
    };

    let page_size = std::cmp::min(query.page_size.unwrap_or(100), 1000);
    let offset = (query.page.unwrap_or(1) - 1) * page_size;

    let table = crate::core::db::schema::table_name("face_detections");

    let count_sql = format!("SELECT COUNT(*) FROM {} WHERE identity_id = $1", table);

    let total: i64 = match sqlx::query_scalar(&count_sql)
        .bind(identity_id)
        .fetch_one(db.pool())
        .await
    {
        Ok(count) => count,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Count error: {}", e),
            ))
        }
    };

    let sql = format!(
        "SELECT id, face_id, file_uuid, frame_number, confidence, bbox, attributes 
         FROM {} 
         WHERE identity_id = $1
         ORDER BY confidence DESC 
         LIMIT $2 OFFSET $3",
        table
    );

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

    let faces: Vec<IdentityFace> = rows
        .into_iter()
        .map(|r| IdentityFace {
            id: r.0,
            face_id: r.1,
            file_uuid: r.2,
            frame_number: r.3,
            confidence: r.4,
            bbox: r.5,
            attributes: r.6,
        })
        .collect();

    Ok(Json(IdentityFacesResponse {
        identity_id,
        faces,
        total,
    }))
}

async fn get_face_thumbnail(
    Path(face_id): Path<i32>,
) -> Result<impl IntoResponse, (StatusCode, String)> {
    let db = match PostgresDb::init().await {
        Ok(db) => db,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Failed to connect to database: {}", e),
            ))
        }
    };

    let table_fd = crate::core::db::schema::table_name("face_detections");
    let table_v = crate::core::db::schema::table_name("videos");

    let sql = format!(
        "SELECT fd.frame_number, fd.bbox, v.file_path, v.fps
         FROM {} fd
         JOIN {} v ON fd.file_uuid = v.uuid
         WHERE fd.id = $1",
        table_fd, table_v
    );

    let row: Option<(i64, Option<serde_json::Value>, String, f64)> = match sqlx::query_as(&sql)
        .bind(face_id)
        .fetch_optional(db.pool())
        .await
    {
        Ok(row) => row,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Query error: {}", e),
            ))
        }
    };

    let (frame_number, bbox_json, file_path, fps) = match row {
        Some(r) => r,
        None => return Err((StatusCode::NOT_FOUND, format!("Face {} not found", face_id))),
    };

    let bbox: Bbox = match bbox_json {
        Some(json) => serde_json::from_value(json).unwrap_or(Bbox {
            x: 0,
            y: 0,
            width: 100,
            height: 100,
        }),
        None => Bbox {
            x: 0,
            y: 0,
            width: 100,
            height: 100,
        },
    };

    let timestamp = frame_number as f64 / fps;

    let crop_filter = format!("crop={}:{}:{}:{}", bbox.width, bbox.height, bbox.x, bbox.y);

    let output = match Command::new("ffmpeg")
        .args(&[
            "-ss",
            &timestamp.to_string(),
            "-i",
            &file_path,
            "-vf",
            &crop_filter,
            "-frames:v",
            "1",
            "-f",
            "image2pipe",
            "-vcodec",
            "mjpeg",
            "-",
        ])
        .output()
    {
        Ok(o) => o,
        Err(e) => {
            return Err((
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("ffmpeg error: {}", e),
            ))
        }
    };

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

    let response = axum::response::Response::builder()
        .status(StatusCode::OK)
        .header(header::CONTENT_TYPE, "image/jpeg")
        .header(header::CACHE_CONTROL, "public, max-age=3600")
        .body(Body::from(output.stdout))
        .map_err(|e| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Response error: {}", e),
            )
        })?;

    Ok(response)
}

#[derive(Debug, Deserialize)]
struct Bbox {
    x: i32,
    y: i32,
    width: i32,
    height: i32,
}