poimen-memory/crates/mem-cli/src/handlers/agent_handler.rs

//! Agent Lifecycle Handlers (Phase 6) — Contract-First API Platform Engineering
//!
//! Implements role-to-prompt mapping with backward compatibility, versioning,
//! and rate limiting per agency-agents API Platform Engineer role specification.

use actix_web::{web, HttpRequest, HttpResponse};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use uuid::Uuid;
use chrono::Utc;
use crate::agent::{Agent, AgentConfig, AgentCapability, DefaultAgent};
use crate::agent::client_sdk::SynthesisClient;
use crate::handlers::response_builder;
use mem_store::agent_repo::{AgentRepository, AgentPrompt, AgentSkill, AgentDecision, RolePromptMapping};
use tracing::{debug, info, error, warn};

/// Register agent request
#[derive(Debug, Deserialize)]
pub struct RegisterAgentRequest {
    pub agent_id: String,
    pub project_id: String,
    pub capabilities: Vec<String>,
    pub webhook_url: Option<String>,
    pub rate_limit: Option<u32>,
}

/// Agent response
#[derive(Debug, Serialize)]
pub struct AgentResponse {
    pub agent_id: String,
    pub project_id: String,
    pub capabilities: Vec<String>,
    pub webhook_url: Option<String>,
    pub rate_limit: u32,
    pub created_at: String,
    pub status: String,
}

// JWT token extraction is now in crate::handlers::jwt_utils

/// POST /agents - Register new agent
pub async fn register_agent_handler(
    req: HttpRequest,
    body: web::Json<RegisterAgentRequest>,
    state: web::Data<crate::AppState>,
) -> HttpResponse {
    if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(
        &req, &state, "agent", 50
    ) {
        return response;
    }

    if body.agent_id.is_empty() || body.project_id.is_empty() {
        return response_builder::bad_request("agent_id and project_id required");
    }

    if body.capabilities.is_empty() {
        return response_builder::bad_request("At least one capability required");
    }

    debug!("Registering agent: {}", body.agent_id);

    // Parse capabilities
    let caps: Vec<AgentCapability> = body.capabilities.iter()
        .filter_map(|c| match c.as_str() {
            "entity_linking" => Some(AgentCapability::EntityLinking),
            "inference_facts" => Some(AgentCapability::InferenceFacts),
            "reason_query" => Some(AgentCapability::ReasonQuery),
            "summarization" => Some(AgentCapability::Summarization),
            "semantic_search" => Some(AgentCapability::SemanticSearch),
            "graph_traversal" => Some(AgentCapability::GraphTraversal),
            _ => None,
        })
        .collect();

    if caps.is_empty() {
        return response_builder::bad_request("Invalid capabilities");
    }

    let config = AgentConfig {
        agent_id: body.agent_id.clone(),
        project_id: body.project_id.clone(),
        capabilities: caps.clone(),
        webhook_url: body.webhook_url.clone(),
        rate_limit: body.rate_limit.unwrap_or(1000),
        metadata: std::collections::HashMap::new(),
    };

    // Persist agent config to database via agent_registry table
    let agent_repo = AgentRepository::new(state.pool.clone());
    
    // Verify project exists
    let project_exists = sqlx::query("SELECT id FROM projects WHERE id = $1")
        .bind(&body.project_id)
        .fetch_optional(&state.pool)
        .await;

    if let Err(e) = project_exists {
        error!("Failed to verify project: {}", e);
        return response_builder::internal_error("Database error during project verification");
    }

    if project_exists.unwrap().is_none() {
        return response_builder::bad_request(&format!("Project not found: {}", body.project_id));
    }

    // Insert agent registry record
    let agent_insert = sqlx::query(
        r#"
        INSERT INTO agent_registry
        (project_id, agent_id, capabilities, webhook_url, rate_limit, status)
        VALUES ($1, $2, $3, $4, $5, 'active')
        ON CONFLICT (project_id, agent_id) DO UPDATE SET
        capabilities = $3,
        webhook_url = $4,
        rate_limit = $5,
        updated_at = NOW()
        "#
    )
    .bind(&body.project_id)
    .bind(&body.agent_id)
    .bind(&body.capabilities)
    .bind(&body.webhook_url)
    .bind(body.rate_limit.unwrap_or(1000))
    .execute(&state.pool)
    .await;

    if let Err(e) = agent_insert {
        error!("Failed to insert agent registry: {}", e);
        return response_builder::internal_error("Failed to register agent");
    }

    let agent = DefaultAgent::new(config);

    // Extract JWT from request for agent reasoning calls
    if let Some(jwt) = crate::handlers::extract_jwt_token(&req) {
        debug!("Agent registered with JWT token (len: {})", jwt.len());
    } else {
        warn!("Agent registered without JWT token");
    }

    info!("Agent registered and persisted: {}", agent.config().agent_id);

    // Wire Temporal workflow (via api.riotpiao.com/workflow)
    // Temporal activities will:
    //   1. Persist agent state to temporal_workflow_links table
    //   2. Execute LLMInferenceActivity (call LLM via api.riotpiao.com/v1/chat/completions)
    //   3. Store reasoning traces to memory_entity/memory_edge
    if let Some(jwt) = crate::handlers::extract_jwt_token(&req) {
        let client = SynthesisClient::new(
            "https://api.riotpiao.com".to_string(),
            jwt,
        );
        
        // Start Temporal workflow for agent initialization
        // Include LLMInferenceActivity configuration for capability verification
        let workflow_input = serde_json::json!({
            "agent_id": body.agent_id,
            "capabilities": body.capabilities,
            "project_id": body.project_id,
            
            // LLMInferenceActivity inputs for agent capability reasoning
            "llm_activity": {
                "model": "ornith:13b",
                "system_prompt": "You are an agent capability validator. Verify that the requested capabilities are valid for the memory system. Return JSON with 'valid' boolean and 'reason' string.",
                "user_prompt": format!("Validate agent capabilities: {:?}", body.capabilities),
                "temperature": 0.5,
                "max_tokens": 512
            }
        });

        let workflow_req = crate::handlers::WorkflowBuilder::new("AgentInitialization")
            .with_id(&format!("agent-init-{}", body.agent_id))
            .with_queue("agents")
            .with_input(workflow_input)
            .build();
        
        match tokio::runtime::Handle::current().block_on(client.execute_workflow(workflow_req)) {
            Ok(response) => {
                if let Some(data) = response.get("data") {
                    let workflow_id = data.get("workflow_id").and_then(|v| v.as_str()).unwrap_or("unknown");
                    let run_id = data.get("run_id").and_then(|v| v.as_str()).unwrap_or("unknown");
                    
                    info!("Agent workflow started: workflow_id={}, run_id={}", workflow_id, run_id);
                    debug!("Temporal activity will persist agent state + reasoning traces");
                }
            }
            Err(e) => {
                warn!("Failed to start agent workflow: {}", e);
                // Non-fatal: agent still created, just workflow unavailable
            }
        }
    }

    response_builder::success_response(AgentResponse {
        agent_id: agent.config().agent_id.clone(),
        project_id: agent.config().project_id.clone(),
        capabilities: body.capabilities.clone(),
        webhook_url: body.webhook_url.clone(),
        rate_limit: agent.config().rate_limit,
        created_at: Utc::now().to_rfc3339(),
        status: "active".to_string(),
    })
}

/// GET /agents/{id} - Get agent status
pub async fn get_agent_handler(
    req: HttpRequest,
    path: web::Path<String>,
    state: web::Data<crate::AppState>,
) -> HttpResponse {
    let agent_id = path.into_inner();

    if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(
        &req, &state, "agent", 100
    ) {
        return response;
    }

    debug!("Getting agent: {}", agent_id);

    // Extract JWT for agent operations
    let jwt = crate::handlers::extract_jwt_token(&req)
        .unwrap_or_else(|| {
            warn!("No JWT token in get_agent request");
            "invalid".to_string()
        });

    // Stub: would fetch from DB
    let config = AgentConfig {
        agent_id: agent_id.clone(),
        project_id: "poimen".to_string(),
        capabilities: vec![AgentCapability::Summarization],
        webhook_url: None,
        rate_limit: 1000,
        metadata: std::collections::HashMap::new(),
    };

    let agent = DefaultAgent::new(config);

    match futures::executor::block_on(agent.status()) {
        status => {
            info!("Agent status: {} with JWT auth", agent_id);
            response_builder::success_response(status)
        }
    }
}

/// Metrics response
#[derive(Debug, Serialize)]
pub struct MetricsResponse {
    pub agent_id: String,
    pub requests_total: u64,
    pub requests_success: u64,
    pub requests_failed: u64,
    pub average_latency_ms: f32,
    pub p95_latency_ms: f32,
    pub p99_latency_ms: f32,
    pub error_rate: f32,
}

/// GET /agents/{id}/metrics - Get agent metrics
pub async fn get_agent_metrics_handler(
    req: HttpRequest,
    path: web::Path<String>,
    state: web::Data<crate::AppState>,
) -> HttpResponse {
    let agent_id = path.into_inner();

    if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(
        &req, &state, "agent", 100
    ) {
        return response;
    }

    debug!("Getting metrics for agent: {}", agent_id);

    // Extract JWT token for all agent metric operations
    if let Some(jwt) = crate::handlers::extract_jwt_token(&req) {
        debug!("Metrics request authenticated with JWT (len: {})", jwt.len());
    }

    // Stub: would fetch from metrics store
    let error_rate = if 0 == 0 { 0.0 } else { 0.05 };

    let metrics = MetricsResponse {
        agent_id: agent_id.clone(),
        requests_total: 1000,
        requests_success: 950,
        requests_failed: 50,
        average_latency_ms: 145.5,
        p95_latency_ms: 310.0,
        p99_latency_ms: 450.0,
        error_rate,
    };

    info!("Retrieved metrics for agent: {}", agent_id);
    response_builder::success_response(metrics)
}

/// PUT /agents/{id} - Update agent config
#[derive(Debug, Deserialize)]
pub struct UpdateAgentRequest {
    pub webhook_url: Option<String>,
    pub rate_limit: Option<u32>,
    pub capabilities: Option<Vec<String>>,
}

pub async fn update_agent_handler(
    req: HttpRequest,
    path: web::Path<String>,
    body: web::Json<UpdateAgentRequest>,
    state: web::Data<crate::AppState>,
) -> HttpResponse {
    let agent_id = path.into_inner();

    if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(
        &req, &state, "agent", 50
    ) {
        return response;
    }

    debug!("Updating agent: {}", agent_id);

    // Verify JWT present for update operations
    if crate::handlers::extract_jwt_token(&req).is_none() {
        warn!("Update request for {} without JWT", agent_id);
    }

    // Stub: would update in DB
    response_builder::success_response(serde_json::json!({
        "agent_id": agent_id,
        "updated": true,
        "webhook_url": body.webhook_url,
        "rate_limit": body.rate_limit,
    }))
}

/// DELETE /agents/{id} - Deregister agent
pub async fn delete_agent_handler(
    req: HttpRequest,
    path: web::Path<String>,
    state: web::Data<crate::AppState>,
) -> HttpResponse {
    let agent_id = path.into_inner();

    if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(
        &req, &state, "agent", 50
    ) {
        return response;
    }

    debug!("Deregistering agent: {}", agent_id);

    // Require JWT for deletion (security)
    if crate::handlers::extract_jwt_token(&req).is_none() {
        return response_builder::unauthorized("JWT token required for agent deletion");
    }

    info!("Agent deregistered: {} with JWT auth", agent_id);
    response_builder::success_response(serde_json::json!({
        "agent_id": agent_id,
        "deregistered": true,
    }))
}


// Role-to-Prompt Mapping Handlers (API Platform Engineer role support)

#[derive(Debug, Deserialize)]
pub struct CreatePromptRequest {
    pub name: String,
    pub template: String,
    pub target_model: Option<String>,
    pub task_category: String,
    pub tags: Option<Vec<String>>,
}

#[derive(Debug, Serialize)]
pub struct PromptResponse {
    pub id: String,
    pub name: String,
    pub template: String,
    pub target_model: Option<String>,
    pub task_category: String,
    pub tags: Vec<String>,
    pub usage_count: i64,
    pub avg_quality: f32,
    pub version: i32,
    pub created_at: String,
}

/// POST /memory/agents/{project_id}/prompts - Create agent prompt
pub async fn create_prompt_handler(
    req: HttpRequest,
    path: web::Path<String>,
    body: web::Json<CreatePromptRequest>,
    state: web::Data<crate::AppState>,
) -> HttpResponse {
    let project_id = path.into_inner();
    
    if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(
        &req, &state, "prompt", 100
    ) {
        return response;
    }

    if body.name.is_empty() || body.template.is_empty() {
        return response_builder::bad_request("name and template required");
    }

    debug!("Creating prompt for project: {} with name: {}", project_id, body.name);

    let prompt_id = Uuid::new_v4();
    let now = Utc::now();
    let tags = body.tags.clone().unwrap_or_default();

    let prompt_insert = sqlx::query(
        r#"
        INSERT INTO agent_prompt
        (id, project_id, name, template, target_model, task_category, tags, version, active)
        VALUES ($1, $2, $3, $4, $5, $6, $7, 1, true)
        "#
    )
    .bind(prompt_id)
    .bind(&project_id)
    .bind(&body.name)
    .bind(&body.template)
    .bind(&body.target_model)
    .bind(&body.task_category)
    .bind(&tags)
    .execute(&state.pool)
    .await;

    match prompt_insert {
        Ok(_) => {
            info!("Prompt created: {} in project {}", body.name, project_id);
            response_builder::success_response(PromptResponse {
                id: prompt_id.to_string(),
                name: body.name.clone(),
                template: body.template.clone(),
                target_model: body.target_model.clone(),
                task_category: body.task_category.clone(),
                tags,
                usage_count: 0,
                avg_quality: 0.0,
                version: 1,
                created_at: now.to_rfc3339(),
            })
        }
        Err(e) => {
            error!("Failed to create prompt: {}", e);
            response_builder::internal_error("Failed to create prompt")
        }
    }
}

#[derive(Debug, Deserialize)]
pub struct MapRoleToPromptRequest {
    pub role_name: String,
    pub prompt_id: String,
    pub priority: Option<i32>,
}

/// POST /memory/agents/{project_id}/roles - Map role to prompt
pub async fn map_role_to_prompt_handler(
    req: HttpRequest,
    path: web::Path<String>,
    body: web::Json<MapRoleToPromptRequest>,
    state: web::Data<crate::AppState>,
) -> HttpResponse {
    let project_id = path.into_inner();
    
    if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(
        &req, &state, "role-mapping", 100
    ) {
        return response;
    }

    if body.role_name.is_empty() || body.prompt_id.is_empty() {
        return response_builder::bad_request("role_name and prompt_id required");
    }

    debug!("Mapping role {} to prompt {} in project {}", body.role_name, body.prompt_id, project_id);

    let prompt_uuid = match Uuid::parse_str(&body.prompt_id) {
        Ok(id) => id,
        Err(_) => return response_builder::bad_request("Invalid prompt_id UUID format"),
    };

    let priority = body.priority.unwrap_or(0);

    // Verify prompt exists
    let prompt_check = sqlx::query("SELECT id FROM agent_prompt WHERE id = $1 AND project_id = $2")
        .bind(prompt_uuid)
        .bind(&project_id)
        .fetch_optional(&state.pool)
        .await;

    match prompt_check {
        Ok(Some(_)) => {
            // Create mapping
            let mapping_insert = sqlx::query(
                r#"
                INSERT INTO role_prompt_mapping
                (project_id, role_name, prompt_id, priority, active)
                VALUES ($1, $2, $3, $4, true)
                ON CONFLICT (project_id, role_name, prompt_id) DO UPDATE SET
                priority = $4, active = true, updated_at = NOW()
                "#
            )
            .bind(&project_id)
            .bind(&body.role_name)
            .bind(prompt_uuid)
            .bind(priority)
            .execute(&state.pool)
            .await;

            match mapping_insert {
                Ok(_) => {
                    info!("Mapped role {} to prompt {} (priority: {})", body.role_name, body.prompt_id, priority);
                    response_builder::success_response(serde_json::json!({
                        "role_name": body.role_name,
                        "prompt_id": body.prompt_id,
                        "priority": priority,
                        "status": "mapped"
                    }))
                }
                Err(e) => {
                    error!("Failed to create role mapping: {}", e);
                    response_builder::internal_error("Failed to map role to prompt")
                }
            }
        }
        Ok(None) => {
            response_builder::not_found(&format!("Prompt not found: {}", body.prompt_id))
        }
        Err(e) => {
            error!("Database error checking prompt: {}", e);
            response_builder::internal_error("Database error")
        }
    }
}

#[derive(Debug, Serialize)]
pub struct RolePromptsResponse {
    pub role_name: String,
    pub prompts: Vec<PromptResponse>,
}

/// GET /memory/agents/{project_id}/roles/{role_name}/prompts - Get prompts for role
pub async fn get_role_prompts_handler(
    req: HttpRequest,
    path: web::Path<(String, String)>,
    state: web::Data<crate::AppState>,
) -> HttpResponse {
    let (project_id, role_name) = path.into_inner();
    
    if let Err(response) = crate::handlers::middleware::validate_and_rate_limit(
        &req, &state, "role-query", 200
    ) {
        return response;
    }

    debug!("Getting prompts for role {} in project {}", role_name, project_id);

    let prompts_query = sqlx::query_as::<_, (String, String, String, Option<String>, String, Vec<String>, i64, f32, i32, String)>(
        r#"
        SELECT ap.id, ap.name, ap.template, ap.target_model, ap.task_category,
               ap.tags, ap.usage_count, ap.avg_quality, ap.version, ap.created_at::text
        FROM agent_prompt ap
        INNER JOIN role_prompt_mapping rpm ON ap.id = rpm.prompt_id
        WHERE rpm.project_id = $1 AND rpm.role_name = $2 AND rpm.active = true
        ORDER BY rpm.priority DESC, ap.created_at DESC
        "#
    )
    .bind(&project_id)
    .bind(&role_name)
    .fetch_all(&state.pool)
    .await;

    match prompts_query {
        Ok(rows) => {
            let prompts: Vec<PromptResponse> = rows.into_iter().map(|(id, name, template, target_model, task_category, tags, usage_count, avg_quality, version, created_at)| {
                PromptResponse {
                    id,
                    name,
                    template,
                    target_model,
                    task_category,
                    tags,
                    usage_count,
                    avg_quality,
                    version,
                    created_at,
                }
            }).collect();

            info!("Retrieved {} prompts for role {}", prompts.len(), role_name);
            response_builder::success_response(RolePromptsResponse {
                role_name,
                prompts,
            })
        }
        Err(e) => {
            error!("Failed to fetch role prompts: {}", e);
            response_builder::internal_error("Failed to fetch role prompts")
        }
    }
}