//! Main SDK implementation
//!
//! This module contains the core GeminiSDK struct and its implementation.

use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};

use base64::Engine;
use reqwest::multipart;
use serde::{Deserialize, Serialize};

use crate::agent::Agent;
use crate::attachment::Attachment;
use crate::error::{Result, SDKError};
use crate::GeminiConfig;

/// Main SDK struct for interacting with Gemini AI services
pub struct GeminiSDK {
    config: GeminiConfig,
    client: reqwest::Client,
    agents: HashMap<String, Agent>,
    access_token: Option<String>,
    token_expires_at: Option<u64>,
}

/// Request for sending a message
#[derive(Debug, Clone)]
pub struct MessageRequest {
    /// Text content of the message
    pub text: String,
    /// File attachments
    pub attachments: Vec<Attachment>,
    /// ID of the agent to use (optional)
    pub agent_id: Option<String>,
}

/// Response from sending a message
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MessageResponse {
    /// Generated content
    pub content: String,
    /// Agent ID that was used
    pub agent_id: Option<String>,
}

impl MessageRequest {
    /// Create a text-only message request
    pub fn text_only(text: impl Into<String>) -> Self {
        Self {
            text: text.into(),
            attachments: Vec::new(),
            agent_id: None,
        }
    }

    /// Create a message request with attachments
    pub fn with_attachments(text: impl Into<String>, attachments: Vec<Attachment>) -> Self {
        Self {
            text: text.into(),
            attachments,
            agent_id: None,
        }
    }

    /// Create a message request with a specific agent
    pub fn with_agent(text: impl Into<String>, agent_id: impl Into<String>) -> Self {
        Self {
            text: text.into(),
            attachments: Vec::new(),
            agent_id: Some(agent_id.into()),
        }
    }

    /// Set the agent ID
    pub fn agent(mut self, agent_id: impl Into<String>) -> Self {
        self.agent_id = Some(agent_id.into());
        self
    }

    /// Add an attachment
    pub fn attach(mut self, attachment: Attachment) -> Self {
        self.attachments.push(attachment);
        self
    }
}

impl GeminiSDK {
    /// Create a new SDK instance with default configuration
    pub fn new() -> Result<Self> {
        Self::with_config(GeminiConfig::default())
    }

    /// Create a new SDK instance with custom configuration
    pub fn with_config(config: GeminiConfig) -> Result<Self> {
        let client = reqwest::Client::builder()
            .timeout(std::time::Duration::from_secs(config.timeout))
            .build()
            .map_err(|e| SDKError::config(format!("Failed to create HTTP client: {}", e)))?;

        Ok(Self {
            config,
            client,
            agents: HashMap::new(),
            access_token: None,
            token_expires_at: None,
        })
    }

    /// Create a new SDK instance from environment variables
    pub fn from_env() -> Result<Self> {
        Self::with_config(GeminiConfig::from_env())
    }

    /// Get the current configuration
    pub fn config(&self) -> &GeminiConfig {
        &self.config
    }

    /// Update the configuration
    pub fn set_config(&mut self, config: GeminiConfig) {
        self.config = config;
    }

    /// Add an agent to the SDK
    pub fn add_agent(&mut self, agent: Agent) {
        self.agents.insert(agent.id.clone(), agent);
    }

    /// Get an agent by ID
    pub fn get_agent(&self, agent_id: &str) -> Option<&Agent> {
        self.agents.get(agent_id)
    }

    /// List all agents
    pub fn list_agents(&self) -> Vec<&Agent> {
        self.agents.values().collect()
    }

    /// Remove an agent
    pub fn remove_agent(&mut self, agent_id: &str) -> Option<Agent> {
        self.agents.remove(agent_id)
    }

    /// Send a message and get a response
    pub async fn send_message(&mut self, request: MessageRequest) -> Result<MessageResponse> {
        // Validate attachments
        for attachment in &request.attachments {
            attachment.validate()?;
        }

        // Get agent configuration if specified
        let agent_config = if let Some(agent_id) = &request.agent_id {
            let agent = self.get_agent(agent_id)
                .ok_or_else(|| SDKError::agent_not_found(agent_id))?;
            Some((agent.temperature, agent.max_tokens, agent.system_prompt.clone()))
        } else {
            None
        };

        // Get access token
        let access_token = self.get_access_token().await?;

        // Process attachments and build parts
        let mut parts = vec![Part::Text { text: request.text.clone() }];

        for attachment in &request.attachments {
            if attachment.is_video() {
                // Upload video and get file URI
                let file_uri = self.upload_video_file(&attachment.path).await?;
                parts.push(Part::FileData {
                    file_data: FileData {
                        mime_type: attachment.effective_mime_type()
                            .unwrap_or("video/mp4").to_string(),
                        file_uri,
                    }
                });
            } else if attachment.is_image() {
                // Encode image as base64
                let image_data = tokio::fs::read(&attachment.path).await
                    .map_err(|e| SDKError::attachment(format!("Failed to read image: {}", e)))?;
                let base64_data = base64::prelude::BASE64_STANDARD.encode(&image_data);

                parts.push(Part::InlineData {
                    inline_data: InlineData {
                        mime_type: attachment.effective_mime_type()
                            .unwrap_or("image/jpeg").to_string(),
                        data: base64_data,
                    }
                });
            } else {
                return Err(SDKError::attachment(format!(
                    "Unsupported attachment type: {:?}",
                    attachment.effective_mime_type()
                )));
            }
        }

        // Build generation config
        let generation_config = GenerationConfig {
            temperature: agent_config.as_ref().map(|(temp, _, _)| *temp).unwrap_or(self.config.temperature),
            top_k: 32,
            top_p: 1.0,
            max_output_tokens: agent_config.as_ref().map(|(_, tokens, _)| *tokens).unwrap_or(self.config.max_tokens),
        };

        // Build request with system prompt merged into user message if needed
        let user_text = if let Some((_, _, system_prompt)) = agent_config {
            format!("{}\n\nUser: {}", system_prompt, request.text)
        } else {
            request.text.clone()
        };

        // Update the first text part with the combined message
        if let Some(Part::Text { text }) = parts.first_mut() {
            *text = user_text;
        }

        let contents = vec![ContentPart {
            role: "user".to_string(),
            parts,
        }];

        let request_data = GenerateContentRequest {
            contents,
            generation_config,
        };

        // Send request
        let response_content = self.send_generate_request(&access_token, &request_data).await?;

        Ok(MessageResponse {
            content: response_content,
            agent_id: request.agent_id,
        })
    }

    /// Get Google access token
    async fn get_access_token(&mut self) -> Result<String> {
        // Check cached token
        let current_time = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .map_err(|e| SDKError::Generic(anyhow::anyhow!("System time error: {}", e)))?
            .as_secs();

        if let (Some(token), Some(expires_at)) = (&self.access_token, self.token_expires_at) {
            if current_time < expires_at - 300 { // Refresh 5 minutes early
                return Ok(token.clone());
            }
        }

        // Get new access token
        let url = format!("{}/google/access-token", self.config.base_url);

        let response = self.client
            .get(&url)
            .header("Authorization", format!("Bearer {}", self.config.bearer_token))
            .send()
            .await?;

        let status = response.status();
        if !status.is_success() {
            let error_body = response.text().await.unwrap_or_default();
            return Err(SDKError::api(status.as_u16(), format!("Failed to get access token: {}", error_body)));
        }

        let response_text = response.text().await?;
        let token_response: TokenResponse = serde_json::from_str(&response_text)
            .map_err(|e| SDKError::Json(e))?;

        // Cache token
        self.access_token = Some(token_response.access_token.clone());
        self.token_expires_at = Some(current_time + token_response.expires_in);

        Ok(token_response.access_token)
    }

    /// Upload video file to Gemini
    async fn upload_video_file(&mut self, video_path: &str) -> Result<String> {
        let access_token = self.get_access_token().await?;

        // Read video file
        let video_data = tokio::fs::read(video_path).await
            .map_err(|e| SDKError::attachment(format!("Failed to read video file: {}", e)))?;

        let file_name = std::path::Path::new(video_path)
            .file_name()
            .and_then(|n| n.to_str())
            .unwrap_or("video.mp4");

        // Create multipart form
        let form = multipart::Form::new()
            .part("file", multipart::Part::bytes(video_data)
                .file_name(file_name.to_string())
                .mime_str("video/mp4")
                .map_err(|e| SDKError::attachment(format!("Invalid MIME type: {}", e)))?);

        // Upload to Vertex AI
        let upload_url = format!("{}/google/vertex-ai/upload", self.config.base_url);
        let query_params = [
            ("bucket", "dy-media-storage"),
            ("prefix", "video-analysis")
        ];

        let response = self.client
            .post(&upload_url)
            .header("Authorization", format!("Bearer {}", access_token))
            .header("x-google-api-key", &access_token)
            .query(&query_params)
            .multipart(form)
            .send()
            .await?;

        let status = response.status();
        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Err(SDKError::api(status.as_u16(), format!("Video upload failed: {}", error_text)));
        }

        let response_text = response.text().await?;
        let upload_response: UploadResponse = serde_json::from_str(&response_text)
            .map_err(|e| SDKError::Json(e))?;

        // Get file URI
        let file_uri = upload_response.urn
            .or(upload_response.file_uri)
            .or_else(|| upload_response.name.map(|name| format!("gs://dy-media-storage/{}", name)))
            .ok_or_else(|| SDKError::attachment("No file URI in upload response".to_string()))?;

        Ok(file_uri)
    }

    /// Send generate request to Gemini API
    async fn send_generate_request(
        &self,
        access_token: &str,
        request_data: &GenerateContentRequest,
    ) -> Result<String> {
        let client_config = self.create_gemini_client(access_token);
        let generate_url = format!("{}/{}:generateContent", client_config.gateway_url, self.config.model_name);

        let mut request_builder = self.client
            .post(&generate_url)
            .timeout(std::time::Duration::from_secs(self.config.timeout))
            .json(request_data);

        // Add headers
        for (key, value) in &client_config.headers {
            request_builder = request_builder.header(key, value);
        }

        let response = request_builder.send().await?;
        let status = response.status();

        if !status.is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Err(SDKError::api(status.as_u16(), format!("API request failed: {}", error_text)));
        }

        let response_text = response.text().await?;
        let gemini_response: GeminiResponse = serde_json::from_str(&response_text)
            .map_err(|e| SDKError::Json(e))?;

        if gemini_response.candidates.is_empty() {
            return Err(SDKError::api(200, "Empty response from API".to_string()));
        }

        // Parse response content
        if let Some(candidate) = gemini_response.candidates.first() {
            if let Some(part) = candidate.content.parts.first() {
                return Ok(part.text.clone());
            }
        }

        Err(SDKError::api(200, "Invalid response format".to_string()))
    }

    /// Create Gemini client configuration
    fn create_gemini_client(&self, access_token: &str) -> ClientConfig {
        let mut headers = HashMap::new();
        headers.insert("Authorization".to_string(), format!("Bearer {}", access_token));
        headers.insert("Content-Type".to_string(), "application/json".to_string());

        let region = self.config.regions.first()
            .unwrap_or(&"us-central1".to_string())
            .clone();

        let gateway_url = format!(
            "https://gateway.ai.cloudflare.com/v1/{}/{}/google-vertex-ai/v1/projects/{}/locations/{}/publishers/google/models",
            self.config.cloudflare_project_id,
            self.config.cloudflare_gateway_id,
            self.config.google_project_id,
            region
        );

        ClientConfig {
            gateway_url,
            headers,
        }
    }
}

// Data structures for Gemini API
#[derive(Debug, Deserialize)]
struct TokenResponse {
    access_token: String,
    expires_in: u64,
}

#[derive(Debug)]
struct ClientConfig {
    gateway_url: String,
    headers: HashMap<String, String>,
}

#[derive(Debug, Deserialize)]
struct UploadResponse {
    file_uri: Option<String>,
    urn: Option<String>,
    name: Option<String>,
}

#[derive(Debug, Serialize)]
struct GenerateContentRequest {
    contents: Vec<ContentPart>,
    #[serde(rename = "generationConfig")]
    generation_config: GenerationConfig,
}

#[derive(Debug, Serialize)]
struct ContentPart {
    role: String,
    parts: Vec<Part>,
}

#[derive(Debug, Serialize)]
#[serde(untagged)]
enum Part {
    Text { text: String },
    FileData {
        #[serde(rename = "fileData")]
        file_data: FileData
    },
    InlineData {
        #[serde(rename = "inlineData")]
        inline_data: InlineData
    },
}

#[derive(Debug, Serialize)]
struct FileData {
    #[serde(rename = "mimeType")]
    mime_type: String,
    #[serde(rename = "fileUri")]
    file_uri: String,
}

#[derive(Debug, Serialize)]
struct InlineData {
    #[serde(rename = "mimeType")]
    mime_type: String,
    data: String,
}

#[derive(Debug, Serialize)]
struct GenerationConfig {
    temperature: f32,
    #[serde(rename = "topK")]
    top_k: u32,
    #[serde(rename = "topP")]
    top_p: f32,
    #[serde(rename = "maxOutputTokens")]
    max_output_tokens: u32,
}

#[derive(Debug, Deserialize)]
struct GeminiResponse {
    candidates: Vec<Candidate>,
}

#[derive(Debug, Deserialize)]
struct Candidate {
    content: Content,
}

#[derive(Debug, Deserialize)]
struct Content {
    parts: Vec<ResponsePart>,
}

#[derive(Debug, Deserialize)]
struct ResponsePart {
    text: String,
}

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

    #[test]
    fn test_sdk_creation() {
        let sdk = GeminiSDK::new().unwrap();
        assert!(!sdk.config.base_url.is_empty());
        assert!(sdk.agents.is_empty());
    }

    #[test]
    fn test_message_request_builders() {
        let req = MessageRequest::text_only("Hello");
        assert_eq!(req.text, "Hello");
        assert!(req.attachments.is_empty());
        assert!(req.agent_id.is_none());

        let req = MessageRequest::with_agent("Hello", "agent-1");
        assert_eq!(req.text, "Hello");
        assert_eq!(req.agent_id, Some("agent-1".to_string()));
    }

    #[tokio::test]
    async fn test_agent_management() {
        let mut sdk = GeminiSDK::new().unwrap();

        // Create agent
        let agent = Agent::new("test", "Test Agent", "You are a test agent");
        sdk.add_agent(agent.clone());

        // Get agent
        let retrieved = sdk.get_agent(&agent.id);
        assert!(retrieved.is_some());
        assert_eq!(retrieved.unwrap().name, "Test Agent");

        // List agents
        let agents = sdk.list_agents();
        assert_eq!(agents.len(), 1);

        // Remove agent
        let removed = sdk.remove_agent(&agent.id);
        assert!(removed.is_some());
        assert!(sdk.get_agent(&agent.id).is_none());
    }
}
