use anyhow::{anyhow, Result};
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::env;
use std::time::Duration;
use mem_core::gated_loop::LlmClient;

/// Completion response from the model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Completion {
    pub text: String,
    pub usage: Usage,
    pub latency_ms: u64,
}

/// Token usage breakdown.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Usage {
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
    pub total_tokens: u32,
}

/// Chat client for the gateway.
pub struct ChatClient {
    base_url: String,
    api_key: String,
    model: String,
    http: Client,
    timeout: Duration,
    max_retries: u32,
}

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

#[derive(Debug, Serialize)]
struct CompletionRequest {
    model: String,
    messages: Vec<Message>,
    max_tokens: u32,
}

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

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

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

#[derive(Debug, Deserialize)]
struct ResponseUsage {
    prompt_tokens: u32,
    completion_tokens: u32,
    total_tokens: u32,
}

impl LlmClient for ChatClient {
    fn complete_blocking(&self, system: &str, user: &str, max_tokens: usize) -> Result<String> {
        ChatClient::complete_blocking(self, system, user, max_tokens as u32)
    }
}

impl ChatClient {
    /// Create a new chat client.
    ///
    /// # Arguments
    /// * `base_url` - Gateway base URL (e.g., `https://api.riotpiao.com/v1`)
    /// * `api_key` - Authentication key
    /// * `model` - Model identifier (e.g., `qwen2.5:3b-instruct`)
    pub fn new(base_url: impl Into<String>, api_key: impl Into<String>, model: impl Into<String>) -> Result<Self> {
        Ok(Self {
            base_url: base_url.into(),
            api_key: api_key.into(),
            model: model.into(),
            http: Client::new(),
            timeout: Duration::from_secs(300),
            max_retries: 3,
        })
    }

    /// Set custom timeout.
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        self
    }

    /// Set max retries for 5xx errors (default: 3).
    pub fn with_max_retries(mut self, retries: u32) -> Self {
        self.max_retries = retries;
        self
    }

    /// Complete synchronously (blocks until response).
    pub fn complete_blocking(&self, system: &str, user: &str, max_tokens: u32) -> Result<String> {
        let rt = tokio::runtime::Runtime::new()?;
        rt.block_on(async { 
            let completion = self.complete(system, user, max_tokens).await?;
            Ok(completion.text)
        })
    }

    /// Complete a prompt.
    pub async fn complete(&self, system: &str, user: &str, max_tokens: u32) -> Result<Completion> {
        let url = format!("{}/chat/completions", self.base_url);

        let request = CompletionRequest {
            model: self.model.clone(),
            messages: vec![
                Message {
                    role: "system".to_string(),
                    content: system.to_string(),
                },
                Message {
                    role: "user".to_string(),
                    content: user.to_string(),
                },
            ],
            max_tokens,
        };

        let body = serde_json::to_string(&request)?;

        // Record request if MEM_LLM_RECORD is set
        if let Ok(record_dir) = env::var("MEM_LLM_RECORD") {
            let filename = format!("{}/request-{}.json", record_dir, chrono::Local::now().timestamp_millis());
            let _ = std::fs::write(&filename, &body);
        }

        let start = std::time::Instant::now();
        let mut last_error: Option<anyhow::Error> = None;

        for attempt in 0..self.max_retries {
            let mut req = self.http.post(&url);
            // Only add apikey header if it's not empty (for backward compatibility)
            if !self.api_key.is_empty() && !self.api_key.starts_with("http") {
                req = req.header("apikey", &self.api_key);
            }
            
            let response = req
                .header("Content-Type", "application/json")
                .body(body.clone())
                .timeout(self.timeout)
                .send()
                .await;

            let response = match response {
                Ok(r) => r,
                Err(e) => {
                    last_error = Some(anyhow!("Request failed: {}", e));
                    if e.is_timeout() || e.is_status() {
                        if attempt < self.max_retries - 1 {
                            tokio::time::sleep(Duration::from_millis(100 * 2_u64.pow(attempt))).await;
                            continue;
                        }
                    }
                    return Err(last_error.unwrap());
                }
            };

            let status = response.status();
            let body_text = response.text().await.unwrap_or_default();

            // Record response if MEM_LLM_RECORD is set
            if let Ok(record_dir) = env::var("MEM_LLM_RECORD") {
                let filename = format!(
                    "{}/response-{}-{}.json",
                    record_dir,
                    chrono::Local::now().timestamp_millis(),
                    status
                );
                let _ = std::fs::write(&filename, &body_text);
            }

            // Handle auth error
            if status == 401 {
                return Err(anyhow!(
                    "Auth error (401): check apikey header format. Response: {}",
                    body_text
                ));
            }

            // 4xx errors should not be retried
            if status.is_client_error() {
                return Err(anyhow!("Client error ({}): {}", status, body_text));
            }

            // 5xx errors should be retried
            if status.is_server_error() {
                if attempt < self.max_retries - 1 {
                    last_error = Some(anyhow!("Server error ({}): {}", status, body_text));
                    tokio::time::sleep(Duration::from_millis(100 * 2_u64.pow(attempt))).await;
                    continue;
                } else {
                    return Err(anyhow!("Server error ({}): {} (after {} retries)", status, body_text, self.max_retries));
                }
            }

            // Parse success response
            if status.is_success() {
                let completion_response: CompletionResponse = serde_json::from_str(&body_text)?;

                if completion_response.choices.is_empty() {
                    return Err(anyhow!("No choices in response"));
                }

                let latency_ms = start.elapsed().as_millis() as u64;
                let text = completion_response.choices[0].message.content.clone();
                let usage = Usage {
                    prompt_tokens: completion_response.usage.prompt_tokens,
                    completion_tokens: completion_response.usage.completion_tokens,
                    total_tokens: completion_response.usage.total_tokens,
                };

                return Ok(Completion {
                    text,
                    usage,
                    latency_ms,
                });
            }

            return Err(anyhow!("Unexpected status {}: {}", status, body_text));
        }

        Err(last_error.unwrap_or_else(|| anyhow!("Max retries exhausted")))
    }
}

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

    #[test]
    fn test_completion_structs_serialize() {
        let usage = Usage {
            prompt_tokens: 100,
            completion_tokens: 50,
            total_tokens: 150,
        };
        let json = serde_json::to_string(&usage).unwrap();
        assert!(json.contains("100"));
    }
}