use crate::core::db::PostgresDb; use crate::core::db::schema; 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, } impl TkgService { pub fn new(db: Arc) -> Self { Self { db } } pub async fn build(&self, file_uuid: &str, output_dir: &str) -> Result { 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 { 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> { 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(()) } }