use futures::StreamExt;
use serde_json::Value;
use std::path::PathBuf;
use tokio::fs;

use crate::proxy::config::DebugLoggingConfig;

fn build_filename(prefix: &str, trace_id: Option<&str>) -> String {
    let ts = chrono::Utc::now().format("%Y%m%d_%H%M%S%.3f");
    let tid = trace_id.unwrap_or("unknown");
    format!("{}_{}_{}.json", ts, tid, prefix)
}

fn resolve_output_dir(cfg: &DebugLoggingConfig) -> Option<PathBuf> {
    if let Some(dir) = cfg.output_dir.as_ref() {
        return Some(PathBuf::from(dir));
    }
    if let Ok(data_dir) = crate::modules::auth::account::get_data_dir() {
        return Some(data_dir.join("debug_logs"));
    }
    None
}

pub async fn write_debug_payload(
    cfg: &DebugLoggingConfig,
    trace_id: Option<&str>,
    prefix: &str,
    payload: &Value,
) {
    if !cfg.enabled {
        return;
    }

    let output_dir = match resolve_output_dir(cfg) {
        Some(dir) => dir,
        None => {
            tracing::warn!("[Debug-Log] Enabled but output_dir is not available.");
            return;
        }
    };

    if let Err(e) = fs::create_dir_all(&output_dir).await {
        tracing::warn!("[Debug-Log] Failed to create output dir: {}", e);
        return;
    }

    let filename = build_filename(prefix, trace_id);
    let path = output_dir.join(filename);

    match serde_json::to_vec_pretty(payload) {
        Ok(bytes) => {
            if let Err(e) = fs::write(&path, bytes).await {
                tracing::warn!("[Debug-Log] Failed to write file: {}", e);
            }
        }
        Err(e) => {
            tracing::warn!("[Debug-Log] Failed to serialize payload: {}", e);
        }
    }
}

pub fn is_enabled(cfg: &DebugLoggingConfig) -> bool {
    cfg.enabled
}
fn parse_sse_stream(raw: &str) -> (String, String) {
    let mut thinking_parts: Vec<String> = Vec::new();
    let mut content_parts: Vec<String> = Vec::new();

    for line in raw.lines() {
        let line = line.trim();
        if !line.starts_with("data: ") {
            continue;
        }
        let json_str = &line[6..];
        if json_str.is_empty() || json_str == "[DONE]" {
            continue;
        }
        if let Ok(parsed) = serde_json::from_str::<Value>(json_str) {
            if let Some(candidates) = parsed
                .get("response")
                .and_then(|r| r.get("candidates"))
                .and_then(|c| c.as_array())
            {
                for candidate in candidates {
                    if let Some(parts) = candidate
                        .get("content")
                        .and_then(|c| c.get("parts"))
                        .and_then(|p| p.as_array())
                    {
                        for part in parts {
                            let text = part.get("text").and_then(|t| t.as_str()).unwrap_or("");
                            let is_thought = part
                                .get("thought")
                                .and_then(|t| t.as_bool())
                                .unwrap_or(false);

                            if !text.is_empty() {
                                if is_thought {
                                    thinking_parts.push(text.to_string());
                                } else {
                                    content_parts.push(text.to_string());
                                }
                            }
                        }
                    }
                }
            } else if let Some(choices) = parsed.get("choices").and_then(|c| c.as_array()) {
                for choice in choices {
                    if let Some(delta) = choice.get("delta") {
                        if let Some(content) = delta.get("content").and_then(|c| c.as_str()) {
                            if !content.is_empty() {
                                content_parts.push(content.to_string());
                            }
                        }
                    }
                }
            }
        }
    }

    (thinking_parts.join(""), content_parts.join(""))
}

pub fn wrap_reqwest_stream_with_debug(
    stream: std::pin::Pin<
        Box<dyn futures::Stream<Item = Result<bytes::Bytes, reqwest::Error>> + Send>,
    >,
    cfg: DebugLoggingConfig,
    trace_id: String,
    prefix: &'static str,
    meta: Value,
) -> std::pin::Pin<Box<dyn futures::Stream<Item = Result<bytes::Bytes, reqwest::Error>> + Send>> {
    if !is_enabled(&cfg) {
        return stream;
    }

    let wrapped = async_stream::stream! {
        let mut collected: Vec<u8> = Vec::new();
        let mut inner = stream;
        while let Some(item) = inner.next().await {
            if let Ok(bytes) = &item {
                collected.extend_from_slice(bytes);
            }
            yield item;
        }

        let raw_text = String::from_utf8_lossy(&collected).to_string();
        let (thinking_content, response_content) = parse_sse_stream(&raw_text);

        let mut payload = serde_json::json!({
            "kind": "upstream_response",
            "trace_id": trace_id,
            "meta": meta,
        });
        if !thinking_content.is_empty() {
            payload["thinking_content"] = serde_json::Value::String(thinking_content);
        }
        if !response_content.is_empty() {
            payload["response_content"] = serde_json::Value::String(response_content);
        }

        let payload_trace_id = payload["trace_id"].as_str().or(Some("unknown"));
        write_debug_payload(&cfg, payload_trace_id, prefix, &payload).await;
    };

    Box::pin(wrapped)
}