use super::thinking::ThinkingStateManager;
use super::{create_http_client, LlmProvider};
use anyhow::{bail, Context, Result};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use std::time::Duration;

/// OpenAI API client with o-series reasoning support
pub struct OpenAiClient {
    base_url: String,
    api_key: String,
    model: String,
    client: reqwest::Client,
    thinking_enabled: bool,
    reasoning_effort: Option<String>,
    max_tokens: u32,
    temperature: f32,
    top_p: Option<f32>,
    thinking_state: Option<Arc<ThinkingStateManager>>,
}

#[derive(Debug, Serialize)]
struct ChatCompletionRequest {
    model: String,
    messages: Vec<Message>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_p: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    reasoning_effort: Option<String>,
    stream: bool,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
struct Message {
    role: String,
    content: String,
}

#[derive(Debug, Deserialize)]
struct ChatCompletionResponse {
    choices: Vec<Choice>,
}

#[derive(Debug, Deserialize)]
struct Choice {
    message: Message,
}

// --- Streaming response structures ---

#[derive(Debug, Deserialize)]
struct StreamChunk {
    choices: Vec<StreamChoice>,
}

#[derive(Debug, Deserialize)]
struct StreamChoice {
    delta: StreamDelta,
    #[serde(default)]
    finish_reason: Option<String>,
}

#[derive(Debug, Deserialize, Default)]
struct StreamDelta {
    #[serde(default)]
    content: Option<String>,
    #[serde(default)]
    reasoning_content: Option<String>,
}

#[derive(Debug, Deserialize)]
struct ErrorResponse {
    error: ApiError,
}

#[derive(Debug, Deserialize)]
struct ApiError {
    message: String,
    #[serde(rename = "type")]
    error_type: String,
}

impl OpenAiClient {
    /// Create new OpenAI client
    pub fn new(base_url: &str, api_key: &str, model: &str) -> Result<Self> {
        let client = create_http_client(Duration::from_secs(60))?;

        Ok(Self {
            base_url: base_url.trim_end_matches('/').to_string(),
            api_key: api_key.to_string(),
            model: model.to_string(),
            client,
            thinking_enabled: false,
            reasoning_effort: None,
            max_tokens: 500,
            temperature: 0.7,
            top_p: None,
            thinking_state: None,
        })
    }

    pub fn with_timeout(mut self, timeout: Duration) -> Result<Self> {
        self.client = create_http_client(timeout)?;
        Ok(self)
    }

    pub fn with_thinking(mut self, enabled: bool) -> Self {
        self.thinking_enabled = enabled;
        self
    }

    pub fn with_reasoning_effort(mut self, effort: Option<String>) -> Self {
        self.reasoning_effort = effort;
        self
    }

    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
        self.max_tokens = max_tokens;
        self
    }

    pub fn with_temperature(mut self, temperature: f32) -> Self {
        self.temperature = temperature;
        self
    }

    pub fn with_top_p(mut self, top_p: f32) -> Self {
        self.top_p = Some(top_p);
        self
    }

    pub fn with_thinking_state(mut self, state: Arc<ThinkingStateManager>) -> Self {
        self.thinking_state = Some(state);
        self
    }

    pub async fn list_models(&self) -> Result<Vec<String>> {
        let url = format!("{}/models", self.base_url);

        let response = self
            .client
            .get(&url)
            .header("Authorization", format!("Bearer {}", self.api_key))
            .send()
            .await
            .context("Failed to list OpenAI models")?;

        if !response.status().is_success() {
            let status = response.status();
            let text = response.text().await.unwrap_or_default();
            bail!("OpenAI API error: {} - {}", status, text);
        }

        #[derive(Deserialize)]
        struct ModelsResponse {
            data: Vec<Model>,
        }

        #[derive(Deserialize)]
        struct Model {
            id: String,
        }

        let result: ModelsResponse = response
            .json()
            .await
            .context("Failed to parse OpenAI response")?;

        Ok(result.data.into_iter().map(|m| m.id).collect())
    }

    pub async fn validate_key(&self) -> Result<bool> {
        match self.list_models().await {
            Ok(_) => Ok(true),
            Err(e) => {
                let err_str = e.to_string();
                if err_str.contains("401") || err_str.contains("Unauthorized") {
                    Ok(false)
                } else {
                    Err(e)
                }
            }
        }
    }
}

#[async_trait]
impl LlmProvider for OpenAiClient {
    async fn generate(&self, prompt: &str) -> Result<String> {
        let messages = vec![Message {
            role: "user".to_string(),
            content: prompt.to_string(),
        }];

        self.chat_completion_with_retry(messages).await
    }

    async fn generate_with_system(&self, system: &str, user: &str) -> Result<String> {
        let mut messages = vec![];

        if !system.is_empty() {
            messages.push(Message {
                role: "system".to_string(),
                content: system.to_string(),
            });
        }

        messages.push(Message {
            role: "user".to_string(),
            content: user.to_string(),
        });

        self.chat_completion_with_retry(messages).await
    }

    async fn is_available(&self) -> bool {
        self.validate_key().await.unwrap_or(false)
    }

    fn name(&self) -> &str {
        "openai"
    }
}

impl OpenAiClient {
    async fn chat_completion_with_retry(&self, messages: Vec<Message>) -> Result<String> {
        let mut last_error = None;

        for attempt in 1..=3 {
            match self.chat_completion(messages.clone()).await {
                Ok(result) => return Ok(result),
                Err(e) => {
                    let err_msg = e.to_string();
                    let is_retryable = err_msg.contains("timeout")
                        || err_msg.contains("connection")
                        || err_msg.contains("temporary")
                        || err_msg.contains("5")
                            && (err_msg.contains("500")
                                || err_msg.contains("502")
                                || err_msg.contains("503")
                                || err_msg.contains("504"));

                    if !is_retryable || attempt == 3 {
                        last_error = Some(e);
                        break;
                    }

                    tokio::time::sleep(Duration::from_millis(500 * 2u64.pow(attempt - 1))).await;
                }
            }
        }

        Err(last_error.unwrap_or_else(|| anyhow::anyhow!("Request failed after retries")))
    }

    async fn chat_completion(&self, messages: Vec<Message>) -> Result<String> {
        if self.thinking_enabled {
            self.streaming_chat_completion(messages).await
        } else {
            self.non_streaming_chat_completion(messages).await
        }
    }

    async fn non_streaming_chat_completion(&self, messages: Vec<Message>) -> Result<String> {
        let url = format!("{}/chat/completions", self.base_url);

        let request = ChatCompletionRequest {
            model: self.model.clone(),
            messages,
            max_tokens: Some(self.max_tokens),
            temperature: Some(self.temperature),
            top_p: self.top_p,
            reasoning_effort: None,
            stream: false,
        };

        let response = self
            .client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .json(&request)
            .send()
            .await
            .context("Failed to send request to OpenAI")?;

        let status = response.status();

        if !status.is_success() {
            let text = response.text().await.unwrap_or_default();

            if let Ok(error) = serde_json::from_str::<ErrorResponse>(&text) {
                bail!(
                    "OpenAI API error: {} ({})",
                    error.error.message,
                    error.error.error_type
                );
            }

            bail!("OpenAI API error: {} - {}", status, text);
        }

        let result: ChatCompletionResponse = response
            .json()
            .await
            .context("Failed to parse OpenAI response")?;

        result
            .choices
            .into_iter()
            .next()
            .map(|c| c.message.content.trim().to_string())
            .filter(|s| !s.is_empty())
            .ok_or_else(|| anyhow::anyhow!("No response from OpenAI"))
    }

    /// Streaming request for reasoning mode, filters reasoning_content from output
    async fn streaming_chat_completion(&self, messages: Vec<Message>) -> Result<String> {
        let url = format!("{}/chat/completions", self.base_url);

        // For reasoning/thinking mode, omit temperature and top_p
        let request = ChatCompletionRequest {
            model: self.model.clone(),
            messages,
            max_tokens: Some(self.max_tokens),
            temperature: None,
            top_p: None,
            reasoning_effort: self.reasoning_effort.clone(),
            stream: true,
        };

        let response = self
            .client
            .post(&url)
            .header("Authorization", format!("Bearer {}", self.api_key))
            .header("Content-Type", "application/json")
            .header("Accept", "text/event-stream")
            .json(&request)
            .send()
            .await
            .context("Failed to send streaming request to OpenAI")?;

        let status = response.status();

        if !status.is_success() {
            let text = response.text().await.unwrap_or_default();

            if let Ok(error) = serde_json::from_str::<ErrorResponse>(&text) {
                bail!(
                    "OpenAI API error: {} ({})",
                    error.error.message,
                    error.error.error_type
                );
            }

            bail!("OpenAI API error: {} - {}", status, text);
        }

        let mut content_buffer = String::new();
        let mut has_reasoning = false;
        let mut has_content = false;

        let thinking_state = self.thinking_state.as_ref();

        let mut byte_stream = response.bytes_stream();
        let mut line_buffer = String::new();

        use futures_util::StreamExt;

        while let Some(chunk) = byte_stream.next().await {
            let chunk = chunk.context("Failed to read streaming response chunk")?;
            let chunk_str =
                String::from_utf8(chunk.to_vec()).context("Invalid UTF-8 in stream chunk")?;

            line_buffer.push_str(&chunk_str);

            while let Some(line_end) = line_buffer.find('\n') {
                let line = line_buffer[..line_end].trim().to_string();
                line_buffer = line_buffer[line_end + 1..].to_string();

                if line.is_empty() {
                    continue;
                }

                if line == "data: [DONE]" {
                    break;
                }

                if let Some(json_str) = line.strip_prefix("data: ") {
                    if let Ok(chunk) = serde_json::from_str::<StreamChunk>(json_str) {
                        for choice in &chunk.choices {
                            // Handle reasoning_content (o-series)
                            if let Some(ref reasoning) = choice.delta.reasoning_content
                                && !reasoning.is_empty()
                            {
                                if !has_reasoning {
                                    has_reasoning = true;
                                    if let Some(state) = thinking_state {
                                        state.start_thinking();
                                    }
                                }
                                continue;
                            }

                            // Handle content
                            if let Some(ref content) = choice.delta.content
                                && !content.is_empty()
                            {
                                if has_reasoning && !has_content
                                    && let Some(state) = thinking_state
                                {
                                    state.end_thinking();
                                }
                                has_content = true;
                                content_buffer.push_str(content);
                            }
                        }
                    }
                }
            }
        }

        if let Some(state) = thinking_state {
            state.end_thinking();
        }

        let result = content_buffer.trim().to_string();

        if result.is_empty() {
            if has_reasoning && !has_content {
                bail!(
                    "OpenAI returned reasoning content but no final answer. \
                     The model may have entered an incomplete reasoning state. \
                     Please try again or disable thinking mode."
                );
            }
            bail!(
                "No response from OpenAI. \
                 If thinking mode is enabled, try disabling it or ensure the model supports reasoning."
            );
        }

        Ok(result)
    }
}

/// Azure OpenAI client (extends OpenAI with Azure-specific config)
pub struct AzureOpenAiClient {
    endpoint: String,
    api_key: String,
    deployment: String,
    api_version: String,
    client: reqwest::Client,
    thinking_enabled: bool,
    reasoning_effort: Option<String>,
    max_tokens: u32,
    temperature: f32,
    top_p: Option<f32>,
    thinking_state: Option<Arc<ThinkingStateManager>>,
}

impl AzureOpenAiClient {
    pub fn new(
        endpoint: &str,
        api_key: &str,
        deployment: &str,
        api_version: &str,
    ) -> Result<Self> {
        let client = create_http_client(Duration::from_secs(60))?;

        Ok(Self {
            endpoint: endpoint.trim_end_matches('/').to_string(),
            api_key: api_key.to_string(),
            deployment: deployment.to_string(),
            api_version: api_version.to_string(),
            client,
            thinking_enabled: false,
            reasoning_effort: None,
            max_tokens: 500,
            temperature: 0.7,
            top_p: None,
            thinking_state: None,
        })
    }

    async fn chat_completion(&self, messages: Vec<Message>) -> Result<String> {
        let url = format!(
            "{}/openai/deployments/{}/chat/completions?api-version={}",
            self.endpoint, self.deployment, self.api_version
        );

        let request = ChatCompletionRequest {
            model: self.deployment.clone(),
            messages,
            max_tokens: Some(self.max_tokens),
            temperature: Some(self.temperature),
            top_p: self.top_p,
            reasoning_effort: self.reasoning_effort.clone(),
            stream: false,
        };

        let response = self
            .client
            .post(&url)
            .header("api-key", &self.api_key)
            .header("Content-Type", "application/json")
            .json(&request)
            .send()
            .await
            .context("Failed to send request to Azure OpenAI")?;

        if !response.status().is_success() {
            let status = response.status();
            let text = response.text().await.unwrap_or_default();
            bail!("Azure OpenAI API error: {} - {}", status, text);
        }

        let result: ChatCompletionResponse = response
            .json()
            .await
            .context("Failed to parse Azure OpenAI response")?;

        result
            .choices
            .into_iter()
            .next()
            .map(|c| c.message.content.trim().to_string())
            .filter(|s| !s.is_empty())
            .ok_or_else(|| anyhow::anyhow!("No response from Azure OpenAI"))
    }
}

#[async_trait]
impl LlmProvider for AzureOpenAiClient {
    async fn generate(&self, prompt: &str) -> Result<String> {
        let messages = vec![Message {
            role: "user".to_string(),
            content: prompt.to_string(),
        }];

        self.chat_completion(messages).await
    }

    async fn generate_with_system(&self, system: &str, user: &str) -> Result<String> {
        let mut messages = vec![];

        if !system.is_empty() {
            messages.push(Message {
                role: "system".to_string(),
                content: system.to_string(),
            });
        }

        messages.push(Message {
            role: "user".to_string(),
            content: user.to_string(),
        });

        self.chat_completion(messages).await
    }

    async fn is_available(&self) -> bool {
        let url = format!(
            "{}/openai/deployments/{}/chat/completions?api-version={}",
            self.endpoint, self.deployment, self.api_version
        );

        let request = ChatCompletionRequest {
            model: self.deployment.clone(),
            messages: vec![Message {
                role: "user".to_string(),
                content: "Hi".to_string(),
            }],
            max_tokens: Some(5),
            temperature: Some(0.0),
            top_p: None,
            reasoning_effort: None,
            stream: false,
        };

        match self
            .client
            .post(&url)
            .header("api-key", &self.api_key)
            .json(&request)
            .send()
            .await
        {
            Ok(response) => response.status().is_success(),
            Err(_) => false,
        }
    }

    fn name(&self) -> &str {
        "azure-openai"
    }
}

/// Available OpenAI models (including o-series with reasoning)
pub const OPENAI_MODELS: &[&str] = &[
    "o4-mini",
    "o3",
    "o3-mini",
    "o1",
    "o1-mini",
    "o1-pro",
    "gpt-4.1",
    "gpt-4.1-mini",
    "gpt-4.1-nano",
    "gpt-4o",
    "gpt-4o-mini",
    "gpt-4-turbo",
    "gpt-4",
    "gpt-3.5-turbo",
];

pub fn is_valid_model(model: &str) -> bool {
    OPENAI_MODELS.contains(&model)
}

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

    #[test]
    fn test_model_validation_o_series() {
        assert!(is_valid_model("o4-mini"));
        assert!(is_valid_model("o3"));
        assert!(is_valid_model("o1"));
        assert!(is_valid_model("gpt-4o"));
        assert!(is_valid_model("gpt-3.5-turbo"));
        assert!(!is_valid_model("invalid-model"));
    }

    #[test]
    fn test_stream_delta_reasoning_parsing() {
        let json = r#"{"content":null,"reasoning_content":"Let me think..."}"#;
        let delta: StreamDelta = serde_json::from_str(json).unwrap();
        assert!(delta.content.is_none());
        assert_eq!(
            delta.reasoning_content,
            Some("Let me think...".to_string())
        );
    }

    #[test]
    fn test_stream_delta_content_parsing() {
        let json = r#"{"content":"Hello","reasoning_content":null}"#;
        let delta: StreamDelta = serde_json::from_str(json).unwrap();
        assert_eq!(delta.content, Some("Hello".to_string()));
        assert!(delta.reasoning_content.is_none());
    }
}
