//! Azure OpenAI / Azure AI Foundry provider
//!
//! This provider supports Azure OpenAI Service and Azure AI Foundry endpoints.
//! It can be configured either programmatically or via config.yml.

use crate::config::{AuthType, Modality, ModelConfig};
use crate::http::{post_json, HttpRequestSpec};
use async_trait::async_trait;
use enact_core::providers::{
    ChatRequest, ChatResponse, EmbeddingData, EmbeddingRequest, EmbeddingResponse, EmbeddingUsage,
    ImageData, ImageGenerationRequest, ImageGenerationResponse, ModelCapabilities, ModelProvider,
};
use reqwest::Client;
use serde_json::Value;
use std::collections::HashSet;
use std::time::Duration;

/// Azure OpenAI / Azure AI Foundry provider.
///
/// Supports both traditional Azure OpenAI and Azure AI Foundry endpoints.
/// All configuration can be loaded from config.yml via `from_model_config()`.
pub struct AzureOpenAI {
    client: Client,
    /// Model identifier (e.g., "azure_foundry:text:gpt-5-nano")
    model_id: String,
    /// Provider-specific model name (from cost.provider_key, used in request body when required)
    provider_model: Option<String>,
    /// API key for authentication
    api_key: String,
    /// Full URL for API calls (from config)
    full_url: String,
    /// Authentication type (bearer or header)
    auth_type: AuthType,
    /// Request timeout
    timeout: Duration,
    /// Whether to use max_completion_tokens instead of max_tokens
    use_max_completion_tokens: bool,
    /// Whether temperature is supported
    has_temperature_supported: bool,
    /// Whether tool calls are supported
    supports_tool_calls: bool,
    /// Whether reasoning is supported
    supports_reasoning: bool,
    /// Whether model param is required in request body (from config)
    requires_model_in_body: bool,
    /// Whether this is a blackforestlabs endpoint (uses width/height instead of size)
    is_blackforestlabs: bool,
    /// Default max tokens (from config runtime.max_tokens_default)
    max_tokens_default: u32,
    /// Cost per 1K input tokens (from config cost.cost_per_1m_input)
    cost_per_1m_input: Option<f64>,
    /// Cost per 1K output tokens (from config cost.cost_per_1m_output)
    cost_per_1m_output: Option<f64>,
    /// Cost per 1K pixels for image (from config cost.cost_per_1m_pixels)
    cost_per_1m_pixels: Option<f64>,
    /// Output modalities supported by this model
    output_modalities: HashSet<Modality>,
}

impl AzureOpenAI {
    /// Create a new Azure provider from model configuration.
    ///
    /// This is the recommended way to create a provider when using config.yml.
    ///
    /// # Arguments
    ///
    /// * `model_id` - The model identifier (e.g., "azure_foundry:text:gpt-5-nano")
    /// * `config` - The model configuration from config.yml
    /// * `api_key` - The resolved API key
    ///
    /// # Example
    ///
    /// ```no_run
    /// use enact_providers::config::{load_config, resolve_secret};
    /// use enact_providers::AzureOpenAI;
    ///
    /// let config = load_config("config.yml").unwrap();
    /// let model_config = config.get_model("azure_foundry:text:gpt-5-nano").unwrap();
    /// let api_key = resolve_secret(model_config.api_key_env_var()).unwrap();
    /// let provider = AzureOpenAI::from_model_config("azure_foundry:text:gpt-5-nano", model_config, &api_key).unwrap();
    /// ```
    pub fn from_model_config(
        model_id: &str,
        config: &ModelConfig,
        api_key: &str,
    ) -> anyhow::Result<Self> {
        // Check if model requires model param in request body (using config helper)
        let requires_model_in_body = config.requires_model_in_body();

        // Check if this is a blackforestlabs endpoint (uses different body format)
        let is_blackforestlabs = config.runtime.full_url.contains("blackforestlabs");

        Ok(Self {
            client: Client::new(),
            model_id: model_id.to_string(),
            provider_model: config.cost.provider_key.clone(),
            api_key: api_key.to_string(),
            full_url: config.runtime.full_url.clone(),
            auth_type: config.auth.auth_type.clone(),
            timeout: Duration::from_secs(config.runtime.timeout_seconds),
            use_max_completion_tokens: config.runtime.use_max_completion_tokens,
            has_temperature_supported: config.runtime.has_temperature_supported,
            supports_tool_calls: config.capabilities.tool_calls,
            supports_reasoning: config.capabilities.reasoning,
            requires_model_in_body,
            is_blackforestlabs,
            max_tokens_default: config.runtime.max_tokens_default,
            cost_per_1m_input: config.cost.cost_per_1m_input,
            cost_per_1m_output: config.cost.cost_per_1m_output,
            cost_per_1m_pixels: config.cost.cost_per_1m_pixels,
            output_modalities: config.output_modality.iter().cloned().collect(),
        })
    }

    /// Build headers based on authentication type.
    fn build_headers(&self) -> Vec<(String, String)> {
        let mut headers = vec![("Content-Type".to_string(), "application/json".to_string())];

        match &self.auth_type {
            AuthType::Header => {
                // Azure-style api-key header
                headers.push(("api-key".to_string(), self.api_key.clone()));
            }
            AuthType::Bearer => {
                // Bearer token in Authorization header
                headers.push((
                    "Authorization".to_string(),
                    format!("Bearer {}", self.api_key),
                ));
            }
        }

        headers
    }

    /// Transform the request based on model configuration.
    fn transform_request(&self, mut request: ChatRequest) -> anyhow::Result<Value> {
        // Handle temperature support
        if !self.has_temperature_supported {
            request.temperature = None;
        }

        // Convert to JSON
        let mut body = serde_json::to_value(&request)?;

        // Handle max_completion_tokens transformation
        if self.use_max_completion_tokens {
            if let Some(obj) = body.as_object_mut() {
                if let Some(max_tokens) = obj.remove("max_tokens") {
                    obj.insert("max_completion_tokens".to_string(), max_tokens);
                }
            }
        }

        Ok(body)
    }
}

#[async_trait]
impl ModelProvider for AzureOpenAI {
    fn name(&self) -> &str {
        "azure-foundry"
    }

    fn model(&self) -> &str {
        &self.model_id
    }

    fn capabilities(&self) -> ModelCapabilities {
        ModelCapabilities {
            max_tokens: self.max_tokens_default,
            max_output_tokens: self.max_tokens_default,
            supports_streaming: true,
            supports_tools: self.supports_tool_calls,
            supports_reasoning: self.supports_reasoning,
            supports_vision: false,
            supports_json_mode: true,
            supports_embeddings: self.output_modalities.contains(&Modality::Embedding),
            supports_image_generation: self.output_modalities.contains(&Modality::Image),
            supports_audio_transcription: self.output_modalities.contains(&Modality::Audio),
            supports_speech: self.output_modalities.contains(&Modality::Audio),
            supports_video_generation: self.output_modalities.contains(&Modality::Video),
            pii_safe: false,
            cost_per_1m_input: self.cost_per_1m_input,
            cost_per_1m_output: self.cost_per_1m_output,
            cost_per_1m_pixels: self.cost_per_1m_pixels,
        }
    }

    async fn chat(&self, request: ChatRequest) -> anyhow::Result<ChatResponse> {
        // Transform request based on config
        let body = self.transform_request(request)?;

        // Build HTTP request spec
        let http_req = HttpRequestSpec {
            url: self.full_url.clone(),
            headers: self.build_headers(),
            body,
            timeout: self.timeout,
        };

        // Execute request using HTTP helper
        let json_response = post_json(&self.client, http_req).await?;

        // Parse response
        let chat_response: ChatResponse = serde_json::from_value(json_response)?;
        Ok(chat_response)
    }

    async fn embed(&self, request: EmbeddingRequest) -> anyhow::Result<EmbeddingResponse> {
        // Check if embeddings are supported
        if !self.output_modalities.contains(&Modality::Embedding) {
            anyhow::bail!("Embeddings not supported by model {}", self.model_id);
        }

        // Build request body (OpenAI-compatible format)
        let body = serde_json::json!({
            "input": request.input,
        });

        // Build HTTP request spec
        let http_req = HttpRequestSpec {
            url: self.full_url.clone(),
            headers: self.build_headers(),
            body,
            timeout: self.timeout,
        };

        // Execute request
        let json_response = post_json(&self.client, http_req).await?;

        // Parse response - Azure returns OpenAI-compatible format
        let data = json_response
            .get("data")
            .and_then(|d| d.as_array())
            .ok_or_else(|| anyhow::anyhow!("Missing 'data' array in embedding response"))?;

        let embeddings: Vec<EmbeddingData> = data
            .iter()
            .enumerate()
            .filter_map(|(i, item)| {
                let embedding = item.get("embedding")?.as_array()?;
                let vec: Vec<f32> = embedding
                    .iter()
                    .filter_map(|v| v.as_f64().map(|f| f as f32))
                    .collect();
                Some(EmbeddingData {
                    embedding: vec,
                    index: i as u32,
                })
            })
            .collect();

        let usage = json_response.get("usage").and_then(|u| {
            Some(EmbeddingUsage {
                prompt_tokens: u.get("prompt_tokens")?.as_u64()? as u32,
                total_tokens: u.get("total_tokens")?.as_u64()? as u32,
            })
        });

        Ok(EmbeddingResponse {
            data: embeddings,
            model: json_response
                .get("model")
                .and_then(|m| m.as_str())
                .unwrap_or(&self.model_id)
                .to_string(),
            usage,
        })
    }

    async fn generate_image(
        &self,
        request: ImageGenerationRequest,
    ) -> anyhow::Result<ImageGenerationResponse> {
        // Check if image generation is supported
        if !self.output_modalities.contains(&Modality::Image) {
            anyhow::bail!("Image generation not supported by model {}", self.model_id);
        }

        // Build request body - blackforestlabs uses different format
        let body = if self.is_blackforestlabs {
            // Blackforestlabs format: uses width/height instead of size, always includes model
            let size = request.size.as_deref().unwrap_or("1024x1024");
            let parts: Vec<&str> = size.split('x').collect();
            let width: u32 = parts.first().and_then(|s| s.parse().ok()).unwrap_or(1024);
            let height: u32 = parts.get(1).and_then(|s| s.parse().ok()).unwrap_or(1024);

            // Extract model name from provider_key
            let model_name = self
                .provider_model
                .as_ref()
                .map(|m| m.rsplit('/').next().unwrap_or(m))
                .unwrap_or("FLUX.2-pro");

            serde_json::json!({
                "prompt": request.prompt,
                "n": request.n.unwrap_or(1),
                "width": width,
                "height": height,
                "model": model_name,
            })
        } else {
            // Standard OpenAI-compatible format
            let mut body = serde_json::json!({
                "prompt": request.prompt,
            });

            // Add optional parameters if provided
            if let Some(n) = request.n {
                body["n"] = serde_json::json!(n);
            }
            if let Some(ref size) = request.size {
                body["size"] = serde_json::json!(size);
            }
            if let Some(ref quality) = request.quality {
                body["quality"] = serde_json::json!(quality);
            }
            if let Some(ref style) = request.style {
                body["style"] = serde_json::json!(style);
            }
            if let Some(ref response_format) = request.response_format {
                body["response_format"] = serde_json::json!(response_format);
            }
            if let Some(ref user) = request.user {
                body["user"] = serde_json::json!(user);
            }

            // Add model param if required by the endpoint
            if self.requires_model_in_body {
                if let Some(ref provider_model) = self.provider_model {
                    // Extract just the model name if provider_key has format "provider/model"
                    let model_name = provider_model.rsplit('/').next().unwrap_or(provider_model);
                    body["model"] = serde_json::json!(model_name);
                }
            }

            body
        };

        // Build HTTP request spec
        let http_req = HttpRequestSpec {
            url: self.full_url.clone(),
            headers: self.build_headers(),
            body,
            timeout: self.timeout,
        };

        // Execute request
        let json_response = post_json(&self.client, http_req).await?;

        // Parse response - Azure returns OpenAI-compatible format
        let created = json_response
            .get("created")
            .and_then(|c| c.as_u64())
            .unwrap_or_else(|| {
                std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .map(|d| d.as_secs())
                    .unwrap_or(0)
            });

        let data_array = json_response
            .get("data")
            .and_then(|d| d.as_array())
            .ok_or_else(|| anyhow::anyhow!("Missing 'data' array in image response"))?;

        let data: Vec<ImageData> = data_array
            .iter()
            .map(|item| ImageData {
                url: item.get("url").and_then(|u| u.as_str()).map(String::from),
                b64_json: item
                    .get("b64_json")
                    .and_then(|b| b.as_str())
                    .map(String::from),
                revised_prompt: item
                    .get("revised_prompt")
                    .and_then(|r| r.as_str())
                    .map(String::from),
            })
            .collect();

        Ok(ImageGenerationResponse { created, data })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::{
        AuthConfig, CapabilitiesConfig, CostConfig, GovernanceConfig, LifecycleStatus, Modality,
        Protocol, RuntimeConfig, SecretType, SecretsConfig,
    };

    fn create_test_config(auth_type: AuthType, use_max_completion_tokens: bool) -> ModelConfig {
        ModelConfig {
            display_name: "Test Model".to_string(),
            provider_name: "azure-foundry".to_string(),
            input_modality: vec![Modality::Text],
            output_modality: vec![Modality::Text],
            runtime: RuntimeConfig {
                base_url: "https://test.azure.com".to_string(),
                full_url: "https://test.azure.com/openai/deployments/test/chat/completions?api-version=2024-05-01".to_string(),
                protocol: Protocol::AzureSdk,
                api_version: "2024-05-01".to_string(),
                use_max_completion_tokens,
                has_temperature_supported: true,
                additional_notes: vec![],
                model_param: false,
                timeout_seconds: 60,
                max_tokens_default: 4096,
                temperature_default: 0.7,
            },
            auth: AuthConfig { auth_type },
            secrets: SecretsConfig {
                secret_type: SecretType::Environment,
                name: "AZURE_API_KEY".to_string(),
            },
            capabilities: CapabilitiesConfig {
                tool_calls: true,
                reasoning: false,
            },
            governance: GovernanceConfig {
                is_use_approve: true,
                approved_by: Some("admin".to_string()),
                approval_date: Some("2024-01-01".to_string()),
                lifecycle_status: LifecycleStatus::Active,
            },
            cost: CostConfig::default(),
            fallbacks: vec![],
        }
    }

    #[test]
    fn test_from_model_config() {
        let config = create_test_config(AuthType::Header, false);
        let provider =
            AzureOpenAI::from_model_config("test:text:model", &config, "test-key").unwrap();

        assert_eq!(provider.name(), "azure-foundry");
        assert_eq!(provider.model(), "test:text:model");
        assert_eq!(
            provider.full_url,
            "https://test.azure.com/openai/deployments/test/chat/completions?api-version=2024-05-01"
        );
    }

    #[test]
    fn test_header_auth_builds_api_key_header() {
        let config = create_test_config(AuthType::Header, false);
        let provider = AzureOpenAI::from_model_config("test", &config, "my-api-key").unwrap();

        let headers = provider.build_headers();
        assert!(headers
            .iter()
            .any(|(k, v)| k == "api-key" && v == "my-api-key"));
        assert!(!headers.iter().any(|(k, _)| k == "Authorization"));
    }

    #[test]
    fn test_bearer_auth_builds_authorization_header() {
        let config = create_test_config(AuthType::Bearer, false);
        let provider = AzureOpenAI::from_model_config("test", &config, "my-api-key").unwrap();

        let headers = provider.build_headers();
        assert!(headers
            .iter()
            .any(|(k, v)| k == "Authorization" && v == "Bearer my-api-key"));
        assert!(!headers.iter().any(|(k, _)| k == "api-key"));
    }

    #[test]
    fn test_max_completion_tokens_transform() {
        let config = create_test_config(AuthType::Header, true);
        let provider = AzureOpenAI::from_model_config("test", &config, "key").unwrap();

        let request = ChatRequest {
            messages: vec![],
            max_tokens: Some(1000),
            temperature: Some(0.5),
            tools: None,
            tool_choice: None,
        };

        let body = provider.transform_request(request).unwrap();
        let obj = body.as_object().unwrap();

        assert!(obj.contains_key("max_completion_tokens"));
        assert!(!obj.contains_key("max_tokens"));
        assert_eq!(obj.get("max_completion_tokens").unwrap(), 1000);
    }

    #[test]
    fn test_no_max_completion_tokens_transform() {
        let config = create_test_config(AuthType::Header, false);
        let provider = AzureOpenAI::from_model_config("test", &config, "key").unwrap();

        let request = ChatRequest {
            messages: vec![],
            max_tokens: Some(1000),
            temperature: Some(0.5),
            tools: None,
            tool_choice: None,
        };

        let body = provider.transform_request(request).unwrap();
        let obj = body.as_object().unwrap();

        assert!(obj.contains_key("max_tokens"));
        assert!(!obj.contains_key("max_completion_tokens"));
    }

    #[test]
    fn test_capabilities_from_config() {
        let config = create_test_config(AuthType::Header, false);
        let provider = AzureOpenAI::from_model_config("test", &config, "key").unwrap();

        let caps = provider.capabilities();
        assert!(caps.supports_tools);
    }
}
