dynamic_grounding_for_github_copilot 0.1.0

use crate::error::{Error, Result};
use std::sync::Arc;

/// Manages API key retrieval and validation
/// Never stores the key in memory longer than necessary
#[derive(Clone)]
pub struct ApiKeyManager {
    /// Channel to communicate with VS Code extension for key retrieval
    key_provider: Arc<dyn ApiKeyProvider + Send + Sync>,
}

#[async_trait::async_trait]
pub trait ApiKeyProvider: Send + Sync {
    /// Retrieve the API key from secure storage
    async fn get_key(&self) -> Result<SecureString>;

    /// Store the API key in secure storage (optional, for setup)
    async fn set_key(&self, _key: SecureString) -> Result<()> {
        Err(Error::ApiKeyError(
            "API key storage not supported by this provider".to_string(),
        ))
    }
}

/// A string that won't be accidentally logged or displayed
/// Implements Drop to zero memory on deallocation
#[derive(Clone)]
pub struct SecureString {
    inner: String,
}

impl SecureString {
    pub fn new(s: String) -> Self {
        Self { inner: s }
    }

    pub fn as_str(&self) -> &str {
        &self.inner
    }
}

impl Drop for SecureString {
    fn drop(&mut self) {
        // Zero out the memory
        unsafe {
            let bytes = self.inner.as_bytes_mut();
            for b in bytes {
                std::ptr::write_volatile(b, 0);
            }
        }
    }
}

// Prevent accidental logging
impl std::fmt::Debug for SecureString {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "SecureString([REDACTED])")
    }
}

impl std::fmt::Display for SecureString {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "[REDACTED]")
    }
}

impl ApiKeyManager {
    pub fn new(provider: Arc<dyn ApiKeyProvider + Send + Sync>) -> Self {
        Self {
            key_provider: provider,
        }
    }

    /// Get the API key from secure storage
    pub async fn get_api_key(&self) -> Result<SecureString> {
        self.key_provider.get_key().await
    }

    /// Validate that an API key has the correct format
    pub fn validate_key_format(key: &str) -> Result<()> {
        // Gemini API keys typically start with "AI" and are 39 characters
        if key.is_empty() {
            return Err(Error::ApiKeyError("API key is empty".to_string()));
        }

        if key.len() < 20 {
            return Err(Error::ConfigError(
                "API key appears to be too short".to_string(),
            ));
        }

        Ok(())
    }
}

/// Provider that reads from stdin (for MCP protocol communication)
pub struct StdinApiKeyProvider;

#[async_trait::async_trait]
impl ApiKeyProvider for StdinApiKeyProvider {
    async fn get_key(&self) -> Result<SecureString> {
        // In MCP protocol, the VS Code extension will provide the key
        // via initialization parameters or environment
        // For now, return an error - this will be implemented by the MCP server
        Err(Error::ApiKeyError("No API key provided".to_string()))
    }
}

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

    struct MockKeyProvider {
        key: String,
    }

    #[async_trait::async_trait]
    impl ApiKeyProvider for MockKeyProvider {
        async fn get_key(&self) -> Result<SecureString> {
            Ok(SecureString::new(self.key.clone()))
        }
    }

    #[tokio::test]
    async fn test_api_key_manager() {
        let provider = Arc::new(MockKeyProvider {
            key: "AIzaSyDMockKey1234567890123456789".to_string(),
        });
        let manager = ApiKeyManager::new(provider);

        let key = manager.get_api_key().await.unwrap();
        assert_eq!(key.as_str(), "AIzaSyDMockKey1234567890123456789");
    }

    #[test]
    fn test_secure_string_debug() {
        let secure = SecureString::new("secret".to_string());
        let debug_str = format!("{:?}", secure);
        assert!(debug_str.contains("REDACTED"));
        assert!(!debug_str.contains("secret"));
    }

    #[test]
    fn test_validate_key_format() {
        assert!(ApiKeyManager::validate_key_format("AIzaSyDMockKey1234567890123456789").is_ok());
        assert!(ApiKeyManager::validate_key_format("").is_err());
        assert!(ApiKeyManager::validate_key_format("short").is_err());
    }
}