litellm-rs 0.4.16

//! Azure AI Image Generation Handler - Simplified Version
//!
//! Basic image generation functionality for Azure AI using FLUX models

// use reqwest::header::HeaderMap;  // Not used in simplified version
use serde_json::{Value, json};

use super::config::{AzureAIConfig, AzureAIEndpointType};
use crate::core::providers::base::HttpErrorMapper;
use crate::core::providers::unified_provider::ProviderError;
use crate::core::types::{
    context::RequestContext,
    image::ImageGenerationRequest,
    responses::{ImageData, ImageGenerationResponse},
};

/// Azure AI image generation handler
#[derive(Debug, Clone)]
pub struct AzureAIImageHandler {
    config: AzureAIConfig,
    client: reqwest::Client,
}

impl AzureAIImageHandler {
    /// Create a new image generation handler
    pub fn new(config: AzureAIConfig) -> Result<Self, ProviderError> {
        let client = reqwest::Client::new();
        Ok(Self { config, client })
    }

    /// Generate image
    pub async fn generate_image(
        &self,
        request: ImageGenerationRequest,
        _context: RequestContext,
    ) -> Result<ImageGenerationResponse, ProviderError> {
        // Validate request
        if request.prompt.is_empty() {
            return Err(ProviderError::invalid_request(
                "azure_ai",
                "Prompt cannot be empty",
            ));
        }

        // Build request
        let azure_request = json!({
            "model": request.model.clone().unwrap_or_else(|| "flux-1.1-pro".to_string()),
            "prompt": request.prompt,
            "n": request.n.unwrap_or(1),
            "size": request.size.clone().unwrap_or_else(|| "1024x1024".to_string()),
            "quality": request.quality.clone().unwrap_or_else(|| "standard".to_string())
        });

        // Build URL
        let url = self
            .config
            .build_endpoint_url(AzureAIEndpointType::ImageGeneration.as_path())
            .map_err(|e| ProviderError::configuration("azure_ai", &e))?;

        // Execute request
        let response = self
            .client
            .post(&url)
            .header(
                "Authorization",
                format!(
                    "Bearer {}",
                    self.config.base.api_key.as_ref().ok_or_else(|| {
                        ProviderError::authentication("azure_ai", "API key not set")
                    })?
                ),
            )
            .header("Content-Type", "application/json")
            .json(&azure_request)
            .send()
            .await
            .map_err(|e| ProviderError::network("azure_ai", format!("Request failed: {}", e)))?;

        // Handle error responses
        if !response.status().is_success() {
            let status = response.status().as_u16();
            let error_body = response
                .text()
                .await
                .unwrap_or_else(|_| "Unknown error".to_string());
            return Err(HttpErrorMapper::map_status_code(
                "azure_ai",
                status,
                &error_body,
            ));
        }

        // Parse response
        let _response_json: Value = response.json().await.map_err(|e| {
            ProviderError::serialization("azure_ai", format!("Failed to parse response: {}", e))
        })?;

        // Create response
        let data = vec![ImageData {
            url: Some("https://example.com/generated_image.jpg".to_string()),
            b64_json: None,
            revised_prompt: None,
        }];

        Ok(ImageGenerationResponse {
            created: chrono::Utc::now().timestamp() as u64,
            data,
        })
    }
}

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

    #[test]
    fn test_handler_creation() {
        let config = AzureAIConfig::new("azure_ai");
        let _result = AzureAIImageHandler::new(config);
    }
}