momentry_core/src/core/tkg/service.rs

use crate::core::db::schema;
use crate::core::db::PostgresDb;
use crate::core::tkg::log::TkgLogger;
use crate::core::tkg::models::{TkgBuildStats, TkgOperationLog};
use anyhow::Result;
use sqlx::PgPool;
use std::sync::Arc;

pub struct TkgService {
    db: Arc<PostgresDb>,
}

impl TkgService {
    pub fn new(db: Arc<PostgresDb>) -> Self {
        Self { db }
    }

    pub async fn build(&self, file_uuid: &str, output_dir: &str) -> Result<TkgBuildStats> {
        let log_id = TkgLogger::start_operation(self.db.pool(), file_uuid, "build").await?;
        let mut logger = TkgLogger::new(self.db.pool().clone());
        logger.log_id = Some(log_id);

        let redis = crate::core::db::RedisClient::new()
            .ok()
            .map(|r| std::sync::Arc::new(r));
        let result =
            crate::core::processor::tkg::build_tkg(&self.db, file_uuid, output_dir, redis).await;

        match result {
            Ok(r) => {
                let stats = TkgBuildStats {
                    face_track_nodes: r.face_track_nodes as i32,
                    gaze_track_nodes: r.gaze_track_nodes as i32,
                    lip_track_nodes: r.lip_track_nodes as i32,
                    text_region_nodes: r.text_region_nodes as i32,
                    appearance_trace_nodes: r.appearance_trace_nodes as i32,
                    accessory_nodes: r.accessory_nodes as i32,
                    object_nodes: r.object_nodes as i32,
                    hand_nodes: r.hand_nodes as i32,
                    speaker_nodes: r.speaker_nodes as i32,
                    co_occurrence_edges: r.co_occurrence_edges as i32,
                    speaker_face_edges: r.speaker_face_edges as i32,
                    face_face_edges: r.face_face_edges as i32,
                    mutual_gaze_edges: r.mutual_gaze_edges as i32,
                    lip_sync_edges: r.lip_sync_edges as i32,
                    has_appearance_edges: r.has_appearance_edges as i32,
                    wears_edges: r.wears_edges as i32,
                    hand_object_edges: r.hand_object_edges as i32,
                };
                logger.complete_operation(None).await?;
                tracing::info!(
                    "[TKG-Service] Build completed for {}: {} nodes, {} edges",
                    file_uuid,
                    stats.total_nodes(),
                    stats.total_edges()
                );
                Ok(stats)
            }
            Err(e) => {
                logger.complete_operation(Some(&e.to_string())).await?;
                tracing::error!("[TKG-Service] Build failed for {}: {}", file_uuid, e);
                Err(e)
            }
        }
    }

    pub async fn rebuild(
        &self,
        file_uuid: &str,
        output_dir: &str,
        force: bool,
    ) -> Result<TkgBuildStats> {
        let operation = if force { "rebuild_force" } else { "rebuild" };
        let log_id = TkgLogger::start_operation(self.db.pool(), file_uuid, operation).await?;

        if force {
            self.delete_internal(file_uuid).await?;
        }

        let result = self.build(file_uuid, output_dir).await;

        // Update the original log entry
        if let Some(id) = Some(log_id) {
            let table = schema::table_name("tkg_operation_log");
            let status = if result.is_ok() {
                "completed"
            } else {
                "failed"
            };
            let error = result.as_ref().err().map(|e| e.to_string());
            sqlx::query(&format!(
                "UPDATE {table} SET status = $1, error_message = $2, completed_at = NOW() WHERE id = $3"
            ))
            .bind(status)
            .bind(error)
            .bind(id)
            .execute(self.db.pool())
            .await?;
        }

        result
    }

    pub async fn delete(&self, file_uuid: &str) -> Result<()> {
        let log_id = TkgLogger::start_operation(self.db.pool(), file_uuid, "delete").await?;

        let result = self.delete_internal(file_uuid).await;

        match result {
            Ok(_) => {
                TkgLogger::new(self.db.pool().clone())
                    .complete_operation(None)
                    .await?;
                Ok(())
            }
            Err(e) => {
                TkgLogger::new(self.db.pool().clone())
                    .complete_operation(Some(&e.to_string()))
                    .await?;
                Err(e)
            }
        }
    }

    async fn delete_internal(&self, file_uuid: &str) -> Result<()> {
        let nodes_table = schema::table_name("tkg_nodes");
        let edges_table = schema::table_name("tkg_edges");

        // Delete edges first (foreign key constraint)
        let edges_deleted = sqlx::query(&format!("DELETE FROM {edges_table} WHERE file_uuid = $1"))
            .bind(file_uuid)
            .execute(self.db.pool())
            .await?;

        // Delete nodes
        let nodes_deleted = sqlx::query(&format!("DELETE FROM {nodes_table} WHERE file_uuid = $1"))
            .bind(file_uuid)
            .execute(self.db.pool())
            .await?;

        tracing::info!(
            "[TKG-Service] Deleted {} nodes and {} edges for {}",
            nodes_deleted.rows_affected(),
            edges_deleted.rows_affected(),
            file_uuid
        );

        Ok(())
    }

    pub async fn get_operations(&self, file_uuid: &str) -> Result<Vec<TkgOperationLog>> {
        TkgLogger::get_operations(self.db.pool(), file_uuid).await
    }

    pub async fn delete_tkg_and_logs(&self, file_uuid: &str) -> Result<()> {
        self.delete_internal(file_uuid).await?;
        TkgLogger::delete_operations(self.db.pool(), file_uuid).await?;
        Ok(())
    }
}