//! Common types used across all AIX providers.
//!
//! This module defines the shared data structures for requests, responses,
//! and configuration that are used by all AI providers.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// The role of a message in a conversation.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    /// System message that sets the behavior of the assistant
    System,
    /// User message
    User,
    /// Assistant message
    Assistant,
    /// Tool/function call message
    Tool,
}

impl Role {
    /// Get the string representation of the role.
    pub fn as_str(&self) -> &'static str {
        match self {
            Role::System => "system",
            Role::User => "user",
            Role::Assistant => "assistant",
            Role::Tool => "tool",
        }
    }
}

/// A single message in a conversation.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ChatMessage {
    /// The role of the message author
    pub role: Role,
    /// The content of the message
    pub content: String,
    /// Optional name of the message author
    pub name: Option<String>,
}

impl ChatMessage {
    /// Create a new user message.
    pub fn user<S: Into<String>>(content: S) -> Self {
        Self {
            role: Role::User,
            content: content.into(),
            name: None,
        }
    }

    /// Create a new assistant message.
    pub fn assistant<S: Into<String>>(content: S) -> Self {
        Self {
            role: Role::Assistant,
            content: content.into(),
            name: None,
        }
    }

    /// Create a new system message.
    pub fn system<S: Into<String>>(content: S) -> Self {
        Self {
            role: Role::System,
            content: content.into(),
            name: None,
        }
    }

    /// Create a new tool message.
    pub fn tool<S: Into<String>>(content: S) -> Self {
        Self {
            role: Role::Tool,
            content: content.into(),
            name: None,
        }
    }

    /// Set the name of the message author.
    pub fn with_name<S: Into<String>>(mut self, name: S) -> Self {
        self.name = Some(name.into());
        self
    }
}

/// Configuration for model behavior.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ModelConfig {
    /// Controls randomness (0.0 = deterministic, 2.0 = maximum randomness)
    pub temperature: Option<f32>,
    /// Maximum number of tokens to generate
    pub max_tokens: Option<u32>,
    /// Nucleus sampling parameter (0.0 = no diversity, 1.0 = maximum diversity)
    pub top_p: Option<f32>,
    /// Frequency penalty for repeated tokens (-2.0 to 2.0)
    pub frequency_penalty: Option<f32>,
    /// Presence penalty for new topics (-2.0 to 2.0)
    pub presence_penalty: Option<f32>,
    /// Sequences that will stop generation when encountered
    pub stop_sequences: Vec<String>,
    /// Additional metadata for the request
    pub metadata: HashMap<String, String>,
}

impl ModelConfig {
    /// Create a new default model configuration.
    pub fn new() -> Self {
        Self::default()
    }

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

    /// Set the maximum number of tokens.
    pub fn with_max_tokens(mut self, max_tokens: u32) -> Self {
        self.max_tokens = Some(max_tokens);
        self
    }

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

    /// Set the frequency penalty.
    pub fn with_frequency_penalty(mut self, frequency_penalty: f32) -> Self {
        self.frequency_penalty = Some(frequency_penalty);
        self
    }

    /// Set the presence penalty.
    pub fn with_presence_penalty(mut self, presence_penalty: f32) -> Self {
        self.presence_penalty = Some(presence_penalty);
        self
    }

    /// Set stop sequences.
    pub fn with_stop_sequences(mut self, stop_sequences: Vec<String>) -> Self {
        self.stop_sequences = stop_sequences;
        self
    }

    /// Add metadata.
    pub fn with_metadata<K: Into<String>, V: Into<String>>(mut self, key: K, value: V) -> Self {
        self.metadata.insert(key.into(), value.into());
        self
    }
}

/// Builder for creating chat requests.
#[derive(Debug, Clone)]
pub struct ChatRequestBuilder {
    model: String,
    messages: Vec<ChatMessage>,
    config: ModelConfig,
    stream: bool,
}

impl ChatRequestBuilder {
    /// Create a new chat request builder.
    pub fn new<S: Into<String>>(model: S) -> Self {
        Self {
            model: model.into(),
            messages: Vec::new(),
            config: ModelConfig::default(),
            stream: false,
        }
    }

    /// Add a message to the request.
    pub fn message(mut self, message: ChatMessage) -> Self {
        self.messages.push(message);
        self
    }

    /// Add multiple messages to the request.
    pub fn messages(mut self, messages: Vec<ChatMessage>) -> Self {
        self.messages.extend(messages);
        self
    }

    /// Set the model configuration.
    pub fn config(mut self, config: ModelConfig) -> Self {
        self.config = config;
        self
    }

    /// Enable or disable streaming.
    pub fn stream(mut self, stream: bool) -> Self {
        self.stream = stream;
        self
    }

    /// Build the chat request.
    pub fn build(self) -> ChatRequest {
        ChatRequest {
            model: self.model,
            messages: self.messages,
            config: self.config,
            stream: self.stream,
        }
    }
}

/// A chat completion request.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatRequest {
    /// The model to use for completion
    pub model: String,
    /// The conversation history
    pub messages: Vec<ChatMessage>,
    /// Model configuration parameters
    pub config: ModelConfig,
    /// Whether to stream the response
    pub stream: bool,
}

impl ChatRequest {
    /// Create a new chat request with the given model.
    pub fn new<S: Into<String>>(model: S) -> ChatRequestBuilder {
        ChatRequestBuilder::new(model)
    }

    /// Create a new chat request with a single user message.
    pub fn simple<S: Into<String>, M: Into<String>>(model: S, message: M) -> Self {
        Self::new(model).message(ChatMessage::user(message)).build()
    }
}

/// Token usage information.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage {
    /// Number of tokens in the prompt
    pub prompt_tokens: u32,
    /// Number of tokens in the completion
    pub completion_tokens: u32,
    /// Total number of tokens
    pub total_tokens: u32,
}

impl Usage {
    /// Create new usage information.
    pub fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
        Self {
            prompt_tokens,
            completion_tokens,
            total_tokens: prompt_tokens + completion_tokens,
        }
    }
}

/// A chat completion response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatResponse {
    /// Unique identifier for the completion
    pub id: String,
    /// The model that generated the completion
    pub model: String,
    /// The generated content
    pub content: String,
    /// The role of the assistant (typically "assistant")
    pub role: Role,
    /// Token usage information
    pub usage: Usage,
    /// The reason the generation stopped
    pub finish_reason: Option<String>,
}

impl ChatResponse {
    /// Create a new chat response.
    pub fn new<S: Into<String>>(
        id: S,
        model: S,
        content: S,
        role: Role,
        usage: Usage,
    ) -> Self {
        Self {
            id: id.into(),
            model: model.into(),
            content: content.into(),
            role,
            usage,
            finish_reason: None,
        }
    }

    /// Set the finish reason.
    pub fn with_finish_reason<S: Into<String>>(mut self, finish_reason: S) -> Self {
        self.finish_reason = Some(finish_reason.into());
        self
    }
}

/// A chunk from a streaming response.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct StreamChunk {
    /// Unique identifier for the chunk
    pub id: String,
    /// The delta content (new tokens)
    pub delta: String,
    /// The reason the generation stopped (if applicable)
    pub finish_reason: Option<String>,
}

impl StreamChunk {
    /// Create a new stream chunk.
    pub fn new<S: Into<String>>(id: S, delta: S) -> Self {
        Self {
            id: id.into(),
            delta: delta.into(),
            finish_reason: None,
        }
    }

    /// Create a final chunk with finish reason.
    pub fn with_finish_reason<S: Into<String>>(id: S, finish_reason: S) -> Self {
        Self {
            id: id.into(),
            delta: String::new(),
            finish_reason: Some(finish_reason.into()),
        }
    }

    /// Check if this is the final chunk.
    pub fn is_final(&self) -> bool {
        self.finish_reason.is_some()
    }
}

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

    #[test]
    fn test_chat_message_creation() {
        let msg = ChatMessage::user("Hello").with_name("Alice");
        assert_eq!(msg.role, Role::User);
        assert_eq!(msg.content, "Hello");
        assert_eq!(msg.name, Some("Alice".to_string()));
    }

    #[test]
    fn test_model_config_builder() {
        let config = ModelConfig::new()
            .with_temperature(0.7)
            .with_max_tokens(1000)
            .with_metadata("key", "value");

        assert_eq!(config.temperature, Some(0.7));
        assert_eq!(config.max_tokens, Some(1000));
        assert_eq!(config.metadata.get("key"), Some(&"value".to_string()));
    }

    #[test]
    fn test_chat_request_builder() {
        let request = ChatRequest::new("gpt-4")
            .message(ChatMessage::system("You are a helpful assistant."))
            .message(ChatMessage::user("Hello!"))
            .stream(true)
            .build();

        assert_eq!(request.model, "gpt-4");
        assert_eq!(request.messages.len(), 2);
        assert!(request.stream);
    }

    #[test]
    fn test_simple_chat_request() {
        let request = ChatRequest::simple("gpt-4", "Hello, world!");
        assert_eq!(request.model, "gpt-4");
        assert_eq!(request.messages.len(), 1);
        assert_eq!(request.messages[0].content, "Hello, world!");
        assert_eq!(request.messages[0].role, Role::User);
    }

    #[test]
    fn test_usage() {
        let usage = Usage::new(10, 20);
        assert_eq!(usage.prompt_tokens, 10);
        assert_eq!(usage.completion_tokens, 20);
        assert_eq!(usage.total_tokens, 30);
    }

    #[test]
    fn test_stream_chunk() {
        let chunk = StreamChunk::new("123", "Hello");
        assert_eq!(chunk.id, "123");
        assert_eq!(chunk.delta, "Hello");
        assert!(!chunk.is_final());

        let final_chunk = StreamChunk::with_finish_reason("123", "stop");
        assert!(final_chunk.is_final());
    }
}
