use std::pin::Pin;

use async_trait::async_trait;
use futures::stream::{self, Stream};
use futures::StreamExt;
use motosan_agent_tool::ToolDef;

use crate::error::Result;
use crate::message::Message;

/// Token usage statistics from a single LLM call.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct TokenUsage {
    pub input_tokens: u64,
    pub output_tokens: u64,
}

impl TokenUsage {
    /// Accumulate another usage into this one.
    pub fn accumulate(&mut self, other: TokenUsage) {
        self.input_tokens += other.input_tokens;
        self.output_tokens += other.output_tokens;
    }
}

/// A single tool call within an [`LlmResponse::ToolCalls`] batch.
#[derive(Debug, Clone, PartialEq)]
pub struct ToolCallItem {
    /// Unique identifier for the tool call (used to correlate the result).
    pub id: String,
    /// Name of the tool to invoke.
    pub name: String,
    /// Parsed JSON arguments for the tool.
    pub args: serde_json::Value,
}

/// The complete (non-streaming) response from an LLM backend.
#[derive(Debug, Clone, PartialEq)]
pub enum LlmResponse {
    /// A plain text reply with no tool invocations.
    Message(String),
    /// The model wants to call one or more tools.
    ToolCalls(Vec<ToolCallItem>),
}

impl LlmResponse {
    /// Convenience constructor for a single tool call.
    pub fn single_tool_call(id: String, name: String, args: serde_json::Value) -> Self {
        Self::ToolCalls(vec![ToolCallItem { id, name, args }])
    }
}

/// The complete output from a single LLM call, including token usage.
#[derive(Debug, Clone)]
pub struct ChatOutput {
    /// The response content (text or tool calls).
    pub response: LlmResponse,
    /// Token usage for this call, if reported by the backend.
    pub usage: Option<TokenUsage>,
}

impl ChatOutput {
    /// Wrap an `LlmResponse` with no usage information.
    pub fn new(response: LlmResponse) -> Self {
        Self {
            response,
            usage: None,
        }
    }

    /// Wrap an `LlmResponse` with usage information.
    pub fn with_usage(response: LlmResponse, usage: TokenUsage) -> Self {
        Self {
            response,
            usage: Some(usage),
        }
    }
}

/// A chunk from a streaming LLM response.
#[derive(Debug, Clone, PartialEq)]
pub enum StreamChunk {
    /// A partial text delta from the LLM.
    TextDelta(String),
    /// The final complete response (text or tool calls).
    Done(LlmResponse),
    /// Token usage reported at the end of a streaming response.
    Usage(TokenUsage),
}

/// Trait that every LLM backend must implement.
///
/// The agent loop is generic over this trait so callers can plug in OpenAI,
/// Anthropic, local models, or test doubles without changing loop logic.
///
/// Designed to work with `Arc<dyn LlmClient>` for shared ownership across
/// async tasks.
#[async_trait]
pub trait LlmClient: Send + Sync {
    /// Send a conversation and available tools, and receive a complete response.
    async fn chat(&self, messages: &[Message], tools: &[ToolDef]) -> Result<ChatOutput>;

    /// Streaming variant of [`chat`](Self::chat).
    ///
    /// Returns a stream of [`StreamChunk`] items — zero or more `TextDelta`
    /// chunks followed by a final `Done` with the complete [`LlmResponse`],
    /// optionally followed by a `Usage` chunk with token counts.
    ///
    /// The default implementation delegates to [`chat`](Self::chat) and
    /// yields `Done` and (if available) `Usage` chunks, so backends only
    /// need to override this if they support true streaming.
    fn chat_stream<'a>(
        &'a self,
        messages: &'a [Message],
        tools: &'a [ToolDef],
    ) -> Pin<Box<dyn Stream<Item = Result<StreamChunk>> + Send + 'a>> {
        Box::pin(
            stream::once(async move {
                let output = self.chat(messages, tools).await?;
                Ok(output)
            })
            .flat_map(|result: Result<ChatOutput>| match result {
                Ok(output) => {
                    let mut chunks = vec![Ok(StreamChunk::Done(output.response))];
                    if let Some(usage) = output.usage {
                        chunks.push(Ok(StreamChunk::Usage(usage)));
                    }
                    stream::iter(chunks)
                }
                Err(e) => stream::iter(vec![Err(e)]),
            }),
        )
    }
}

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

    /// Compile-time test: `Arc<dyn LlmClient>` must be valid.
    /// This function is never called but ensures the trait is object-safe
    /// and compatible with `Arc`.
    #[allow(dead_code)]
    fn assert_arc_compatible(_client: Arc<dyn LlmClient>) {}

    #[test]
    fn llm_response_message_variant() {
        let resp = LlmResponse::Message("hello".to_string());
        assert_eq!(resp, LlmResponse::Message("hello".to_string()));
    }

    #[test]
    fn llm_response_single_tool_call() {
        let resp = LlmResponse::single_tool_call(
            "call_1".to_string(),
            "search".to_string(),
            serde_json::json!({"query": "rust"}),
        );
        match &resp {
            LlmResponse::ToolCalls(items) => {
                assert_eq!(items.len(), 1);
                assert_eq!(items[0].id, "call_1");
                assert_eq!(items[0].name, "search");
                assert_eq!(items[0].args, serde_json::json!({"query": "rust"}));
            }
            _ => panic!("expected ToolCalls variant"),
        }
    }

    #[test]
    fn llm_response_tool_calls_variant() {
        let calls = vec![
            ToolCallItem {
                id: "call_1".to_string(),
                name: "search".to_string(),
                args: serde_json::json!({"q": "a"}),
            },
            ToolCallItem {
                id: "call_2".to_string(),
                name: "fetch".to_string(),
                args: serde_json::json!({"url": "http://example.com"}),
            },
        ];
        let resp = LlmResponse::ToolCalls(calls.clone());
        match resp {
            LlmResponse::ToolCalls(items) => {
                assert_eq!(items.len(), 2);
                assert_eq!(items[0].name, "search");
                assert_eq!(items[1].name, "fetch");
            }
            _ => panic!("expected ToolCalls variant"),
        }
    }

    #[test]
    fn tool_call_item_equality() {
        let a = ToolCallItem {
            id: "1".to_string(),
            name: "test".to_string(),
            args: serde_json::json!(null),
        };
        let b = a.clone();
        assert_eq!(a, b);
    }

    #[test]
    fn token_usage_default_is_zero() {
        let u = TokenUsage::default();
        assert_eq!(u.input_tokens, 0);
        assert_eq!(u.output_tokens, 0);
    }

    #[test]
    fn token_usage_accumulate() {
        let mut total = TokenUsage {
            input_tokens: 10,
            output_tokens: 20,
        };
        total.accumulate(TokenUsage {
            input_tokens: 5,
            output_tokens: 15,
        });
        assert_eq!(total.input_tokens, 15);
        assert_eq!(total.output_tokens, 35);
    }

    #[test]
    fn chat_output_new_has_no_usage() {
        let output = ChatOutput::new(LlmResponse::Message("hi".into()));
        assert!(output.usage.is_none());
    }

    #[test]
    fn chat_output_with_usage() {
        let usage = TokenUsage {
            input_tokens: 100,
            output_tokens: 50,
        };
        let output = ChatOutput::with_usage(LlmResponse::Message("hi".into()), usage);
        assert_eq!(output.usage, Some(usage));
    }
}
