poimen-memory/crates/mem-cli/src/gateway_queue_adapter.rs

//! M8.2 — Gateway Queue Adapter
//!
//! Calls SQS via `api.riotpiao.com` gateway with JWT authentication.
//! Uses X-Service routing to reach kmsvc backend.

use crate::queue_adapter::{QueueAdapter, QueueMessage, QueueStats};
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use uuid::Uuid;
use std::sync::Arc;

/// Token provider trait (async)
#[async_trait]
pub trait TokenProvider: Send + Sync {
    async fn token(&self) -> Result<String>;
}

/// Static JWT token provider (for testing)
pub struct StaticTokenProvider {
    token: String,
}

impl StaticTokenProvider {
    pub fn new(token: String) -> Self {
        Self { token }
    }
}

#[async_trait]
impl TokenProvider for StaticTokenProvider {
    async fn token(&self) -> Result<String> {
        Ok(self.token.clone())
    }
}

/// Authentik token provider (production)
pub struct AuthentikTokenProvider {
    issuer: String,
    client_id: String,
    client_secret: String,
    http_client: reqwest::Client,
    cached_token: Arc<tokio::sync::RwLock<CachedToken>>,
}

#[derive(Clone)]
struct CachedToken {
    token: Option<String>,
    expires_at: i64,
}

impl AuthentikTokenProvider {
    pub fn new(issuer: String, client_id: String, client_secret: String) -> Self {
        Self {
            issuer,
            client_id,
            client_secret,
            http_client: reqwest::Client::new(),
            cached_token: Arc::new(tokio::sync::RwLock::new(CachedToken {
                token: None,
                expires_at: 0,
            })),
        }
    }

    async fn refresh_token(&self) -> Result<String> {
        let token_url = format!("{}/application/o/token/", self.issuer);

        let params = [
            ("grant_type", "client_credentials"),
            ("client_id", &self.client_id),
            ("client_secret", &self.client_secret),
            ("scope", "openid"),
        ];

        let resp = self
            .http_client
            .post(&token_url)
            .form(&params)
            .send()
            .await?;

        if !resp.status().is_success() {
            return Err(anyhow!("Failed to get token from Authentik: {}", resp.status()));
        }

        let token_resp: serde_json::Value = resp.json().await?;
        let token = token_resp["access_token"]
            .as_str()
            .ok_or_else(|| anyhow!("No access_token in Authentik response"))?
            .to_string();

        let expires_in = token_resp["expires_in"]
            .as_i64()
            .unwrap_or(3600);
        let expires_at = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64 + expires_in;

        let mut cached = self.cached_token.write().await;
        cached.token = Some(token.clone());
        cached.expires_at = expires_at;

        tracing::debug!("Token refreshed from Authentik, expires in {}s", expires_in);

        Ok(token)
    }
}

#[async_trait]
impl TokenProvider for AuthentikTokenProvider {
    async fn token(&self) -> Result<String> {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_secs() as i64;

        // Check cache
        {
            let cached = self.cached_token.read().await;
            if let Some(token) = cached.token.as_ref() {
                if now < cached.expires_at - 60 {
                    return Ok(token.clone());
                }
            }
        }

        // Refresh
        self.refresh_token().await
    }
}

/// SQS SendMessage request
#[derive(Debug, Serialize)]
struct SendMessageRequest {
    #[serde(rename = "messageBody")]
    message_body: String,
    #[serde(rename = "messageAttributes")]
    message_attributes: MessageAttributes,
    #[serde(rename = "delaySeconds")]
    delay_seconds: i32,
}

/// SQS SendMessage response
#[derive(Debug, Deserialize)]
struct SendMessageResponse {
    #[serde(rename = "messageId")]
    message_id: String,
}

/// SQS ReceiveMessage response
#[derive(Debug, Deserialize)]
struct ReceiveMessageResponse {
    messages: Option<Vec<SqsMessage>>,
}

/// SQS Message from ReceiveMessage response
#[derive(Debug, Deserialize)]
struct SqsMessage {
    #[serde(rename = "messageId")]
    message_id: String,
    #[serde(rename = "receiptHandle")]
    receipt_handle: String,
    body: String,
    attributes: Option<std::collections::HashMap<String, String>>,
    #[serde(rename = "receiveCount")]
    receive_count: i32,
}

/// SQS DeleteMessage request
#[derive(Debug, Serialize)]
struct DeleteMessageRequest {
    #[serde(rename = "receiptHandle")]
    receipt_handle: String,
}

/// Message attributes wrapper
#[derive(Debug, Serialize)]
struct MessageAttributes {
    values: std::collections::HashMap<String, String>,
}

/// Gateway Queue Adapter
///
/// Routes through api.riotpiao.com gateway to kmsvc backend.
pub struct GatewayQueueAdapter {
    gateway_url: String,
    token_source: Arc<dyn TokenProvider>,
    http_client: reqwest::Client,
    default_queue_prefix: String,
}

impl GatewayQueueAdapter {
    /// Create with static token (testing)
    pub fn with_static_token(gateway_url: String, token: String) -> Self {
        Self {
            gateway_url,
            token_source: Arc::new(StaticTokenProvider::new(token)),
            http_client: reqwest::Client::new(),
            default_queue_prefix: "poimen-chunks".to_string(),
        }
    }

    /// Create with Authentik provider (production)
    pub fn with_authentik(
        gateway_url: String,
        issuer: String,
        client_id: String,
        client_secret: String,
    ) -> Self {
        Self {
            gateway_url,
            token_source: Arc::new(AuthentikTokenProvider::new(issuer, client_id, client_secret)),
            http_client: reqwest::Client::new(),
            default_queue_prefix: "poimen-chunks".to_string(),
        }
    }

    fn queue_name(&self, project: &str) -> String {
        format!("{}-{}", self.default_queue_prefix, project)
    }
}

#[async_trait]
impl QueueAdapter for GatewayQueueAdapter {
    async fn send_chunk(
        &self,
        chunk_id: Uuid,
        body: String,
        project: String,
        attributes: std::collections::HashMap<String, String>,
    ) -> Result<String> {
        let token = self.token_source.token().await?;

        // Base64 encode body
        let encoded_body = base64::encode(body.as_bytes());

        // Build request
        let mut attrs = attributes;
        attrs.insert("chunk_id".to_string(), chunk_id.to_string());
        attrs.insert("project".to_string(), project.clone());

        let req = SendMessageRequest {
            message_body: encoded_body,
            message_attributes: MessageAttributes { values: attrs },
            delay_seconds: 0,
        };

        let resp = self
            .http_client
            .post(&self.gateway_url)
            .header("X-Service", "sqs")
            .header("Authorization", format!("Bearer {}", token))
            .header("Content-Type", "application/json")
            .json(&req)
            .send()
            .await?;

        if !resp.status().is_success() {
            let status = resp.status();
            let error = resp.text().await.unwrap_or_default();
            return Err(anyhow!("SendMessage failed: {} {}", status, error));
        }

        let sqs_resp: SendMessageResponse = resp.json().await?;

        tracing::debug!(
            "Chunk queued via gateway: message_id={}, chunk_id={}, project={}",
            sqs_resp.message_id, chunk_id, project
        );

        Ok(sqs_resp.message_id)
    }

    async fn receive_chunks(
        &self,
        max_messages: i32,
        visibility_timeout_secs: i32,
        project: Option<&str>,
    ) -> Result<Vec<QueueMessage>> {
        let token = self.token_source.token().await?;
        let project = project.unwrap_or("default");
        let max = max_messages.min(10).max(1);

        // Build query string
        let queue_name = self.queue_name(project);
        let query = format!(
            "X-Service=sqs&queue={}&maxNumberOfMessages={}&waitTimeSeconds=20&visibilityTimeoutSeconds={}",
            urlencoding::encode(&queue_name),
            max,
            visibility_timeout_secs
        );

        let resp = self
            .http_client
            .get(&format!("{}?{}", self.gateway_url, query))
            .header("Authorization", format!("Bearer {}", token))
            .send()
            .await?;

        if !resp.status().is_success() {
            let status = resp.status();
            let error = resp.text().await.unwrap_or_default();
            return Err(anyhow!("ReceiveMessage failed: {} {}", status, error));
        }

        let sqs_resp: ReceiveMessageResponse = resp.json().await?;

        let mut messages = Vec::new();
        if let Some(sqs_msgs) = sqs_resp.messages {
            for msg in sqs_msgs {
                // Decode body from base64
                let body_bytes = base64::decode(msg.body.as_bytes())?;
                let body = String::from_utf8(body_bytes)?;

                let chunk_id = msg
                    .attributes
                    .as_ref()
                    .and_then(|a| a.get("chunk_id"))
                    .and_then(|s| Uuid::parse_str(s).ok())
                    .unwrap_or_else(Uuid::nil);

                messages.push(QueueMessage {
                    message_id: msg.message_id,
                    chunk_id,
                    body,
                    receive_count: msg.receive_count,
                    receipt_handle: msg.receipt_handle,
                    project: project.to_string(),
                    attributes: msg.attributes.unwrap_or_default(),
                });
            }
        }

        tracing::debug!(
            "Received {} messages from queue via gateway: project={}",
            messages.len(),
            project
        );

        Ok(messages)
    }

    async fn delete_chunk(&self, message_id: &str, receipt_handle: &str) -> Result<()> {
        let token = self.token_source.token().await?;

        let req = DeleteMessageRequest {
            receipt_handle: receipt_handle.to_string(),
        };

        let resp = self
            .http_client
            .delete(&self.gateway_url)
            .header("X-Service", "sqs")
            .header("Authorization", format!("Bearer {}", token))
            .header("Content-Type", "application/json")
            .json(&req)
            .send()
            .await?;

        if !resp.status().is_success() && resp.status().as_u16() != 204 {
            let status = resp.status();
            let error = resp.text().await.unwrap_or_default();
            return Err(anyhow!("DeleteMessage failed: {} {}", status, error));
        }

        tracing::debug!("Message deleted via gateway: message_id={}", message_id);

        Ok(())
    }

    async fn change_visibility(
        &self,
        message_id: &str,
        _receipt_handle: &str,
        visibility_timeout_secs: i32,
    ) -> Result<()> {
        // TODO: Implement when gateway adds support for ChangeMessageVisibility
        
        tracing::warn!(
            "ChangeMessageVisibility not yet supported via gateway: message_id={}, timeout={}s",
            message_id,
            visibility_timeout_secs
        );

        Ok(())
    }

    async fn send_to_dlq(&self, message_id: &str, receipt_handle: &str, reason: &str) -> Result<()> {
        // Delete from main queue
        self.delete_chunk(message_id, receipt_handle).await?;

        // Send to DLQ
        let token = self.token_source.token().await?;

        let dlq_body = serde_json::json!({
            "message_id": message_id,
            "reason": reason,
            "failed_at": std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs()
        })
        .to_string();

        let encoded_body = base64::encode(dlq_body.as_bytes());

        let req = SendMessageRequest {
            message_body: encoded_body,
            message_attributes: MessageAttributes {
                values: std::collections::HashMap::new(),
            },
            delay_seconds: 0,
        };

        let resp = self
            .http_client
            .post(&self.gateway_url)
            .header("X-Service", "sqs")
            .header("Authorization", format!("Bearer {}", token))
            .header("Content-Type", "application/json")
            .json(&req)
            .send()
            .await?;

        if !resp.status().is_success() {
            return Err(anyhow!("SendToDLQ failed: {}", resp.status()));
        }

        tracing::warn!(
            "Message sent to DLQ via gateway: message_id={}, reason={}",
            message_id,
            reason
        );

        Ok(())
    }

    async fn get_stats(&self, project: Option<&str>) -> Result<QueueStats> {
        let _token = self.token_source.token().await?;
        let _project = project.unwrap_or("default");

        Ok(QueueStats {
            available_messages: 0,
            in_flight_messages: 0,
            dead_letter_messages: 0,
            total_processed: 0,
            average_delay_secs: 0,
        })
    }

    async fn purge(&self, project: Option<&str>) -> Result<usize> {
        let _token = self.token_source.token().await?;
        let _project = project.unwrap_or("default");

        tracing::warn!("Purge not yet supported via gateway");

        Ok(0)
    }

    async fn health_check(&self) -> Result<()> {
        let token = self.token_source.token().await?;

        let query = format!(
            "X-Service=sqs&queue=health-check&maxNumberOfMessages=0&waitTimeSeconds=0&visibilityTimeoutSeconds=0"
        );

        let resp = self
            .http_client
            .get(&format!("{}?{}", self.gateway_url, query))
            .header("Authorization", format!("Bearer {}", token))
            .timeout(std::time::Duration::from_secs(5))
            .send()
            .await?;

        if resp.status().is_success() || resp.status().as_u16() == 404 {
            tracing::debug!("Gateway health check passed");
            Ok(())
        } else {
            Err(anyhow!("Gateway health check failed: {}", resp.status()))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_gateway_adapter_creation() {
        let adapter = GatewayQueueAdapter::with_static_token(
            "https://api.riotpiao.com".to_string(),
            "test-token".to_string(),
        );

        assert_eq!(adapter.gateway_url, "https://api.riotpiao.com");
        assert_eq!(adapter.default_queue_prefix, "poimen-chunks");
    }

    #[test]
    fn test_queue_name_formatting() {
        let adapter = GatewayQueueAdapter::with_static_token(
            "https://api.riotpiao.com".to_string(),
            "test-token".to_string(),
        );

        assert_eq!(adapter.queue_name("myproject"), "poimen-chunks-myproject");
    }

    #[test]
    fn test_base64_roundtrip() {
        let original = "hello world";
        let encoded = base64::encode(original.as_bytes());
        let decoded = String::from_utf8(base64::decode(encoded.as_bytes()).unwrap()).unwrap();
        assert_eq!(decoded, original);
    }

    #[tokio::test]
    async fn test_static_token_provider() {
        let provider = StaticTokenProvider::new("my-token".to_string());
        let token = provider.token().await.unwrap();
        assert_eq!(token, "my-token");
    }
}