tori/src/sink.rs

use std::collections::HashMap;
use std::sync::Arc;
use std::sync::atomic::{AtomicU16, Ordering};
use serde::{Deserialize, Serialize};
use sqlx::sqlite::SqlitePool;
use tokio::sync::{RwLock, broadcast};

use crate::agent::{PlanStepInfo, WsMessage, ServiceInfo};
use crate::state::{AgentState, Artifact};

/// All updates produced by the agent loop. This is the single output interface
/// that decouples the agent logic from DB persistence and WebSocket broadcasting.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "kind")]
pub enum AgentUpdate {
    PlanUpdate { workflow_id: String, steps: Vec<PlanStepInfo> },
    WorkflowStatus { workflow_id: String, status: String, #[serde(default)] reason: String },
    Activity { workflow_id: String, activity: String },
    ExecutionLog { workflow_id: String, step_order: i32, tool_name: String, tool_input: String, output: String, status: String },
    LlmCallLog { workflow_id: String, step_order: i32, phase: String, messages_count: i32, tools_count: i32, tool_calls: String, text_response: String, prompt_tokens: Option<u32>, completion_tokens: Option<u32>, latency_ms: i64 },
    StateSnapshot { workflow_id: String, step_order: i32, state: AgentState },
    WorkflowComplete { workflow_id: String, status: String, #[serde(default)] reason: String },
    ArtifactSave { workflow_id: String, step_order: i32, artifact: Artifact },
    RequirementUpdate { workflow_id: String, requirement: String },
    /// base64-encoded file content
    FileSync { project_id: String, path: String, data_b64: String },
    Error { message: String },
}

/// Manages local services (start_service / stop_service tools).
pub struct ServiceManager {
    pub services: RwLock<HashMap<String, ServiceInfo>>,
    next_port: AtomicU16,
}

impl ServiceManager {
    pub fn new(start_port: u16) -> Arc<Self> {
        Arc::new(Self {
            services: RwLock::new(HashMap::new()),
            next_port: AtomicU16::new(start_port),
        })
    }

    pub fn allocate_port(&self) -> u16 {
        self.next_port.fetch_add(1, Ordering::Relaxed)
    }
}

/// Helper: broadcast if sender is available.
fn bcast(tx: Option<&broadcast::Sender<WsMessage>>, msg: WsMessage) {
    if let Some(tx) = tx { let _ = tx.send(msg); }
}

/// Process a single AgentUpdate: persist to DB and broadcast to frontend.
pub async fn handle_single_update(
    update: &AgentUpdate,
    pool: &SqlitePool,
    broadcast_tx: Option<&broadcast::Sender<WsMessage>>,
) {
    match update {
        AgentUpdate::PlanUpdate { workflow_id, steps } => {
            bcast(broadcast_tx, WsMessage::PlanUpdate { workflow_id: workflow_id.clone(), steps: steps.clone() });
        }
        AgentUpdate::WorkflowStatus { workflow_id, status, reason } => {
            let _ = sqlx::query("UPDATE workflows SET status = ?, status_reason = ? WHERE id = ?")
                .bind(status).bind(reason).bind(workflow_id).execute(pool).await;
            bcast(broadcast_tx, WsMessage::WorkflowStatusUpdate { workflow_id: workflow_id.clone(), status: status.clone() });
        }
        AgentUpdate::Activity { workflow_id, activity } => {
            bcast(broadcast_tx, WsMessage::ActivityUpdate { workflow_id: workflow_id.clone(), activity: activity.clone() });
        }
        AgentUpdate::ExecutionLog { workflow_id, step_order, tool_name, tool_input, output, status } => {
            let id = uuid::Uuid::new_v4().to_string();
            let _ = sqlx::query(
                "INSERT INTO execution_log (id, workflow_id, step_order, tool_name, tool_input, output, status, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, datetime('now'))"
            ).bind(&id).bind(workflow_id).bind(step_order).bind(tool_name).bind(tool_input).bind(output).bind(status)
            .execute(pool).await;
            bcast(broadcast_tx, WsMessage::StepStatusUpdate { step_id: id, status: status.clone(), output: output.clone() });
        }
        AgentUpdate::LlmCallLog { workflow_id, step_order, phase, messages_count, tools_count, tool_calls, text_response, prompt_tokens, completion_tokens, latency_ms } => {
            let id = uuid::Uuid::new_v4().to_string();
            let _ = sqlx::query(
                "INSERT INTO llm_call_log (id, workflow_id, step_order, phase, messages_count, tools_count, tool_calls, text_response, prompt_tokens, completion_tokens, latency_ms, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, datetime('now'))"
            ).bind(&id).bind(workflow_id).bind(step_order).bind(phase).bind(messages_count).bind(tools_count)
            .bind(tool_calls).bind(text_response).bind(prompt_tokens.map(|v| v as i32)).bind(completion_tokens.map(|v| v as i32)).bind(*latency_ms as i32)
            .execute(pool).await;
            let entry = crate::db::LlmCallLogEntry {
                id, workflow_id: workflow_id.clone(), step_order: *step_order, phase: phase.clone(),
                messages_count: *messages_count, tools_count: *tools_count, tool_calls: tool_calls.clone(),
                text_response: text_response.clone(), prompt_tokens: prompt_tokens.map(|v| v as i32),
                completion_tokens: completion_tokens.map(|v| v as i32), latency_ms: *latency_ms as i32,
                created_at: String::new(),
            };
            bcast(broadcast_tx, WsMessage::LlmCallLog { workflow_id: workflow_id.clone(), entry });
        }
        AgentUpdate::StateSnapshot { workflow_id, step_order, state } => {
            let id = uuid::Uuid::new_v4().to_string();
            let json = serde_json::to_string(state).unwrap_or_default();
            let _ = sqlx::query(
                "INSERT INTO agent_state_snapshots (id, workflow_id, step_order, state_json, created_at) VALUES (?, ?, ?, ?, datetime('now'))"
            ).bind(&id).bind(workflow_id).bind(step_order).bind(&json).execute(pool).await;
        }
        AgentUpdate::WorkflowComplete { workflow_id, status, reason } => {
            let _ = sqlx::query("UPDATE workflows SET status = ?, status_reason = ? WHERE id = ?")
                .bind(status).bind(reason).bind(workflow_id).execute(pool).await;
            bcast(broadcast_tx, WsMessage::WorkflowStatusUpdate { workflow_id: workflow_id.clone(), status: status.clone() });
        }
        AgentUpdate::ArtifactSave { workflow_id, step_order, artifact } => {
            let id = uuid::Uuid::new_v4().to_string();
            let _ = sqlx::query(
                "INSERT INTO step_artifacts (id, workflow_id, step_order, name, path, artifact_type, description) VALUES (?, ?, ?, ?, ?, ?, ?)"
            ).bind(&id).bind(workflow_id).bind(step_order).bind(&artifact.name).bind(&artifact.path)
            .bind(&artifact.artifact_type).bind(&artifact.description).execute(pool).await;
        }
        AgentUpdate::RequirementUpdate { workflow_id, requirement } => {
            let _ = sqlx::query("UPDATE workflows SET requirement = ? WHERE id = ?")
                .bind(requirement).bind(workflow_id).execute(pool).await;
            bcast(broadcast_tx, WsMessage::RequirementUpdate { workflow_id: workflow_id.clone(), requirement: requirement.clone() });
        }
        AgentUpdate::FileSync { project_id, path, data_b64 } => {
            use base64::Engine;
            let base = format!("/app/data/workspaces/{}", project_id);
            let full = std::path::Path::new(&base).join(path);
            if let Some(parent) = full.parent() {
                let _ = tokio::fs::create_dir_all(parent).await;
            }
            if let Ok(bytes) = base64::engine::general_purpose::STANDARD.decode(data_b64) {
                let _ = tokio::fs::write(&full, &bytes).await;
            }
        }
        AgentUpdate::Error { message } => {
            bcast(broadcast_tx, WsMessage::Error { message: message.clone() });
        }
    }
}