use std::collections::BTreeMap;
use std::error::Error;
use std::pin::Pin;

use futures_util::{Stream, StreamExt};
use serde::{Deserialize, Serialize};

use crate::bus::ProviderReasoningState;

use super::{ChatCompletionResponse, ToolCall, Usage};

pub type DynProviderError = Box<dyn Error + Send + Sync>;
pub type ProviderStreamItem = Result<ProviderChunk, DynProviderError>;
pub type ProviderStream = Pin<Box<dyn Stream<Item = ProviderStreamItem> + Send>>;

pub(crate) fn diagnostic_request_summary(
    model: &str,
    message_count: usize,
    tool_count: usize,
) -> String {
    serde_json::json!({
        "model": model,
        "message_count": message_count,
        "tool_count": tool_count,
        "stream": true,
    })
    .to_string()
}

/// Incremental framing shared by SSE-based providers. It accepts arbitrary
/// byte/UTF-8 boundaries and returns only joined `data:` payloads.
#[derive(Default)]
pub(crate) struct SseFramer {
    buffer: Vec<u8>,
}

impl SseFramer {
    pub(crate) fn push(&mut self, bytes: &[u8]) -> Result<Vec<String>, std::string::FromUtf8Error> {
        self.buffer.extend_from_slice(bytes);
        self.drain_frames(false)
    }

    pub(crate) fn finish(&mut self) -> Result<Vec<String>, std::string::FromUtf8Error> {
        self.drain_frames(true)
    }

    fn drain_frames(&mut self, finish: bool) -> Result<Vec<String>, std::string::FromUtf8Error> {
        let mut frames = Vec::new();
        while let Some((position, delimiter_len)) = find_sse_delimiter(&self.buffer) {
            let frame = self.buffer.drain(..position).collect::<Vec<_>>();
            self.buffer.drain(..delimiter_len);
            if let Some(data) = sse_data(frame)? {
                frames.push(data);
            }
        }
        if finish && !self.buffer.is_empty() {
            let frame = std::mem::take(&mut self.buffer);
            if let Some(data) = sse_data(frame)? {
                frames.push(data);
            }
        }
        Ok(frames)
    }
}

fn find_sse_delimiter(buffer: &[u8]) -> Option<(usize, usize)> {
    let lf = buffer.windows(2).position(|window| window == b"\n\n");
    let crlf = buffer.windows(4).position(|window| window == b"\r\n\r\n");
    match (lf, crlf) {
        (Some(left), Some(right)) if left <= right => Some((left, 2)),
        (Some(_), Some(right)) => Some((right, 4)),
        (Some(position), None) => Some((position, 2)),
        (None, Some(position)) => Some((position, 4)),
        (None, None) => None,
    }
}

fn sse_data(frame: Vec<u8>) -> Result<Option<String>, std::string::FromUtf8Error> {
    let frame = String::from_utf8(frame)?;
    let data = frame
        .lines()
        .filter_map(|line| {
            line.strip_prefix("data:")
                .map(|value| value.strip_prefix(' ').unwrap_or(value))
        })
        .collect::<Vec<_>>()
        .join("\n");
    Ok((!data.is_empty()).then_some(data))
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FinishReason {
    Stop,
    ToolCalls,
    Length,
    ContentFilter,
    Other(String),
}

impl FinishReason {
    pub fn from_provider(value: &str) -> Self {
        match value {
            "stop" | "end_turn" | "stop_sequence" => Self::Stop,
            "tool_calls" | "tool_use" => Self::ToolCalls,
            "length" | "max_tokens" => Self::Length,
            "content_filter" => Self::ContentFilter,
            other => Self::Other(other.to_string()),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum ProviderChunk {
    Metadata {
        id: String,
        model: String,
    },
    Text(String),
    Reasoning(String),
    ToolCallStart {
        index: usize,
        id: Option<String>,
        name: Option<String>,
    },
    ToolCallArguments {
        index: usize,
        delta: String,
    },
    ProviderState(ProviderReasoningState),
    Usage(Usage),
    Done(FinishReason),
}

#[derive(Default)]
struct PartialToolCall {
    id: Option<String>,
    name: Option<String>,
    arguments: String,
}

#[derive(Default)]
pub struct ProviderResponseAccumulator {
    id: String,
    model: String,
    content: String,
    reasoning_content: String,
    provider_state: Option<ProviderReasoningState>,
    usage: Usage,
    tool_calls: BTreeMap<usize, PartialToolCall>,
}

impl ProviderResponseAccumulator {
    pub fn push(&mut self, chunk: ProviderChunk) {
        match chunk {
            ProviderChunk::Metadata {
                id: response_id,
                model: response_model,
            } => {
                self.id = response_id;
                self.model = response_model;
            }
            ProviderChunk::Text(delta) => self.content.push_str(&delta),
            ProviderChunk::Reasoning(delta) => self.reasoning_content.push_str(&delta),
            ProviderChunk::ToolCallStart {
                index,
                id: call_id,
                name,
            } => {
                let partial = self.tool_calls.entry(index).or_default();
                if call_id.is_some() {
                    partial.id = call_id;
                }
                if name.is_some() {
                    partial.name = name;
                }
            }
            ProviderChunk::ToolCallArguments { index, delta } => {
                self.tool_calls
                    .entry(index)
                    .or_default()
                    .arguments
                    .push_str(&delta);
            }
            ProviderChunk::ProviderState(state) => self.provider_state = Some(state),
            ProviderChunk::Usage(value) => self.usage = value,
            ProviderChunk::Done(_) => {}
        }
    }

    pub fn finish(self) -> ChatCompletionResponse {
        let tool_calls = self
            .tool_calls
            .into_iter()
            .map(|(index, partial)| ToolCall {
                id: partial.id.unwrap_or_else(|| format!("tool_call_{index}")),
                name: partial.name.unwrap_or_default(),
                arguments: serde_json::from_str(&partial.arguments)
                    .unwrap_or(serde_json::Value::Null),
            })
            .collect();

        ChatCompletionResponse {
            id: self.id,
            model: self.model,
            content: self.content,
            reasoning_content: (!self.reasoning_content.is_empty())
                .then_some(self.reasoning_content),
            provider_state: self.provider_state,
            tool_calls,
            usage: self.usage,
        }
    }
}

pub async fn collect_provider_stream(
    mut provider_stream: ProviderStream,
) -> Result<ChatCompletionResponse, DynProviderError> {
    let mut accumulator = ProviderResponseAccumulator::default();
    while let Some(chunk) = provider_stream.next().await {
        accumulator.push(chunk?);
    }
    Ok(accumulator.finish())
}

#[cfg(test)]
pub fn provider_stream_for_test(response: ChatCompletionResponse) -> ProviderStream {
    let finish_reason = if response.tool_calls.is_empty() {
        FinishReason::Stop
    } else {
        FinishReason::ToolCalls
    };
    let mut chunks = vec![ProviderChunk::Metadata {
        id: response.id,
        model: response.model,
    }];
    if let Some(reasoning) = response.reasoning_content {
        chunks.push(ProviderChunk::Reasoning(reasoning));
    }
    if !response.content.is_empty() {
        chunks.push(ProviderChunk::Text(response.content));
    }
    for (index, call) in response.tool_calls.into_iter().enumerate() {
        chunks.push(ProviderChunk::ToolCallStart {
            index,
            id: Some(call.id),
            name: Some(call.name),
        });
        chunks.push(ProviderChunk::ToolCallArguments {
            index,
            delta: serde_json::to_string(&call.arguments).unwrap_or_else(|_| "null".to_string()),
        });
    }
    if let Some(state) = response.provider_state {
        chunks.push(ProviderChunk::ProviderState(state));
    }
    chunks.push(ProviderChunk::Usage(response.usage));
    chunks.push(ProviderChunk::Done(finish_reason));
    Box::pin(futures_util::stream::iter(chunks.into_iter().map(Ok)))
}

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

    #[test]
    fn diagnostic_summary_contains_counts_without_message_content() {
        let summary = diagnostic_request_summary("model", 3, 2);

        assert_eq!(
            serde_json::from_str::<serde_json::Value>(&summary).unwrap(),
            serde_json::json!({
                "model": "model",
                "message_count": 3,
                "tool_count": 2,
                "stream": true
            })
        );
        assert!(!summary.contains("reasoning"));
        assert!(!summary.contains("signature"));
    }

    #[tokio::test]
    async fn collector_assembles_interleaved_tool_argument_fragments() {
        let chunks = vec![
            ProviderChunk::Metadata {
                id: "response".into(),
                model: "model".into(),
            },
            ProviderChunk::Reasoning("why".into()),
            ProviderChunk::Text("answer".into()),
            ProviderChunk::ToolCallStart {
                index: 1,
                id: Some("second".into()),
                name: Some("b".into()),
            },
            ProviderChunk::ToolCallArguments {
                index: 1,
                delta: "{\"n\":".into(),
            },
            ProviderChunk::ToolCallStart {
                index: 0,
                id: Some("first".into()),
                name: Some("a".into()),
            },
            ProviderChunk::ToolCallArguments {
                index: 0,
                delta: "{}".into(),
            },
            ProviderChunk::ToolCallArguments {
                index: 1,
                delta: "2}".into(),
            },
            ProviderChunk::Done(FinishReason::ToolCalls),
        ];

        let response = collect_provider_stream(Box::pin(stream::iter(chunks.into_iter().map(Ok))))
            .await
            .unwrap();

        assert_eq!(response.content, "answer");
        assert_eq!(response.reasoning_content.as_deref(), Some("why"));
        assert_eq!(response.tool_calls.len(), 2);
        assert_eq!(response.tool_calls[0].id, "first");
        assert_eq!(
            response.tool_calls[1].arguments,
            serde_json::json!({"n": 2})
        );
    }
}
