use secrecy::{ExposeSecret, SecretString};
use serde::{Deserialize, Serialize};

use crate::error::AgentError;

#[derive(Clone)]
pub struct LlmClient {
    base_url: String,
    api_key: SecretString,
    model: String,
    embed_model: String,
    http: reqwest::Client,
}

// ── Request types ──────────────────────────────────────────────

#[derive(Serialize, Clone, Debug)]
pub struct ChatMessage {
    pub role: String,
    pub content: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCallRequest>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_call_id: Option<String>,
}

#[derive(Serialize)]
struct ChatCompletionRequest {
    model: String,
    messages: Vec<ChatMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    tools: Option<Vec<ToolDefinitionPayload>>,
}

#[derive(Serialize)]
struct ToolDefinitionPayload {
    r#type: String,
    function: ToolFunctionPayload,
}

#[derive(Serialize)]
struct ToolFunctionPayload {
    name: String,
    description: String,
    parameters: serde_json::Value,
}

// ── Response types ─────────────────────────────────────────────

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

#[derive(Deserialize)]
struct ChatChoice {
    message: ChatResponseMessage,
}

#[derive(Deserialize)]
struct ChatResponseMessage {
    #[serde(default)]
    content: Option<String>,
    #[serde(default)]
    tool_calls: Option<Vec<ToolCallResponse>>,
}

#[derive(Deserialize)]
struct ToolCallResponse {
    id: String,
    function: ToolCallFunction,
}

#[derive(Deserialize)]
struct ToolCallFunction {
    name: String,
    arguments: String,
}

// ── Public types for tool calling ──────────────────────────────

/// Definition of a tool that the LLM can invoke
#[derive(Debug, Clone, Serialize)]
pub struct ToolDefinition {
    pub name: String,
    pub description: String,
    pub parameters: serde_json::Value,
}

/// A tool call request from the LLM
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LlmToolCall {
    pub id: String,
    pub name: String,
    pub arguments: serde_json::Value,
}

/// A tool call in the request message format (for sending back tool_calls in assistant messages)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCallRequest {
    pub id: String,
    pub r#type: String,
    pub function: ToolCallRequestFunction,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCallRequestFunction {
    pub name: String,
    pub arguments: String,
}

/// Response from the LLM — either content or tool calls
#[derive(Debug, Clone)]
pub enum LlmResponse {
    Content(String),
    ToolCalls(Vec<LlmToolCall>),
}

// ── Embedding types ────────────────────────────────────────────

#[derive(Serialize)]
struct EmbeddingRequest {
    model: String,
    input: Vec<String>,
}

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

#[derive(Deserialize)]
struct EmbeddingData {
    embedding: Vec<f64>,
    index: usize,
}

// ── Implementation ─────────────────────────────────────────────

impl LlmClient {
    pub fn new(
        base_url: String,
        api_key: SecretString,
        model: String,
        embed_model: String,
    ) -> Self {
        Self {
            base_url,
            api_key,
            model,
            embed_model,
            http: reqwest::Client::new(),
        }
    }

    pub fn embed_model(&self) -> &str {
        &self.embed_model
    }

    fn chat_url(&self) -> String {
        format!(
            "{}/v1/chat/completions",
            self.base_url.trim_end_matches('/')
        )
    }

    fn auth_header(&self) -> Option<String> {
        let key = self.api_key.expose_secret();
        if key.is_empty() {
            None
        } else {
            Some(format!("Bearer {key}"))
        }
    }

    /// Simple chat: system + user prompt → text response
    pub async fn chat(
        &self,
        system_prompt: &str,
        user_prompt: &str,
        temperature: Option<f64>,
    ) -> Result<String, AgentError> {
        let messages = vec![
            ChatMessage {
                role: "system".to_string(),
                content: Some(system_prompt.to_string()),
                tool_calls: None,
                tool_call_id: None,
            },
            ChatMessage {
                role: "user".to_string(),
                content: Some(user_prompt.to_string()),
                tool_calls: None,
                tool_call_id: None,
            },
        ];

        let request_body = ChatCompletionRequest {
            model: self.model.clone(),
            messages,
            temperature,
            max_tokens: Some(4096),
            tools: None,
        };

        self.send_chat_request(&request_body).await.map(|resp| {
            match resp {
                LlmResponse::Content(c) => c,
                LlmResponse::ToolCalls(_) => String::new(), // shouldn't happen without tools
            }
        })
    }

    /// Chat with a list of (role, content) messages → text response
    #[allow(dead_code)]
    pub async fn chat_with_messages(
        &self,
        messages: Vec<(String, String)>,
        temperature: Option<f64>,
    ) -> Result<String, AgentError> {
        let messages = messages
            .into_iter()
            .map(|(role, content)| ChatMessage {
                role,
                content: Some(content),
                tool_calls: None,
                tool_call_id: None,
            })
            .collect();

        let request_body = ChatCompletionRequest {
            model: self.model.clone(),
            messages,
            temperature,
            max_tokens: Some(4096),
            tools: None,
        };

        self.send_chat_request(&request_body).await.map(|resp| {
            match resp {
                LlmResponse::Content(c) => c,
                LlmResponse::ToolCalls(_) => String::new(),
            }
        })
    }

    /// Chat with tool definitions — returns either content or tool calls.
    /// Use this for the AI pentest orchestrator loop.
    pub async fn chat_with_tools(
        &self,
        messages: Vec<ChatMessage>,
        tools: &[ToolDefinition],
        temperature: Option<f64>,
        max_tokens: Option<u32>,
    ) -> Result<LlmResponse, AgentError> {
        let tool_payloads: Vec<ToolDefinitionPayload> = tools
            .iter()
            .map(|t| ToolDefinitionPayload {
                r#type: "function".to_string(),
                function: ToolFunctionPayload {
                    name: t.name.clone(),
                    description: t.description.clone(),
                    parameters: t.parameters.clone(),
                },
            })
            .collect();

        let request_body = ChatCompletionRequest {
            model: self.model.clone(),
            messages,
            temperature,
            max_tokens: Some(max_tokens.unwrap_or(8192)),
            tools: if tool_payloads.is_empty() {
                None
            } else {
                Some(tool_payloads)
            },
        };

        self.send_chat_request(&request_body).await
    }

    /// Internal method to send a chat completion request and parse the response
    async fn send_chat_request(
        &self,
        request_body: &ChatCompletionRequest,
    ) -> Result<LlmResponse, AgentError> {
        let mut req = self
            .http
            .post(&self.chat_url())
            .header("content-type", "application/json")
            .json(request_body);

        if let Some(auth) = self.auth_header() {
            req = req.header("Authorization", auth);
        }

        let resp = req
            .send()
            .await
            .map_err(|e| AgentError::Other(format!("LiteLLM request failed: {e}")))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().await.unwrap_or_default();
            return Err(AgentError::Other(format!(
                "LiteLLM returned {status}: {body}"
            )));
        }

        let body: ChatCompletionResponse = resp
            .json()
            .await
            .map_err(|e| AgentError::Other(format!("Failed to parse LiteLLM response: {e}")))?;

        let choice = body
            .choices
            .first()
            .ok_or_else(|| AgentError::Other("Empty response from LiteLLM".to_string()))?;

        // Check for tool calls first
        if let Some(tool_calls) = &choice.message.tool_calls {
            if !tool_calls.is_empty() {
                let calls: Vec<LlmToolCall> = tool_calls
                    .iter()
                    .map(|tc| {
                        let arguments = serde_json::from_str(&tc.function.arguments)
                            .unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
                        LlmToolCall {
                            id: tc.id.clone(),
                            name: tc.function.name.clone(),
                            arguments,
                        }
                    })
                    .collect();
                return Ok(LlmResponse::ToolCalls(calls));
            }
        }

        // Otherwise return content
        let content = choice
            .message
            .content
            .clone()
            .unwrap_or_default();
        Ok(LlmResponse::Content(content))
    }

    /// Generate embeddings for a batch of texts
    pub async fn embed(&self, texts: Vec<String>) -> Result<Vec<Vec<f64>>, AgentError> {
        let url = format!("{}/v1/embeddings", self.base_url.trim_end_matches('/'));

        let request_body = EmbeddingRequest {
            model: self.embed_model.clone(),
            input: texts,
        };

        let mut req = self
            .http
            .post(&url)
            .header("content-type", "application/json")
            .json(&request_body);

        if let Some(auth) = self.auth_header() {
            req = req.header("Authorization", auth);
        }

        let resp = req
            .send()
            .await
            .map_err(|e| AgentError::Other(format!("Embedding request failed: {e}")))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().await.unwrap_or_default();
            return Err(AgentError::Other(format!(
                "Embedding API returned {status}: {body}"
            )));
        }

        let body: EmbeddingResponse = resp
            .json()
            .await
            .map_err(|e| AgentError::Other(format!("Failed to parse embedding response: {e}")))?;

        let mut data = body.data;
        data.sort_by_key(|d| d.index);

        Ok(data.into_iter().map(|d| d.embedding).collect())
    }
}
