//! Request builder + SSE decoder for Anthropic's `/v1/messages` streaming API.

use std::collections::HashMap;

use async_stream::try_stream;
use eventsource_stream::Eventsource;
use futures::Stream;
use harness_core::llm::{
    FinishReason, Initiator, LlmEvent, LlmEventStream, LlmRequest, ProviderError, Role, WireContent,
};
use harness_core::types::TokenUsage;
use serde_json::{json, Value};

const DEFAULT_MAX_TOKENS: u32 = 8192;
/// opencode's `applyCaching`: cache breakpoints on the first 2 system blocks + last 2
/// non-system messages.
const CACHE_BREAKPOINTS: usize = 2;

fn cache_control() -> Value {
    json!({"type": "ephemeral"})
}

fn build_system(system: &[String]) -> Vec<Value> {
    system
        .iter()
        .enumerate()
        .map(|(i, block)| {
            let mut b = json!({"type": "text", "text": block});
            if i < CACHE_BREAKPOINTS {
                b["cache_control"] = cache_control();
            }
            b
        })
        .collect()
}

fn wire_role_to_anthropic(role: Role) -> &'static str {
    match role {
        Role::User | Role::Tool => "user",
        Role::Assistant => "assistant",
        Role::System => "user", // system blocks are carried separately in `system`, not here
    }
}

fn build_messages(messages: &[harness_core::llm::WireMessage]) -> Vec<Value> {
    let mut built: Vec<Value> = messages
        .iter()
        .map(|m| {
            let content: Vec<Value> = m
                .content
                .iter()
                .map(|c| match c {
                    WireContent::Text { text } => json!({"type": "text", "text": text}),
                    WireContent::ToolCall { call_id, name, input } => {
                        json!({"type": "tool_use", "id": call_id, "name": name, "input": input})
                    }
                    WireContent::ToolResult { call_id, output, is_error } => {
                        json!({"type": "tool_result", "tool_use_id": call_id, "content": output, "is_error": is_error})
                    }
                    WireContent::Image { mime_type, data } => {
                        json!({"type": "image", "source": {"type": "base64", "media_type": mime_type, "data": data}})
                    }
                })
                .collect();
            json!({"role": wire_role_to_anthropic(m.role), "content": content})
        })
        .collect();

    let n = built.len();
    for msg in built.iter_mut().skip(n.saturating_sub(CACHE_BREAKPOINTS)) {
        if let Some(content) = msg["content"].as_array_mut() {
            if let Some(last) = content.last_mut() {
                last["cache_control"] = cache_control();
            }
        }
    }
    built
}

pub fn build_request(req: &LlmRequest) -> Value {
    let mut body = json!({
        "model": req.model,
        "system": build_system(&req.system),
        "messages": build_messages(&req.messages),
        "max_tokens": req.max_tokens.unwrap_or(DEFAULT_MAX_TOKENS),
        "stream": true,
    });

    if !req.tools.is_empty() {
        body["tools"] = Value::Array(
            req.tools
                .iter()
                .map(|t| json!({"name": t.name, "description": t.description, "input_schema": t.parameters}))
                .collect(),
        );
    }
    if let Some(temp) = req.temperature {
        body["temperature"] = json!(temp);
    }
    if let Some(reasoning) = &req.reasoning {
        if let Some(budget) = reasoning.budget_tokens {
            body["thinking"] = json!({"type": "enabled", "budget_tokens": budget});
        }
    }
    body
}

pub fn initiator_header(initiator: Initiator) -> &'static str {
    match initiator {
        Initiator::User => "user",
        Initiator::Agent => "agent",
    }
}

fn map_stop_reason(reason: Option<&str>) -> FinishReason {
    match reason {
        Some("end_turn") | Some("stop_sequence") => FinishReason::Stop,
        Some("tool_use") => FinishReason::ToolCalls,
        Some("max_tokens") => FinishReason::Length,
        Some(other) => FinishReason::Unknown(other.to_string()),
        None => FinishReason::Unknown("none".to_string()),
    }
}

#[derive(Default)]
struct BlockState {
    kind: BlockKind,
    call_id: String,
    name: String,
    partial_json: String,
    signature: Option<String>,
}

#[derive(Default, PartialEq, Eq)]
enum BlockKind {
    #[default]
    Text,
    Thinking,
    ToolUse,
}

/// Decodes a raw SSE byte stream into our normalized `LlmEvent` stream. Errors from the
/// underlying HTTP stream and any Anthropic `error` event both surface as `Err`.
pub fn decode<S, E>(byte_stream: S) -> LlmEventStream
where
    S: Stream<Item = Result<bytes::Bytes, E>> + Send + 'static,
    E: std::error::Error + Send + Sync + 'static,
{
    let events = byte_stream.eventsource();
    Box::pin(try_stream! {
        futures::pin_mut!(events);
        let mut blocks: HashMap<u64, BlockState> = HashMap::new();
        let mut usage = TokenUsage::default();

        while let Some(item) = futures::StreamExt::next(&mut events).await {
            let event = item.map_err(|e| ProviderError::Decode(e.to_string()))?;
            if event.data.is_empty() {
                continue;
            }
            let value: Value = serde_json::from_str(&event.data)
                .map_err(|e| ProviderError::Decode(format!("{e}: {}", event.data)))?;
            let kind = value["type"].as_str().unwrap_or_default();

            match kind {
                "message_start" => {
                    let u = &value["message"]["usage"];
                    usage.input = u["input_tokens"].as_u64().unwrap_or(0);
                    usage.cache_write = u["cache_creation_input_tokens"].as_u64().unwrap_or(0);
                    usage.cache_read = u["cache_read_input_tokens"].as_u64().unwrap_or(0);
                }
                "content_block_start" => {
                    let index = value["index"].as_u64().unwrap_or(0);
                    let block = &value["content_block"];
                    match block["type"].as_str().unwrap_or_default() {
                        "text" => {
                            blocks.insert(index, BlockState { kind: BlockKind::Text, ..Default::default() });
                            yield LlmEvent::TextStart { id: index.to_string() };
                        }
                        "thinking" => {
                            blocks.insert(index, BlockState { kind: BlockKind::Thinking, ..Default::default() });
                            yield LlmEvent::ReasoningStart { id: index.to_string() };
                        }
                        "tool_use" => {
                            let call_id = block["id"].as_str().unwrap_or_default().to_string();
                            let name = block["name"].as_str().unwrap_or_default().to_string();
                            blocks.insert(index, BlockState {
                                kind: BlockKind::ToolUse,
                                call_id: call_id.clone(),
                                name: name.clone(),
                                ..Default::default()
                            });
                            yield LlmEvent::ToolInputStart { call_id, name };
                        }
                        _ => {}
                    }
                }
                "content_block_delta" => {
                    let index = value["index"].as_u64().unwrap_or(0);
                    let delta = &value["delta"];
                    match delta["type"].as_str().unwrap_or_default() {
                        "text_delta" => {
                            let text = delta["text"].as_str().unwrap_or_default().to_string();
                            yield LlmEvent::TextDelta { id: index.to_string(), text };
                        }
                        "thinking_delta" => {
                            let text = delta["thinking"].as_str().unwrap_or_default().to_string();
                            yield LlmEvent::ReasoningDelta { id: index.to_string(), text };
                        }
                        "signature_delta" => {
                            if let Some(state) = blocks.get_mut(&index) {
                                state.signature = Some(delta["signature"].as_str().unwrap_or_default().to_string());
                            }
                        }
                        "input_json_delta" => {
                            let partial = delta["partial_json"].as_str().unwrap_or_default();
                            if let Some(state) = blocks.get_mut(&index) {
                                state.partial_json.push_str(partial);
                                yield LlmEvent::ToolInputDelta { call_id: state.call_id.clone(), json: partial.to_string() };
                            }
                        }
                        _ => {}
                    }
                }
                "content_block_stop" => {
                    let index = value["index"].as_u64().unwrap_or(0);
                    if let Some(state) = blocks.remove(&index) {
                        match state.kind {
                            BlockKind::Text => yield LlmEvent::TextEnd { id: index.to_string() },
                            BlockKind::Thinking => {
                                yield LlmEvent::ReasoningEnd { id: index.to_string(), signature: state.signature };
                            }
                            BlockKind::ToolUse => {
                                let input: Value = if state.partial_json.trim().is_empty() {
                                    json!({})
                                } else {
                                    serde_json::from_str(&state.partial_json).unwrap_or(Value::Null)
                                };
                                yield LlmEvent::ToolCall { call_id: state.call_id, name: state.name, input };
                            }
                        }
                    }
                }
                "message_delta" => {
                    if let Some(out) = value["usage"]["output_tokens"].as_u64() {
                        usage.output = out;
                    }
                    let stop_reason = value["delta"]["stop_reason"].as_str();
                    yield LlmEvent::Finish { reason: map_stop_reason(stop_reason), usage };
                }
                "error" => {
                    let message = value["error"]["message"].as_str().unwrap_or("unknown error").to_string();
                    let err_type = value["error"]["type"].as_str().unwrap_or("");
                    Err(match err_type {
                        "overloaded_error" => ProviderError::Overloaded,
                        "rate_limit_error" => ProviderError::RateLimited { retry_after: None },
                        "authentication_error" | "permission_error" => ProviderError::Auth(message),
                        _ => ProviderError::Http { status: 0, body: message },
                    })?;
                }
                _ => {} // ping, message_stop: nothing to emit
            }
        }
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures::StreamExt;
    use harness_core::llm::{ReasoningOpts, ToolSchema, WireMessage};

    fn sse_stream(raw: &'static str) -> LlmEventStream {
        let chunks: Vec<Result<bytes::Bytes, std::io::Error>> =
            vec![Ok(bytes::Bytes::from_static(raw.as_bytes()))];
        decode(futures::stream::iter(chunks))
    }

    #[tokio::test]
    async fn decodes_text_only_response() {
        let raw = concat!(
            "event: message_start\n",
            "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":10}}}\n\n",
            "event: content_block_start\n",
            "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
            "event: content_block_delta\n",
            "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\n",
            "event: content_block_stop\n",
            "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
            "event: message_delta\n",
            "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":5}}\n\n",
            "event: message_stop\n",
            "data: {\"type\":\"message_stop\"}\n\n",
        );
        let events: Vec<LlmEvent> = sse_stream(raw).map(|e| e.unwrap()).collect().await;
        assert_eq!(
            events,
            vec![
                LlmEvent::TextStart { id: "0".into() },
                LlmEvent::TextDelta {
                    id: "0".into(),
                    text: "Hello".into()
                },
                LlmEvent::TextEnd { id: "0".into() },
                LlmEvent::Finish {
                    reason: FinishReason::Stop,
                    usage: TokenUsage {
                        input: 10,
                        output: 5,
                        ..Default::default()
                    },
                },
            ]
        );
    }

    #[tokio::test]
    async fn decodes_tool_call_with_streamed_json_input() {
        let raw = concat!(
            "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n",
            "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"call_1\",\"name\":\"read\"}}\n\n",
            "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"file\\\"\"}}\n\n",
            "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\":\\\"a.txt\\\"}\"}}\n\n",
            "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
            "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\"},\"usage\":{\"output_tokens\":8}}\n\n",
        );
        let events: Vec<LlmEvent> = sse_stream(raw).map(|e| e.unwrap()).collect().await;
        assert_eq!(
            events,
            vec![
                LlmEvent::ToolInputStart {
                    call_id: "call_1".into(),
                    name: "read".into()
                },
                LlmEvent::ToolInputDelta {
                    call_id: "call_1".into(),
                    json: "{\"file\"".into()
                },
                LlmEvent::ToolInputDelta {
                    call_id: "call_1".into(),
                    json: ":\"a.txt\"}".into()
                },
                LlmEvent::ToolCall {
                    call_id: "call_1".into(),
                    name: "read".into(),
                    input: json!({"file": "a.txt"}),
                },
                LlmEvent::Finish {
                    reason: FinishReason::ToolCalls,
                    usage: TokenUsage {
                        input: 1,
                        output: 8,
                        ..Default::default()
                    },
                },
            ]
        );
    }

    #[tokio::test]
    async fn decodes_thinking_block_with_signature() {
        let raw = concat!(
            "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n",
            "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\"}}\n\n",
            "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"pondering\"}}\n\n",
            "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"signature_delta\",\"signature\":\"sig123\"}}\n\n",
            "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
            "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":2}}\n\n",
        );
        let events: Vec<LlmEvent> = sse_stream(raw).map(|e| e.unwrap()).collect().await;
        assert_eq!(
            events,
            vec![
                LlmEvent::ReasoningStart { id: "0".into() },
                LlmEvent::ReasoningDelta {
                    id: "0".into(),
                    text: "pondering".into()
                },
                LlmEvent::ReasoningEnd {
                    id: "0".into(),
                    signature: Some("sig123".into())
                },
                LlmEvent::Finish {
                    reason: FinishReason::Stop,
                    usage: TokenUsage {
                        input: 1,
                        output: 2,
                        ..Default::default()
                    },
                },
            ]
        );
    }

    #[tokio::test]
    async fn mid_stream_error_event_surfaces_as_err() {
        let raw = concat!(
            "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":1}}}\n\n",
            "data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"overloaded\"}}\n\n",
        );
        let events: Vec<Result<LlmEvent, ProviderError>> = sse_stream(raw).collect().await;
        assert!(matches!(
            events.last(),
            Some(Err(ProviderError::Overloaded))
        ));
    }

    #[test]
    fn build_request_applies_cache_control_to_first_two_system_blocks() {
        let req = LlmRequest {
            model: "claude-sonnet".into(),
            system: vec!["env".into(), "agent".into(), "instructions".into()],
            messages: vec![],
            tools: vec![],
            temperature: None,
            max_tokens: None,
            reasoning: None,
            initiator: Initiator::User,
        };
        let body = build_request(&req);
        let system = body["system"].as_array().unwrap();
        assert_eq!(system.len(), 3);
        assert!(system[0]["cache_control"].is_object());
        assert!(system[1]["cache_control"].is_object());
        assert!(system[2].get("cache_control").is_none());
    }

    #[test]
    fn build_request_applies_cache_control_to_last_two_messages() {
        let messages = vec![
            WireMessage {
                role: Role::User,
                content: vec![WireContent::Text { text: "1".into() }],
            },
            WireMessage {
                role: Role::Assistant,
                content: vec![WireContent::Text { text: "2".into() }],
            },
            WireMessage {
                role: Role::User,
                content: vec![WireContent::Text { text: "3".into() }],
            },
        ];
        let req = LlmRequest {
            model: "claude-sonnet".into(),
            system: vec![],
            messages,
            tools: vec![],
            temperature: None,
            max_tokens: None,
            reasoning: None,
            initiator: Initiator::User,
        };
        let body = build_request(&req);
        let msgs = body["messages"].as_array().unwrap();
        assert!(msgs[0]["content"][0].get("cache_control").is_none());
        assert!(msgs[1]["content"][0]["cache_control"].is_object());
        assert!(msgs[2]["content"][0]["cache_control"].is_object());
    }

    #[test]
    fn build_request_maps_tool_schema_to_input_schema_key() {
        let req = LlmRequest {
            model: "claude-sonnet".into(),
            system: vec![],
            messages: vec![],
            tools: vec![ToolSchema {
                name: "read".into(),
                description: "reads a file".into(),
                parameters: json!({"type": "object"}),
            }],
            temperature: None,
            max_tokens: None,
            reasoning: None,
            initiator: Initiator::User,
        };
        let body = build_request(&req);
        assert_eq!(body["tools"][0]["name"], "read");
        assert_eq!(body["tools"][0]["input_schema"], json!({"type": "object"}));
    }

    #[test]
    fn build_request_includes_thinking_budget_when_reasoning_set() {
        let req = LlmRequest {
            model: "claude-sonnet".into(),
            system: vec![],
            messages: vec![],
            tools: vec![],
            temperature: None,
            max_tokens: None,
            reasoning: Some(ReasoningOpts {
                effort: None,
                budget_tokens: Some(2048),
            }),
            initiator: Initiator::User,
        };
        let body = build_request(&req);
        assert_eq!(body["thinking"]["budget_tokens"], 2048);
    }
}