chimera-opencode 0.2.1

use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, OnceLock};
use std::time::Duration;

use chimera_core::*;

use chimera_core::ToolChoice;

use crate::config::{OpenCodeConfig, OpenCodeProvider};
use crate::translate;
use crate::wire::*;

const DEFAULT_TIMEOUT_SECS: u64 = 300;
const DEFAULT_MAX_RETRIES: u32 = 2;
const DEFAULT_MODEL: &str = "glm-5";
const BASE_RETRY_DELAY: Duration = Duration::from_secs(5);
const MAX_RETRY_DELAY: Duration = Duration::from_secs(120);

pub struct HttpSession {
    api_key: Option<String>,
    provider: OpenCodeProvider,
    session_id: Arc<OnceLock<String>>,
    active_interrupt: Arc<tokio::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>>,
    config: SessionConfig<OpenCodeConfig>,
    history: Arc<tokio::sync::Mutex<Vec<WireMessage>>>,
    client: Option<reqwest::Client>,
}

impl HttpSession {
    pub(crate) fn new(
        api_key: Option<String>,
        provider: OpenCodeProvider,
        config: SessionConfig<OpenCodeConfig>,
    ) -> Self {
        Self {
            api_key,
            provider,
            session_id: Arc::new(OnceLock::new()),
            active_interrupt: Arc::new(tokio::sync::Mutex::new(None)),
            config,
            history: Arc::new(tokio::sync::Mutex::new(Vec::new())),
            client: None,
        }
    }

    pub(crate) fn with_session_id(self, id: String) -> Self {
        let _ = self.session_id.set(id);
        self
    }

    #[cfg(test)]
    pub(crate) fn api_key(&self) -> Option<&str> {
        self.api_key.as_deref()
    }

    fn model(&self) -> &str {
        let m = self.config.model.as_deref().unwrap_or(DEFAULT_MODEL);
        // Strip "provider-id/" prefix (e.g. "opencode-go/glm-5" -> "glm-5")
        m.split_once('/').map(|(_, model)| model).unwrap_or(m)
    }

    fn timeout(&self) -> Duration {
        Duration::from_secs(
            self.config
                .backend
                .timeout_secs
                .unwrap_or(DEFAULT_TIMEOUT_SECS),
        )
    }

    fn max_retries(&self) -> u32 {
        self.config
            .backend
            .max_retries
            .unwrap_or(DEFAULT_MAX_RETRIES)
    }

    fn base_url(&self) -> &str {
        self.provider.base_url()
    }

    fn get_or_build_client(&mut self) -> Result<reqwest::Client> {
        if let Some(ref client) = self.client {
            return Ok(client.clone());
        }
        let client = reqwest::Client::builder()
            .timeout(self.timeout())
            .build()
            .map_err(|e| AgentError::Other {
                message: "failed to build HTTP client".into(),
                source: Some(Box::new(e)),
            })?;
        self.client = Some(client.clone());
        Ok(client)
    }

    fn build_request(
        &self,
        input: &Input,
        options: &TurnOptions,
        history: &[WireMessage],
    ) -> Result<ChatCompletionRequest> {
        let mut messages = Vec::new();

        if let Some(sp) = &self.config.system_prompt {
            messages.push(WireMessage::system(sp.clone()));
        }

        messages.extend_from_slice(history);

        let prompt_text = prompt_text_from_input(input)?;
        messages.push(WireMessage::user(prompt_text));

        let response_format = options.output_schema.as_ref().map(|schema| ResponseFormat {
            format_type: "json_schema".into(),
            json_schema: Some(serde_json::json!({
                "name": "response",
                "strict": true,
                "schema": schema,
            })),
        });

        let stream = if self.config.backend.stream {
            Some(true)
        } else {
            None
        };

        let tools: Vec<WireTool> = options
            .tools
            .iter()
            .map(|t| WireTool {
                kind: "function".into(),
                function: WireToolFunction {
                    name: t.name.clone(),
                    description: t.description.clone(),
                    parameters: t.parameters.clone(),
                },
            })
            .collect();

        let tool_choice = options.tool_choice.as_ref().map(|tc| match tc {
            ToolChoice::Auto => serde_json::json!("auto"),
            ToolChoice::Required => serde_json::json!("required"),
            ToolChoice::None => serde_json::json!("none"),
            ToolChoice::Function(name) => {
                serde_json::json!({"type": "function", "function": {"name": name}})
            }
        });

        Ok(ChatCompletionRequest {
            model: self.model().to_string(),
            messages,
            max_tokens: self.config.backend.max_tokens,
            temperature: self.config.backend.temperature,
            stream,
            response_format,
            tools,
            tool_choice,
        })
    }

    /// Make a single-turn completion where the model is forced to call `tool`.
    ///
    /// Returns the raw JSON arguments string from the first tool call.
    pub async fn complete_with_tool(
        &mut self,
        system_prompt: &str,
        user_message: &str,
        tool: chimera_core::ToolDefinition,
    ) -> chimera_core::Result<String> {
        use chimera_core::{Input, Session, ToolChoice, TurnOptions};
        // Override system prompt for this call.
        let original_system = self.config.system_prompt.clone();
        if !system_prompt.is_empty() {
            self.config.system_prompt = Some(system_prompt.to_string());
        }
        // Force non-streaming for structured extraction.
        let original_stream = self.config.backend.stream;
        self.config.backend.stream = false;

        let options = TurnOptions {
            tools: vec![tool],
            tool_choice: Some(ToolChoice::Required),
            ..TurnOptions::default()
        };
        let output = self
            .turn(Input::Text(user_message.to_string()), options)
            .await;

        // Restore config.
        self.config.system_prompt = original_system;
        self.config.backend.stream = original_stream;

        let output = output?;
        output
            .tool_calls
            .into_iter()
            .next()
            .map(|tc| tc.arguments)
            .ok_or_else(|| chimera_core::AgentError::TurnFailed {
                message: "model did not produce a tool call".into(),
            })
    }
}

fn prompt_text_from_input(input: &Input) -> Result<String> {
    match input {
        Input::Text(s) => Ok(s.clone()),
        Input::Structured(parts) => {
            let parts = parts
                .iter()
                .map(|p| match p {
                    InputPart::Text(t) => Ok(t.as_str()),
                    InputPart::Image(path) => Err(AgentError::Other {
                        message: format!(
                            "opencode HTTP mode does not support image input yet: {}",
                            path.display()
                        ),
                        source: None,
                    }),
                    _ => Err(AgentError::Other {
                        message: "opencode HTTP mode received an unsupported structured input part"
                            .into(),
                        source: None,
                    }),
                })
                .collect::<Result<Vec<_>>>()?;
            Ok(parts.join("\n\n"))
        }
        _ => Err(AgentError::Other {
            message: "opencode HTTP mode received an unsupported input shape".into(),
            source: None,
        }),
    }
}

/// Public session type -- dispatches between HTTP (direct API) and Agent (subprocess) modes.
pub enum OpenCodeSession {
    /// HTTP mode: calls the OpenCode Zen API (or compatible endpoint) directly.
    Http(Box<HttpSession>),
    /// Agent mode: spawns `opencode serve` and communicates via the local HTTP API.
    Agent(Box<crate::agent_session::OpenCodeAgentSession>),
}

impl Session for OpenCodeSession {
    fn turn_stream(
        &mut self,
        input: Input,
        options: TurnOptions,
    ) -> Pin<Box<dyn Future<Output = Result<EventStream>> + Send + '_>> {
        match self {
            OpenCodeSession::Http(s) => s.turn_stream(input, options),
            OpenCodeSession::Agent(s) => s.turn_stream(input, options),
        }
    }

    fn session_id(&self) -> Option<&str> {
        match self {
            OpenCodeSession::Http(s) => s.session_id(),
            OpenCodeSession::Agent(s) => s.session_id(),
        }
    }

    fn interrupt(&mut self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>> {
        match self {
            OpenCodeSession::Http(s) => s.interrupt(),
            OpenCodeSession::Agent(s) => s.interrupt(),
        }
    }
}

impl Session for HttpSession {
    fn turn_stream(
        &mut self,
        input: Input,
        options: TurnOptions,
    ) -> Pin<Box<dyn Future<Output = Result<EventStream>> + Send + '_>> {
        Box::pin(async move {
            if self.active_interrupt.lock().await.is_some() {
                return Err(AgentError::Other {
                    message: "turn already in progress".into(),
                    source: None,
                });
            }

            let client = self.get_or_build_client()?;
            let history_snapshot = self.history.lock().await.clone();
            let request = self.build_request(&input, &options, &history_snapshot)?;
            let url = format!("{}/chat/completions", self.base_url());
            let api_key = self.api_key.clone();
            let max_retries = self.max_retries();
            let is_streaming = request.stream == Some(true);

            // Append user message to history for multi-turn.
            let prompt_text = prompt_text_from_input(&input)?;
            self.history
                .lock()
                .await
                .push(WireMessage::user(prompt_text));

            let (tx, rx) =
                tokio::sync::mpsc::channel::<std::result::Result<AgentEvent, AgentError>>(128);
            let session_id = Arc::clone(&self.session_id);
            let (interrupt_tx, interrupt_rx) = tokio::sync::oneshot::channel();
            {
                let mut slot = self.active_interrupt.lock().await;
                *slot = Some(interrupt_tx);
            }
            let active_interrupt = Arc::clone(&self.active_interrupt);
            let history = Arc::clone(&self.history);
            let timeout_duration = options.timeout;

            tokio::spawn(async move {
                execute_turn(
                    client,
                    url,
                    api_key,
                    request,
                    max_retries,
                    is_streaming,
                    tx,
                    interrupt_rx,
                    session_id,
                    active_interrupt,
                    history,
                    timeout_duration,
                )
                .await;
            });

            Ok(EventStream::from_receiver(rx))
        })
    }

    fn session_id(&self) -> Option<&str> {
        self.session_id.get().map(|s| s.as_str())
    }

    fn interrupt(&mut self) -> Pin<Box<dyn Future<Output = Result<()>> + Send + '_>> {
        let active_interrupt = Arc::clone(&self.active_interrupt);
        Box::pin(async move {
            if let Some(tx) = active_interrupt.lock().await.take() {
                let _ = tx.send(());
            }
            Ok(())
        })
    }
}

#[allow(clippy::too_many_arguments)]
async fn execute_turn(
    client: reqwest::Client,
    url: String,
    api_key: Option<String>,
    request: ChatCompletionRequest,
    max_retries: u32,
    is_streaming: bool,
    tx: tokio::sync::mpsc::Sender<std::result::Result<AgentEvent, AgentError>>,
    mut interrupt_rx: tokio::sync::oneshot::Receiver<()>,
    session_id: Arc<OnceLock<String>>,
    active_interrupt: Arc<tokio::sync::Mutex<Option<tokio::sync::oneshot::Sender<()>>>>,
    history: Arc<tokio::sync::Mutex<Vec<WireMessage>>>,
    timeout_duration: Option<Duration>,
) {
    let timeout_sleep = tokio::time::sleep(timeout_duration.unwrap_or(Duration::MAX));
    tokio::pin!(timeout_sleep);
    let has_timeout = timeout_duration.is_some();

    let _ = tx.send(Ok(AgentEvent::TurnStarted)).await;

    // Retry loop for the HTTP request.
    let mut attempt = 0u32;
    let response = loop {
        let mut req = client.post(&url).json(&request);
        if let Some(ref key) = api_key {
            req = req.bearer_auth(key);
        }

        tokio::select! {
            _ = &mut interrupt_rx => {
                let _ = tx.send(Err(AgentError::Interrupted)).await;
                active_interrupt.lock().await.take();
                return;
            }
            _ = &mut timeout_sleep, if has_timeout => {
                let _ = tx.send(Err(AgentError::Timeout {
                    duration: timeout_duration.unwrap(),
                })).await;
                active_interrupt.lock().await.take();
                return;
            }
            result = req.send() => {
                match result {
                    Ok(resp) if resp.status().is_success() => break resp,
                    Ok(resp) => {
                        let status = resp.status();
                        let retry_after = resp.headers()
                            .get("retry-after")
                            .and_then(|v| v.to_str().ok())
                            .and_then(|s| s.parse::<u64>().ok())
                            .unwrap_or(60);
                        let body = resp.text().await.unwrap_or_default();

                        if attempt >= max_retries {
                            let _ = tx.send(Ok(AgentEvent::TurnFailed {
                                message: format!("HTTP {status}: {body}"),
                            })).await;
                            active_interrupt.lock().await.take();
                            return;
                        }

                        let delay = if status.as_u16() == 429 {
                            Duration::from_secs(retry_after.min(MAX_RETRY_DELAY.as_secs()))
                        } else if status.is_server_error() {
                            (BASE_RETRY_DELAY * (1 << attempt)).min(MAX_RETRY_DELAY)
                        } else {
                            // Client error (4xx other than 429) -- don't retry.
                            let _ = tx.send(Ok(AgentEvent::TurnFailed {
                                message: format!("HTTP {status}: {body}"),
                            })).await;
                            active_interrupt.lock().await.take();
                            return;
                        };

                        attempt += 1;
                        tokio::time::sleep(delay).await;
                        continue;
                    }
                    Err(e) => {
                        if attempt >= max_retries {
                            let _ = tx.send(Err(AgentError::Other {
                                message: format!("HTTP request failed: {e}"),
                                source: Some(Box::new(e)),
                            })).await;
                            active_interrupt.lock().await.take();
                            return;
                        }
                        attempt += 1;
                        let delay = (BASE_RETRY_DELAY * (1 << attempt)).min(MAX_RETRY_DELAY);
                        tokio::time::sleep(delay).await;
                        continue;
                    }
                }
            }
        }
    };

    if is_streaming {
        handle_sse_stream(
            response,
            &tx,
            interrupt_rx,
            &session_id,
            &history,
            timeout_duration,
        )
        .await;
    } else {
        handle_non_streaming(response, &tx, &session_id, &history).await;
    }

    active_interrupt.lock().await.take();
}

async fn handle_non_streaming(
    response: reqwest::Response,
    tx: &tokio::sync::mpsc::Sender<std::result::Result<AgentEvent, AgentError>>,
    session_id: &Arc<OnceLock<String>>,
    history: &Arc<tokio::sync::Mutex<Vec<WireMessage>>>,
) {
    match response.json::<ChatCompletionResponse>().await {
        Ok(resp) => {
            // Append assistant response to history for multi-turn.
            if let Some(choice) = resp.choices.first() {
                let content = choice.message.content.as_deref().unwrap_or("");
                history.lock().await.push(WireMessage::assistant(content));
            }

            let events = translate::translate_response(resp);
            for event in events {
                if let AgentEvent::TurnCompleted {
                    result: Some(ref r),
                    ..
                } = event
                    && let Some(ref sid) = r.session_id
                {
                    let _ = session_id.set(sid.clone());
                }
                if tx.send(Ok(event)).await.is_err() {
                    return;
                }
            }
        }
        Err(e) => {
            let _ = tx
                .send(Err(AgentError::Other {
                    message: format!("failed to parse response: {e}"),
                    source: Some(Box::new(e)),
                }))
                .await;
        }
    }
}

#[allow(clippy::too_many_arguments)]
/// Accumulated state for a tool call being built across SSE chunks.
#[derive(Default)]
struct PendingToolCall {
    id: String,
    name: String,
    arguments: String,
}

async fn handle_sse_stream(
    response: reqwest::Response,
    tx: &tokio::sync::mpsc::Sender<std::result::Result<AgentEvent, AgentError>>,
    mut interrupt_rx: tokio::sync::oneshot::Receiver<()>,
    session_id: &Arc<OnceLock<String>>,
    history: &Arc<tokio::sync::Mutex<Vec<WireMessage>>>,
    timeout_duration: Option<Duration>,
) {
    let timeout_sleep = tokio::time::sleep(timeout_duration.unwrap_or(Duration::MAX));
    tokio::pin!(timeout_sleep);
    let has_timeout = timeout_duration.is_some();
    let mut buf = String::new();
    let mut full_content = String::new();
    // Tool call accumulator: index -> PendingToolCall
    let mut pending_tool_calls: std::collections::BTreeMap<usize, PendingToolCall> =
        std::collections::BTreeMap::new();
    let mut response = response;

    loop {
        tokio::select! {
            _ = &mut interrupt_rx => {
                let _ = tx.send(Err(AgentError::Interrupted)).await;
                return;
            }
            _ = &mut timeout_sleep, if has_timeout => {
                let _ = tx.send(Err(AgentError::Timeout {
                    duration: timeout_duration.unwrap(),
                })).await;
                return;
            }
            chunk_result = response.chunk() => {
                match chunk_result {
                    Ok(Some(bytes)) => {
                        buf.push_str(&String::from_utf8_lossy(&bytes));

                        while let Some(pos) = buf.find('\n') {
                            let line = buf[..pos].to_string();
                            buf = buf[pos + 1..].to_string();

                            let line = line.trim();
                            if line.is_empty() {
                                continue;
                            }

                            if let Some(data) = line.strip_prefix("data: ") {
                                if data == "[DONE]" {
                                    // Append accumulated content to history.
                                    if !full_content.is_empty() {
                                        history
                                            .lock()
                                            .await
                                            .push(WireMessage::assistant(&full_content));
                                    }
                                    // Emit accumulated tool calls.
                                    for (_, tc) in pending_tool_calls {
                                        if tx
                                            .send(Ok(AgentEvent::ToolCall {
                                                id: tc.id,
                                                name: tc.name,
                                                arguments: tc.arguments,
                                            }))
                                            .await
                                            .is_err()
                                        {
                                            return;
                                        }
                                    }
                                    return;
                                }

                                if let Ok(chunk) =
                                    serde_json::from_str::<ChatCompletionChunk>(data)
                                {
                                    // Accumulate content for history.
                                    for choice in &chunk.choices {
                                        if let Some(ref c) = choice.delta.content {
                                            full_content.push_str(c);
                                        }
                                        // Accumulate tool call deltas.
                                        for tc_delta in &choice.delta.tool_calls {
                                            let entry = pending_tool_calls
                                                .entry(tc_delta.index)
                                                .or_default();
                                            if let Some(ref id) = tc_delta.id {
                                                entry.id.clone_from(id);
                                            }
                                            if let Some(ref f) = tc_delta.function {
                                                if let Some(ref name) = f.name {
                                                    entry.name.clone_from(name);
                                                }
                                                if let Some(ref args) = f.arguments {
                                                    entry.arguments.push_str(args);
                                                }
                                            }
                                        }
                                    }

                                    let events = translate::translate_chunk(chunk);
                                    for event in events {
                                        if let AgentEvent::TurnCompleted {
                                            result: Some(ref r),
                                            ..
                                        } = event
                                            && let Some(ref sid) = r.session_id
                                        {
                                            let _ = session_id.set(sid.clone());
                                        }
                                        if tx.send(Ok(event)).await.is_err() {
                                            return;
                                        }
                                    }
                                }
                            }
                        }
                    }
                    Ok(None) => {
                        // Stream ended without [DONE]. Append any accumulated content.
                        if !full_content.is_empty() {
                            history.lock().await.push(
                                WireMessage::assistant(&full_content),
                            );
                        }
                        return;
                    }
                    Err(e) => {
                        let _ = tx.send(Err(AgentError::Other {
                            message: format!("SSE stream error: {e}"),
                            source: Some(Box::new(e)),
                        })).await;
                        return;
                    }
                }
            }
        }
    }
}

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

    fn make_session(config: SessionConfig<OpenCodeConfig>) -> HttpSession {
        HttpSession::new(Some("test-key".into()), OpenCodeProvider::Zen, config)
    }

    fn default_config() -> SessionConfig<OpenCodeConfig> {
        SessionConfig::builder()
            .backend(OpenCodeConfig::default())
            .build()
    }

    #[test]
    fn session_id_none_initially() {
        let session = make_session(default_config());
        assert!(session.session_id().is_none());
    }

    #[test]
    fn session_id_after_set() {
        let session = make_session(default_config()).with_session_id("sess-1".into());
        assert_eq!(session.session_id(), Some("sess-1"));
    }

    #[test]
    fn model_defaults_to_glm5() {
        let session = make_session(default_config());
        assert_eq!(session.model(), "glm-5");
    }

    #[test]
    fn model_from_config() {
        let config = SessionConfig::builder()
            .model("claude-sonnet-4-20250514")
            .backend(OpenCodeConfig::default())
            .build();
        let session = make_session(config);
        assert_eq!(session.model(), "claude-sonnet-4-20250514");
    }

    #[test]
    fn model_strips_provider_prefix() {
        let config = SessionConfig::builder()
            .model("opencode-go/glm-5")
            .backend(OpenCodeConfig::default())
            .build();
        let session = make_session(config);
        assert_eq!(session.model(), "glm-5");
    }

    #[tokio::test]
    async fn build_request_basic() {
        let session = make_session(default_config());
        let history = session.history.lock().await.clone();
        let req = session
            .build_request(
                &Input::Text("hello".into()),
                &TurnOptions::default(),
                &history,
            )
            .unwrap();

        assert_eq!(req.model, "glm-5");
        assert_eq!(req.messages.len(), 1);
        assert_eq!(req.messages[0].role, "user");
        assert_eq!(req.messages[0].content, "hello");
        assert_eq!(req.stream, Some(true));
        assert!(req.response_format.is_none());
    }

    #[tokio::test]
    async fn build_request_with_system_prompt_and_schema() {
        let config = SessionConfig::builder()
            .system_prompt("Be helpful")
            .backend(OpenCodeConfig::builder().stream(false).build())
            .build();
        let session = make_session(config);
        let options = TurnOptions {
            output_schema: Some(serde_json::json!({"type": "object"})),
            ..Default::default()
        };
        let history = session.history.lock().await.clone();
        let req = session
            .build_request(&Input::Text("test".into()), &options, &history)
            .unwrap();

        assert_eq!(req.messages[0].role, "system");
        assert_eq!(req.messages[0].content, "Be helpful");
        assert_eq!(req.messages[1].role, "user");
        assert!(req.stream.is_none());
        assert!(req.response_format.is_some());
        assert_eq!(
            req.response_format.as_ref().unwrap().format_type,
            "json_schema"
        );
    }

    #[tokio::test]
    async fn build_request_includes_history() {
        let config = default_config();
        let session = make_session(config);
        {
            let mut h = session.history.lock().await;
            h.push(WireMessage::user("first"));
            h.push(WireMessage::assistant("reply"));
        }
        let history = session.history.lock().await.clone();
        let req = session
            .build_request(
                &Input::Text("second".into()),
                &TurnOptions::default(),
                &history,
            )
            .unwrap();

        assert_eq!(req.messages.len(), 3);
        assert_eq!(req.messages[0].content, "first");
        assert_eq!(req.messages[1].content, "reply");
        assert_eq!(req.messages[2].content, "second");
    }

    #[tokio::test]
    async fn turn_stream_rejects_when_turn_active() {
        let mut session = make_session(default_config());
        let (tx, _rx) = tokio::sync::oneshot::channel();
        *session.active_interrupt.lock().await = Some(tx);

        let err = session
            .turn_stream(Input::Text("hi".into()), TurnOptions::default())
            .await;
        match err {
            Err(AgentError::Other { message, .. }) => {
                assert_eq!(message, "turn already in progress")
            }
            _ => panic!("expected Other error"),
        }
    }

    #[test]
    fn prompt_text_rejects_images() {
        let err = prompt_text_from_input(&Input::Structured(vec![InputPart::Image(
            std::path::PathBuf::from("/tmp/example.png"),
        )]))
        .unwrap_err();

        assert!(
            err.to_string()
                .contains("opencode HTTP mode does not support image input yet")
        );
    }
}