// SSE (Server-Sent Events) Transport
use anyhow::{anyhow, Result};
use async_trait::async_trait;
use bytes::Bytes;
use futures::StreamExt;
use reqwest::{header, Client};
use serde_json::{json, Value};
use std::collections::HashMap;
use tokio::sync::mpsc;

use crate::auth::AuthConfig;
use crate::providers::base::Provider;
use crate::providers::sse::SseProvider;
use crate::tools::Tool;
use crate::transports::{
    stream::{boxed_channel_stream, StreamResult},
    ClientTransport,
};

pub struct SseTransport {
    client: Client,
}

impl SseTransport {
    pub fn new() -> Self {
        Self {
            client: Client::new(),
        }
    }

    fn build_payload(&self, prov: &SseProvider, args: HashMap<String, Value>) -> Value {
        if let Some(body_field) = &prov.body_field {
            json!({ body_field: args })
        } else {
            json!(args)
        }
    }

    fn apply_headers(
        &self,
        request: reqwest::RequestBuilder,
        prov: &SseProvider,
        extra_accept: Option<&str>,
    ) -> reqwest::RequestBuilder {
        let mut builder = request;
        builder = builder.header("Accept", extra_accept.unwrap_or("application/json"));
        if let Some(headers) = &prov.headers {
            for (k, v) in headers {
                builder = builder.header(k, v);
            }
        }
        builder
    }

    fn apply_auth(
        &self,
        builder: reqwest::RequestBuilder,
        auth: &AuthConfig,
    ) -> Result<reqwest::RequestBuilder> {
        match auth {
            AuthConfig::ApiKey(api_key) => {
                let location = api_key.location.to_ascii_lowercase();
                match location.as_str() {
                    "header" => Ok(builder.header(&api_key.var_name, &api_key.api_key)),
                    "query" => {
                        Ok(builder.query(&[(api_key.var_name.clone(), api_key.api_key.clone())]))
                    }
                    "cookie" => {
                        let cookie_value = format!("{}={}", api_key.var_name, api_key.api_key);
                        Ok(builder.header(header::COOKIE, cookie_value))
                    }
                    other => Err(anyhow!("Unsupported API key location: {}", other)),
                }
            }
            AuthConfig::Basic(basic) => {
                Ok(builder.basic_auth(&basic.username, Some(&basic.password)))
            }
            AuthConfig::OAuth2(_) => Err(anyhow!(
                "OAuth2 auth is not yet supported by the SSE transport"
            )),
        }
    }

    fn parse_tools_from_body(&self, body: &str) -> Vec<Tool> {
        if let Ok(manifest) = serde_json::from_str::<Value>(body) {
            if let Some(tools) = manifest.get("tools").and_then(|v| v.as_array()) {
                let mut parsed = Vec::new();
                for t in tools {
                    if let Ok(tool) = serde_json::from_value::<Tool>(t.clone()) {
                        parsed.push(tool);
                    }
                }
                return parsed;
            }
        }
        vec![]
    }

    fn spawn_sse_reader(
        &self,
        mut stream: impl futures::Stream<Item = Result<Bytes, reqwest::Error>> + Send + Unpin + 'static,
    ) -> mpsc::Receiver<Result<Value>> {
        let (tx, rx) = mpsc::channel(16);
        tokio::spawn(async move {
            let mut buffer = String::new();
            let mut data_buf = String::new();

            while let Some(chunk_res) = stream.next().await {
                match chunk_res {
                    Ok(bytes) => {
                        buffer.push_str(&String::from_utf8_lossy(&bytes));
                        loop {
                            if let Some(pos) = buffer.find('\n') {
                                let mut line = buffer[..pos].to_string();
                                buffer.drain(..=pos);
                                line = line.trim_end_matches('\r').to_string();

                                if line.starts_with("data: ") {
                                    if !data_buf.is_empty() {
                                        data_buf.push('\n');
                                    }
                                    data_buf.push_str(&line[6..]);
                                } else if line.is_empty() {
                                    if !data_buf.is_empty() {
                                        let parsed = serde_json::from_str::<Value>(&data_buf)
                                            .map_err(|e| {
                                                anyhow!("Failed to parse SSE data: {}", e)
                                            });
                                        if tx.send(parsed).await.is_err() {
                                            return;
                                        }
                                        data_buf.clear();
                                    }
                                }
                            } else {
                                break;
                            }
                        }
                    }
                    Err(err) => {
                        let _ = tx
                            .send(Err(anyhow!("Error reading SSE stream: {}", err)))
                            .await;
                        return;
                    }
                }
            }

            // Flush trailing data if present
            if !data_buf.is_empty() {
                let _ = tx
                    .send(
                        serde_json::from_str::<Value>(&data_buf)
                            .map_err(|e| anyhow!("Failed to parse SSE data: {}", e)),
                    )
                    .await;
            }
        });
        rx
    }
}

#[async_trait]
impl ClientTransport for SseTransport {
    async fn register_tool_provider(&self, prov: &dyn Provider) -> Result<Vec<Tool>> {
        let sse_prov = prov
            .as_any()
            .downcast_ref::<SseProvider>()
            .ok_or_else(|| anyhow!("Provider is not an SseProvider"))?;

        let mut request = self
            .client
            .get(&sse_prov.url)
            .header("Accept", "application/json");
        request = self.apply_headers(request, sse_prov, None);
        if let Some(auth) = &sse_prov.base.auth {
            request = self.apply_auth(request, auth)?;
        }

        let response = request.send().await?;
        if !response.status().is_success() {
            return Err(anyhow!(
                "Failed to fetch tools from {}: {}",
                sse_prov.url,
                response.status()
            ));
        }

        let body = response.text().await?;
        Ok(self.parse_tools_from_body(&body))
    }

    async fn deregister_tool_provider(&self, _prov: &dyn Provider) -> Result<()> {
        Ok(())
    }

    async fn call_tool(
        &self,
        tool_name: &str,
        args: HashMap<String, Value>,
        prov: &dyn Provider,
    ) -> Result<Value> {
        // Use streaming parser and eagerly collect values.
        let mut stream = self.call_tool_stream(tool_name, args, prov).await?;
        let mut items = Vec::new();
        while let Some(item) = stream.next().await? {
            items.push(item);
        }
        stream.close().await?;
        Ok(Value::Array(items))
    }

    async fn call_tool_stream(
        &self,
        tool_name: &str,
        args: HashMap<String, Value>,
        prov: &dyn Provider,
    ) -> Result<Box<dyn StreamResult>> {
        let sse_prov = prov
            .as_any()
            .downcast_ref::<SseProvider>()
            .ok_or_else(|| anyhow!("Provider is not an SseProvider"))?;

        let url = format!("{}/{}", sse_prov.url.trim_end_matches('/'), tool_name);
        let payload = self.build_payload(sse_prov, args);

        let mut request = self
            .client
            .post(url)
            .header("Content-Type", "application/json");
        request = self.apply_headers(request, sse_prov, Some("text/event-stream"));
        if let Some(auth) = &sse_prov.base.auth {
            request = self.apply_auth(request, auth)?;
        }
        let response = request.json(&payload).send().await?;

        if !response.status().is_success() {
            return Err(anyhow!("SSE request failed: {}", response.status()));
        }

        let rx = self.spawn_sse_reader(response.bytes_stream());
        Ok(boxed_channel_stream(rx, None))
    }
}
