provider-core 0.1.1

// SPDX-License-Identifier: AGPL-3.0-or-later

//! Transport-independent OpenRouter Chat Completions request and SSE contracts.

use crate::ollama_chat::{
    BoundedRetryAfter, HttpMethod, NetworkPolicy, ReplayBoundary, ReplayDecision,
};
use router_protocol::{
    AttemptId, CanonicalRequest, ContentBlock, InputItem, MessageRole, PhaseKind, RouteIdentity,
    RouteIdentityError, ToolChoice, TransportKind,
};
use serde_json::{Map, Value, json};
use std::collections::{BTreeMap, HashSet};
use std::fmt;

pub const DEFAULT_MAX_SSE_EVENT_BYTES: usize = 256 * 1024;

#[derive(Clone, PartialEq, Eq)]
pub struct OpenRouterSecretReferenceMarker(String);

impl OpenRouterSecretReferenceMarker {
    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl fmt::Debug for OpenRouterSecretReferenceMarker {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("OpenRouterSecretReferenceMarker([REDACTED])")
    }
}

#[derive(Clone, PartialEq)]
pub struct OpenRouterRequestSpec {
    pub method: HttpMethod,
    pub origin: String,
    pub path: String,
    pub content_type: &'static str,
    pub accept: &'static str,
    pub bearer_secret: OpenRouterSecretReferenceMarker,
    pub network: NetworkPolicy,
    body: Value,
}

impl OpenRouterRequestSpec {
    pub fn body(&self) -> &Value {
        &self.body
    }

    pub fn body_bytes(&self) -> Result<Vec<u8>, serde_json::Error> {
        serde_json::to_vec(&self.body)
    }
}

impl fmt::Debug for OpenRouterRequestSpec {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("OpenRouterRequestSpec")
            .field("method", &self.method)
            .field("origin", &self.origin)
            .field("path", &self.path)
            .field("content_type", &self.content_type)
            .field("accept", &self.accept)
            .field("bearer_secret", &self.bearer_secret)
            .field("network", &self.network)
            .field("body", &"[REDACTED]")
            .finish()
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OpenRouterRequestErrorKind {
    InvalidRouteIdentity,
    UnsupportedRoute,
    InvalidCanonicalRequest,
    AmbiguousToolName,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct OpenRouterRequestError {
    pub kind: OpenRouterRequestErrorKind,
}

impl fmt::Display for OpenRouterRequestError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("the OpenRouter chat request could not be prepared")
    }
}

impl std::error::Error for OpenRouterRequestError {}

impl From<RouteIdentityError> for OpenRouterRequestError {
    fn from(_: RouteIdentityError) -> Self {
        Self {
            kind: OpenRouterRequestErrorKind::InvalidRouteIdentity,
        }
    }
}

pub fn build_request(
    route: &RouteIdentity,
    request: &CanonicalRequest,
) -> Result<OpenRouterRequestSpec, OpenRouterRequestError> {
    route.validate()?;
    if route.phase != PhaseKind::OpenrouterApi || route.transport != TransportKind::DirectHttps {
        return Err(OpenRouterRequestError {
            kind: OpenRouterRequestErrorKind::UnsupportedRoute,
        });
    }
    request.validate().map_err(|_| OpenRouterRequestError {
        kind: OpenRouterRequestErrorKind::InvalidCanonicalRequest,
    })?;
    let secret_ref = route.credential_ref.clone().ok_or(OpenRouterRequestError {
        kind: OpenRouterRequestErrorKind::InvalidRouteIdentity,
    })?;

    Ok(OpenRouterRequestSpec {
        method: HttpMethod::Post,
        origin: route.endpoint.origin.clone(),
        path: route.endpoint.path.clone(),
        content_type: "application/json",
        accept: "text/event-stream",
        bearer_secret: OpenRouterSecretReferenceMarker(secret_ref),
        network: NetworkPolicy {
            allow_proxy: false,
            allow_dns: true,
            follow_redirects: false,
            require_tls: true,
        },
        body: build_body(route, request)?,
    })
}

fn build_body(
    route: &RouteIdentity,
    request: &CanonicalRequest,
) -> Result<Value, OpenRouterRequestError> {
    let mut object = Map::new();
    object.insert("model".to_owned(), Value::String(route.model_id.clone()));
    object.insert("messages".to_owned(), Value::Array(build_messages(request)));
    object.insert("stream".to_owned(), Value::Bool(true));
    object.insert("stream_options".to_owned(), json!({"include_usage": true}));
    object.insert(
        "provider".to_owned(),
        json!({
            "require_parameters": true,
            "allow_fallbacks": false,
            "data_collection": "deny",
            "zdr": true,
        }),
    );
    object.insert("tool_choice".to_owned(), openrouter_tool_choice(request));

    if !request.tools.is_empty() && request.tool_choice != ToolChoice::None {
        let mut names = HashSet::new();
        let mut tools = Vec::with_capacity(request.tools.len());
        for tool in &request.tools {
            let name = provider_tool_name(tool.namespace.as_deref(), &tool.name);
            if !names.insert(name.clone()) {
                return Err(OpenRouterRequestError {
                    kind: OpenRouterRequestErrorKind::AmbiguousToolName,
                });
            }
            tools.push(json!({
                "type": "function",
                "function": {
                    "name": name,
                    "description": tool.description,
                    "parameters": tool.parameters,
                    "strict": tool.strict,
                }
            }));
        }
        object.insert("tools".to_owned(), Value::Array(tools));
    }
    if let Some(maximum) = request.max_output_tokens {
        object.insert("max_tokens".to_owned(), Value::from(maximum));
    }
    if let Some(format) = &request.structured_output {
        object.insert(
            "response_format".to_owned(),
            json!({
                "type": "json_schema",
                "json_schema": {
                    "name": format.name,
                    "strict": format.strict,
                    "schema": format.schema,
                }
            }),
        );
    }
    Ok(Value::Object(object))
}

fn build_messages(request: &CanonicalRequest) -> Vec<Value> {
    request
        .input
        .iter()
        .map(|item| match item {
            InputItem::Message { role, content } => json!({
                "role": role_name(*role),
                "content": openrouter_content(content),
            }),
            InputItem::FunctionCall {
                call_id,
                namespace,
                name,
                arguments,
            } => json!({
                "role": "assistant",
                "content": null,
                "tool_calls": [{
                    "id": call_id.as_str(),
                    "type": "function",
                    "function": {
                        "name": provider_tool_name(namespace.as_deref(), name),
                        "arguments": arguments.to_string(),
                    }
                }]
            }),
            InputItem::FunctionCallOutput { call_id, output } => json!({
                "role": "tool",
                "tool_call_id": call_id.as_str(),
                "content": output.to_string(),
            }),
        })
        .collect()
}

fn openrouter_content(content: &[ContentBlock]) -> Value {
    if content
        .iter()
        .all(|block| matches!(block, ContentBlock::Text { .. }))
    {
        let text = content
            .iter()
            .filter_map(|block| match block {
                ContentBlock::Text { text } => Some(text.as_str()),
                ContentBlock::ImageReference { .. } => None,
            })
            .collect::<String>();
        return Value::String(text);
    }
    Value::Array(
        content
            .iter()
            .map(|block| match block {
                ContentBlock::Text { text } => json!({"type": "text", "text": text}),
                ContentBlock::ImageReference { url } => {
                    json!({"type": "image_url", "image_url": {"url": url}})
                }
            })
            .collect(),
    )
}

const fn role_name(role: MessageRole) -> &'static str {
    match role {
        MessageRole::Developer => "developer",
        MessageRole::System => "system",
        MessageRole::User => "user",
        MessageRole::Assistant => "assistant",
    }
}

fn provider_tool_name(namespace: Option<&str>, name: &str) -> String {
    namespace.map_or_else(
        || name.to_owned(),
        |namespace| format!("{namespace}__{name}"),
    )
}

fn openrouter_tool_choice(request: &CanonicalRequest) -> Value {
    match &request.tool_choice {
        ToolChoice::Auto => Value::String("auto".to_owned()),
        ToolChoice::None => Value::String("none".to_owned()),
        ToolChoice::Required => Value::String("required".to_owned()),
        ToolChoice::Named(name) => {
            let provider_name = request
                .tools
                .iter()
                .find(|tool| tool.name == *name)
                .map_or_else(
                    || name.clone(),
                    |tool| provider_tool_name(tool.namespace.as_deref(), &tool.name),
                );
            json!({"type":"function","function":{"name":provider_name}})
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct OpenRouterUsage {
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub total_tokens: u64,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum OpenRouterStreamEvent {
    ContentDelta {
        text: String,
    },
    ToolCallDelta {
        id: String,
        ordinal: u64,
        index: u64,
        name: Option<String>,
        arguments: Option<String>,
    },
    FinishReason {
        reason: String,
    },
    Usage(OpenRouterUsage),
    Done,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OpenRouterStreamErrorKind {
    InvalidUtf8,
    InvalidJson,
    InvalidFrame,
    EventTooLarge,
    ProviderError,
    MissingUsage,
    DuplicateUsage,
    MissingTerminal,
    DuplicateTerminal,
    DataAfterTerminal,
    ParserFailed,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct OpenRouterStreamError {
    pub kind: OpenRouterStreamErrorKind,
    pub byte_offset: u64,
}

impl fmt::Display for OpenRouterStreamError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str("the OpenRouter response stream was invalid")
    }
}

impl std::error::Error for OpenRouterStreamError {}

#[derive(Clone, Debug)]
struct ToolState {
    id: String,
    ordinal: u64,
}

pub struct OpenRouterSseParser {
    attempt_ref: String,
    line_buffer: Vec<u8>,
    data_lines: Vec<String>,
    data_bytes: usize,
    max_event_bytes: usize,
    byte_offset: u64,
    tool_states: BTreeMap<u64, ToolState>,
    next_tool_ordinal: u64,
    usage_seen: bool,
    terminal: bool,
    failed: bool,
    replay_guard: OpenRouterReplayGuard,
}

impl fmt::Debug for OpenRouterSseParser {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("OpenRouterSseParser")
            .field("buffered_line_bytes", &self.line_buffer.len())
            .field("data_bytes", &self.data_bytes)
            .field("max_event_bytes", &self.max_event_bytes)
            .field("byte_offset", &self.byte_offset)
            .field("tool_count", &self.tool_states.len())
            .field("usage_seen", &self.usage_seen)
            .field("terminal", &self.terminal)
            .field("failed", &self.failed)
            .field("replay_decision", &self.replay_guard.decision())
            .finish()
    }
}

impl OpenRouterSseParser {
    pub fn new(attempt_id: &AttemptId) -> Self {
        Self::with_max_event_bytes(attempt_id, DEFAULT_MAX_SSE_EVENT_BYTES)
    }

    pub fn with_max_event_bytes(attempt_id: &AttemptId, max_event_bytes: usize) -> Self {
        Self {
            attempt_ref: attempt_id.as_str().to_owned(),
            line_buffer: Vec::new(),
            data_lines: Vec::new(),
            data_bytes: 0,
            max_event_bytes: max_event_bytes.max(1),
            byte_offset: 0,
            tool_states: BTreeMap::new(),
            next_tool_ordinal: 0,
            usage_seen: false,
            terminal: false,
            failed: false,
            replay_guard: OpenRouterReplayGuard::default(),
        }
    }

    pub fn mark_body_written(&mut self) {
        self.replay_guard.mark_body_written();
    }

    pub fn mark_forwarded_tool_result(&mut self) {
        self.replay_guard.mark_forwarded_tool_result();
    }

    pub fn mark_transport_outcome_unknown(&mut self) {
        self.replay_guard.mark_transport_outcome_unknown();
    }

    pub fn replay_decision(&self) -> ReplayDecision {
        self.replay_guard.decision()
    }

    pub fn feed(
        &mut self,
        bytes: &[u8],
    ) -> Result<Vec<OpenRouterStreamEvent>, OpenRouterStreamError> {
        if self.failed {
            return Err(self.error(OpenRouterStreamErrorKind::ParserFailed));
        }
        let mut events = Vec::new();
        for &byte in bytes {
            self.byte_offset = self.byte_offset.saturating_add(1);
            if byte == b'\n' {
                let mut line = std::mem::take(&mut self.line_buffer);
                if line.last() == Some(&b'\r') {
                    line.pop();
                }
                if let Err(error) = self.process_line(&line, &mut events) {
                    self.failed = true;
                    return Err(error);
                }
            } else {
                self.line_buffer.push(byte);
                if self.line_buffer.len() > self.max_event_bytes {
                    self.failed = true;
                    return Err(self.error(OpenRouterStreamErrorKind::EventTooLarge));
                }
            }
        }
        Ok(events)
    }

    pub fn finish(&mut self) -> Result<Vec<OpenRouterStreamEvent>, OpenRouterStreamError> {
        if self.failed {
            return Err(self.error(OpenRouterStreamErrorKind::ParserFailed));
        }
        let mut events = Vec::new();
        if !self.line_buffer.is_empty() {
            let mut line = std::mem::take(&mut self.line_buffer);
            if line.last() == Some(&b'\r') {
                line.pop();
            }
            if let Err(error) = self.process_line(&line, &mut events) {
                self.failed = true;
                return Err(error);
            }
        }
        if !self.data_lines.is_empty()
            && let Err(error) = self.dispatch_event(&mut events)
        {
            self.failed = true;
            return Err(error);
        }
        if !self.terminal {
            self.failed = true;
            return Err(self.error(OpenRouterStreamErrorKind::MissingTerminal));
        }
        Ok(events)
    }

    fn process_line(
        &mut self,
        line: &[u8],
        events: &mut Vec<OpenRouterStreamEvent>,
    ) -> Result<(), OpenRouterStreamError> {
        let line = std::str::from_utf8(line)
            .map_err(|_| self.error(OpenRouterStreamErrorKind::InvalidUtf8))?;
        if line.is_empty() {
            return self.dispatch_event(events);
        }
        if line.starts_with(':') {
            return Ok(());
        }
        let (field, value) = line.split_once(':').unwrap_or((line, ""));
        if field != "data" {
            return Ok(());
        }
        let value = value.strip_prefix(' ').unwrap_or(value);
        let separator = usize::from(!self.data_lines.is_empty());
        self.data_bytes = self
            .data_bytes
            .saturating_add(value.len())
            .saturating_add(separator);
        if self.data_bytes > self.max_event_bytes {
            return Err(self.error(OpenRouterStreamErrorKind::EventTooLarge));
        }
        self.data_lines.push(value.to_owned());
        Ok(())
    }

    fn dispatch_event(
        &mut self,
        events: &mut Vec<OpenRouterStreamEvent>,
    ) -> Result<(), OpenRouterStreamError> {
        if self.data_lines.is_empty() {
            return Ok(());
        }
        let data = self.data_lines.join("\n");
        self.data_lines.clear();
        self.data_bytes = 0;
        if data == "[DONE]" {
            if self.terminal {
                return Err(self.error(OpenRouterStreamErrorKind::DuplicateTerminal));
            }
            if !self.usage_seen {
                return Err(self.error(OpenRouterStreamErrorKind::MissingUsage));
            }
            self.terminal = true;
            events.push(OpenRouterStreamEvent::Done);
            return Ok(());
        }
        if self.terminal {
            return Err(self.error(OpenRouterStreamErrorKind::DataAfterTerminal));
        }
        let value: Value = serde_json::from_str(&data)
            .map_err(|_| self.error(OpenRouterStreamErrorKind::InvalidJson))?;
        let object = value
            .as_object()
            .ok_or_else(|| self.error(OpenRouterStreamErrorKind::InvalidFrame))?;
        if object.contains_key("error") {
            return Err(self.error(OpenRouterStreamErrorKind::ProviderError));
        }
        self.parse_usage(object, events)?;
        self.parse_choices(object, events)
    }

    fn parse_usage(
        &mut self,
        object: &Map<String, Value>,
        events: &mut Vec<OpenRouterStreamEvent>,
    ) -> Result<(), OpenRouterStreamError> {
        let Some(usage) = object.get("usage") else {
            return Ok(());
        };
        if usage.is_null() {
            return Ok(());
        }
        if self.usage_seen {
            return Err(self.error(OpenRouterStreamErrorKind::DuplicateUsage));
        }
        let usage = usage
            .as_object()
            .ok_or_else(|| self.error(OpenRouterStreamErrorKind::InvalidFrame))?;
        let prompt_tokens = required_u64(usage, "prompt_tokens", self.byte_offset)?;
        let completion_tokens = required_u64(usage, "completion_tokens", self.byte_offset)?;
        let total_tokens = required_u64(usage, "total_tokens", self.byte_offset)?;
        self.usage_seen = true;
        events.push(OpenRouterStreamEvent::Usage(OpenRouterUsage {
            prompt_tokens,
            completion_tokens,
            total_tokens,
        }));
        Ok(())
    }

    fn parse_choices(
        &mut self,
        object: &Map<String, Value>,
        events: &mut Vec<OpenRouterStreamEvent>,
    ) -> Result<(), OpenRouterStreamError> {
        let Some(choices) = object.get("choices") else {
            return Ok(());
        };
        let choices = choices
            .as_array()
            .ok_or_else(|| self.error(OpenRouterStreamErrorKind::InvalidFrame))?;
        for choice in choices {
            let choice = choice
                .as_object()
                .ok_or_else(|| self.error(OpenRouterStreamErrorKind::InvalidFrame))?;
            if choice.contains_key("error") {
                return Err(self.error(OpenRouterStreamErrorKind::ProviderError));
            }
            if let Some(delta) = choice.get("delta") {
                let delta = delta
                    .as_object()
                    .ok_or_else(|| self.error(OpenRouterStreamErrorKind::InvalidFrame))?;
                if delta.contains_key("error") {
                    return Err(self.error(OpenRouterStreamErrorKind::ProviderError));
                }
                if let Some(content) = delta.get("content")
                    && !content.is_null()
                {
                    let text = content
                        .as_str()
                        .ok_or_else(|| self.error(OpenRouterStreamErrorKind::InvalidFrame))?;
                    if !text.is_empty() {
                        let event = OpenRouterStreamEvent::ContentDelta {
                            text: text.to_owned(),
                        };
                        self.replay_guard.observe(&event);
                        events.push(event);
                    }
                }
                if let Some(tool_calls) = delta.get("tool_calls") {
                    self.parse_tool_calls(tool_calls, events)?;
                }
            }
            if let Some(reason) = choice.get("finish_reason")
                && !reason.is_null()
            {
                let reason = reason
                    .as_str()
                    .ok_or_else(|| self.error(OpenRouterStreamErrorKind::InvalidFrame))?;
                if reason == "error" {
                    return Err(self.error(OpenRouterStreamErrorKind::ProviderError));
                }
                events.push(OpenRouterStreamEvent::FinishReason {
                    reason: reason.to_owned(),
                });
            }
        }
        Ok(())
    }

    fn parse_tool_calls(
        &mut self,
        value: &Value,
        events: &mut Vec<OpenRouterStreamEvent>,
    ) -> Result<(), OpenRouterStreamError> {
        let calls = value
            .as_array()
            .ok_or_else(|| self.error(OpenRouterStreamErrorKind::InvalidFrame))?;
        for call in calls {
            let call = call
                .as_object()
                .ok_or_else(|| self.error(OpenRouterStreamErrorKind::InvalidFrame))?;
            let index = call
                .get("index")
                .and_then(Value::as_u64)
                .ok_or_else(|| self.error(OpenRouterStreamErrorKind::InvalidFrame))?;
            let incoming_id = call.get("id").and_then(Value::as_str);
            let state = if let Some(state) = self.tool_states.get(&index) {
                if incoming_id.is_some_and(|id| id != state.id) {
                    return Err(self.error(OpenRouterStreamErrorKind::InvalidFrame));
                }
                state.clone()
            } else {
                let ordinal = self.next_tool_ordinal;
                self.next_tool_ordinal = self.next_tool_ordinal.saturating_add(1);
                let state = ToolState {
                    id: incoming_id.map_or_else(
                        || synthesized_tool_id(&self.attempt_ref, ordinal),
                        str::to_owned,
                    ),
                    ordinal,
                };
                self.tool_states.insert(index, state.clone());
                state
            };
            let function = call.get("function").and_then(Value::as_object);
            let name = function
                .and_then(|function| function.get("name"))
                .and_then(Value::as_str)
                .map(str::to_owned);
            let arguments = function
                .and_then(|function| function.get("arguments"))
                .and_then(Value::as_str)
                .map(str::to_owned);
            if name.is_some() || arguments.is_some() || incoming_id.is_some() {
                let event = OpenRouterStreamEvent::ToolCallDelta {
                    id: state.id,
                    ordinal: state.ordinal,
                    index,
                    name,
                    arguments,
                };
                self.replay_guard.observe(&event);
                events.push(event);
            }
        }
        Ok(())
    }

    fn error(&self, kind: OpenRouterStreamErrorKind) -> OpenRouterStreamError {
        OpenRouterStreamError {
            kind,
            byte_offset: self.byte_offset,
        }
    }
}

fn required_u64(
    object: &Map<String, Value>,
    name: &str,
    byte_offset: u64,
) -> Result<u64, OpenRouterStreamError> {
    object
        .get(name)
        .and_then(Value::as_u64)
        .ok_or(OpenRouterStreamError {
            kind: OpenRouterStreamErrorKind::InvalidFrame,
            byte_offset,
        })
}

fn synthesized_tool_id(attempt_ref: &str, ordinal: u64) -> String {
    const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
    const FNV_PRIME: u64 = 0x0000_0100_0000_01b3;
    let mut value = FNV_OFFSET;
    for byte in attempt_ref
        .bytes()
        .chain(*b":")
        .chain(ordinal.to_be_bytes())
    {
        value ^= u64::from(byte);
        value = value.wrapping_mul(FNV_PRIME);
    }
    format!("call_{value:016x}")
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct OpenRouterReplayGuard {
    boundary: Option<ReplayBoundary>,
    body_written: bool,
}

impl OpenRouterReplayGuard {
    pub fn mark_body_written(&mut self) {
        self.body_written = true;
    }

    pub fn observe(&mut self, event: &OpenRouterStreamEvent) {
        match event {
            OpenRouterStreamEvent::ContentDelta { text } if !text.is_empty() => {
                self.cross(ReplayBoundary::PublicContent);
            }
            OpenRouterStreamEvent::ToolCallDelta { .. } => {
                self.cross(ReplayBoundary::VisibleToolCall);
            }
            OpenRouterStreamEvent::ContentDelta { .. }
            | OpenRouterStreamEvent::FinishReason { .. }
            | OpenRouterStreamEvent::Usage(_)
            | OpenRouterStreamEvent::Done => {}
        }
    }

    pub fn mark_forwarded_tool_result(&mut self) {
        self.cross(ReplayBoundary::ForwardedToolResult);
    }

    pub fn mark_transport_outcome_unknown(&mut self) {
        if self.body_written {
            self.cross(ReplayBoundary::AmbiguousDispatch);
        }
    }

    pub fn decision(&self) -> ReplayDecision {
        self.boundary
            .map_or(ReplayDecision::Allowed, ReplayDecision::Denied)
    }

    fn cross(&mut self, boundary: ReplayBoundary) {
        if self.boundary.is_none() {
            self.boundary = Some(boundary);
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OpenRouterFailureClass {
    InvalidRequest,
    Authentication,
    PaymentRequired,
    Authorization,
    NotFound,
    Timeout,
    Unprocessable,
    RateLimited,
    UpstreamServer,
    UnexpectedStatus,
    MismatchedContentType,
    RedirectRejected,
    Cancelled,
    TransportTimeout,
    AmbiguousDispatch,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct OpenRouterFailure {
    pub class: OpenRouterFailureClass,
    pub status: Option<u16>,
    pub retry_after: Option<BoundedRetryAfter>,
}

/// The only retry authorization exposed by the OpenRouter adapter.  This is
/// intentionally a policy decision rather than a transport convenience: the
/// caller still owns the budget reservation and must record the result.
pub const AUTHORIZED_RATE_LIMIT_RETRY_AFTER_MAX_SECONDS: u32 = 2;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OpenRouterRetryDecision {
    /// Exactly one retry may be attempted before any visible or ambiguous
    /// boundary. The value is the provider-supplied bounded delay.
    RetryOnce { retry_after_seconds: u32 },
    /// The request must become a terminal/exclusion receipt or fall through
    /// according to the surrounding pre-visible cascade.
    NoRetry,
}

/// Applies the candidate's deliberately narrow 429 policy.  A missing,
/// capped, or long `Retry-After` is not a retry authorization.  Likewise,
/// this never grants a retry after an ambiguous write, visible output, a tool
/// effect, or an already terminal attempt.
pub fn authorized_rate_limit_retry(
    failure: OpenRouterFailure,
    replay: ReplayDecision,
    retry_already_used: bool,
) -> OpenRouterRetryDecision {
    let RetryDecisionInputs {
        retry_after_seconds,
    } = RetryDecisionInputs::from(failure);
    if retry_already_used
        || !matches!(failure.class, OpenRouterFailureClass::RateLimited)
        || !matches!(replay, ReplayDecision::Allowed)
        || retry_after_seconds
            .is_none_or(|seconds| seconds > AUTHORIZED_RATE_LIMIT_RETRY_AFTER_MAX_SECONDS)
    {
        return OpenRouterRetryDecision::NoRetry;
    }
    match retry_after_seconds {
        Some(retry_after_seconds) => OpenRouterRetryDecision::RetryOnce {
            retry_after_seconds,
        },
        None => OpenRouterRetryDecision::NoRetry,
    }
}

struct RetryDecisionInputs {
    retry_after_seconds: Option<u32>,
}

impl From<OpenRouterFailure> for RetryDecisionInputs {
    fn from(failure: OpenRouterFailure) -> Self {
        Self {
            retry_after_seconds: failure
                .retry_after
                .filter(|retry_after| !retry_after.capped)
                .map(|retry_after| retry_after.seconds),
        }
    }
}

pub fn normalize_http_response(
    status: u16,
    content_type: Option<&str>,
    retry_after: Option<&str>,
    redirected: bool,
) -> Result<(), OpenRouterFailure> {
    if redirected || (300..400).contains(&status) {
        return Err(failure(
            OpenRouterFailureClass::RedirectRejected,
            Some(status),
            retry_after,
        ));
    }
    if status != 200 {
        let class = match status {
            400 => OpenRouterFailureClass::InvalidRequest,
            401 => OpenRouterFailureClass::Authentication,
            402 => OpenRouterFailureClass::PaymentRequired,
            403 => OpenRouterFailureClass::Authorization,
            404 => OpenRouterFailureClass::NotFound,
            408 => OpenRouterFailureClass::Timeout,
            422 => OpenRouterFailureClass::Unprocessable,
            429 => OpenRouterFailureClass::RateLimited,
            500..=599 => OpenRouterFailureClass::UpstreamServer,
            _ => OpenRouterFailureClass::UnexpectedStatus,
        };
        return Err(failure(class, Some(status), retry_after));
    }
    let media_type = content_type
        .and_then(|value| value.split(';').next())
        .map(str::trim);
    if media_type != Some("text/event-stream") {
        return Err(failure(
            OpenRouterFailureClass::MismatchedContentType,
            Some(status),
            None,
        ));
    }
    Ok(())
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BodyWriteState {
    NotWritten,
    Written,
    Unknown,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TransportFailureSignal {
    Cancelled,
    Timeout,
}

pub fn normalize_transport_failure(
    signal: TransportFailureSignal,
    body_write: BodyWriteState,
) -> OpenRouterFailure {
    let class = match body_write {
        BodyWriteState::Written | BodyWriteState::Unknown => {
            OpenRouterFailureClass::AmbiguousDispatch
        }
        BodyWriteState::NotWritten => match signal {
            TransportFailureSignal::Cancelled => OpenRouterFailureClass::Cancelled,
            TransportFailureSignal::Timeout => OpenRouterFailureClass::TransportTimeout,
        },
    };
    failure(class, None, None)
}

fn failure(
    class: OpenRouterFailureClass,
    status: Option<u16>,
    retry_after: Option<&str>,
) -> OpenRouterFailure {
    OpenRouterFailure {
        class,
        status,
        retry_after: BoundedRetryAfter::parse(retry_after),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use router_protocol::{
        BillingAuthorityClass, CapabilityRequirements, CatalogSourceKind, CredentialSourceKind,
        EndpointIdentity, ExecutionLocation, FunctionTool, HarnessInstanceId, M1_WIRE_VERSION,
        ModelNamespace, PhasePlacement, ProviderService, ReasoningIntent, SessionId, TurnId,
    };

    fn route() -> RouteIdentity {
        RouteIdentity {
            phase: PhaseKind::OpenrouterApi,
            placement: PhasePlacement::Builtin,
            provider: ProviderService::Openrouter,
            transport: TransportKind::DirectHttps,
            endpoint: EndpointIdentity {
                origin: "https://openrouter.ai".to_owned(),
                path: "/api/v1/chat/completions".to_owned(),
            },
            credential_source: CredentialSourceKind::SecretRef,
            credential_ref: Some("keychain:openrouter-secret-canary".to_owned()),
            credential_refs: vec!["keychain:openrouter-secret-canary".to_owned()],
            catalog_source: CatalogSourceKind::OpenrouterCatalog,
            catalog_ref: "catalog:openrouter:verified".to_owned(),
            billing_authority: BillingAuthorityClass::OpenrouterMeteredApi,
            model_namespace: ModelNamespace::Openrouter,
            model_id: "anthropic/claude-sonnet-4.5".to_owned(),
            execution_location: ExecutionLocation::RemoteCloud,
            account_ref: "account:openrouter".to_owned(),
            capacity_domain_ref: "capacity:openrouter".to_owned(),
            local_execution_proof: false,
            approved_local_inventory: Vec::new(),
            reasoning_modes: Vec::new(),
        }
    }

    fn request(prompt: &str) -> CanonicalRequest {
        CanonicalRequest {
            wire_version: M1_WIRE_VERSION.to_owned(),
            model: "router-auto".to_owned(),
            stream: true,
            harness_instance: HarnessInstanceId::parse("fixture-cli").expect("fixture"),
            session_id: SessionId::parse("session-openrouter").expect("fixture"),
            turn_id: TurnId::parse("turn-openrouter").expect("fixture"),
            input: vec![InputItem::Message {
                role: MessageRole::User,
                content: vec![ContentBlock::Text {
                    text: prompt.to_owned(),
                }],
            }],
            tools: vec![FunctionTool {
                namespace: Some("weather".to_owned()),
                name: "lookup".to_owned(),
                description: "lookup weather".to_owned(),
                parameters: json!({"type":"object"}),
                strict: true,
            }],
            tool_choice: ToolChoice::Auto,
            structured_output: None,
            reasoning: ReasoningIntent::None,
            max_output_tokens: Some(64),
            max_spend_micro_usd: None,
            capabilities: CapabilityRequirements::default(),
        }
    }

    #[test]
    fn exact_request_has_fixed_provider_policy_and_redacted_secret_and_prompt() {
        let spec = build_request(&route(), &request("prompt-canary")).expect("request");
        assert_eq!(spec.method, HttpMethod::Post);
        assert_eq!(spec.origin, "https://openrouter.ai");
        assert_eq!(spec.path, "/api/v1/chat/completions");
        assert_eq!(spec.accept, "text/event-stream");
        assert!(spec.network.require_tls);
        assert!(!spec.network.follow_redirects);
        assert_eq!(
            spec.body(),
            &json!({
                "model":"anthropic/claude-sonnet-4.5",
                "messages":[{"role":"user","content":"prompt-canary"}],
                "stream":true,
                "stream_options":{"include_usage":true},
                "provider":{
                    "require_parameters":true,
                    "allow_fallbacks":false,
                    "data_collection":"deny",
                    "zdr":true
                },
                "tool_choice":"auto",
                "tools":[{"type":"function","function":{
                    "name":"weather__lookup",
                    "description":"lookup weather",
                    "parameters":{"type":"object"},
                    "strict":true
                }}],
                "max_tokens":64
            })
        );
        assert_eq!(
            spec.bearer_secret.as_str(),
            "keychain:openrouter-secret-canary"
        );
        let first_payload = spec.body_bytes().expect("serialize exact payload");
        let second_payload = spec.body_bytes().expect("serialize exact payload again");
        assert_eq!(first_payload, second_payload);
        assert!(
            !first_payload
                .windows("openrouter-secret-canary".len())
                .any(|window| window == b"openrouter-secret-canary")
        );
        let debug = format!("{spec:?}");
        assert!(!debug.contains("openrouter-secret-canary"));
        assert!(!debug.contains("prompt-canary"));
    }

    #[test]
    fn public_summary_reasoning_keeps_strict_provider_policy_without_reasoning_param() {
        let mut request = request("prompt-canary");
        request.reasoning = ReasoningIntent::PublicSummary;
        let spec = build_request(&route(), &request).expect("request");
        let body = spec.body().as_object().expect("body object");
        assert!(body.contains_key("provider"));
        assert_eq!(
            body.get("provider"),
            Some(&json!({
                "require_parameters": true,
                "allow_fallbacks": false,
                "data_collection": "deny",
                "zdr": true,
            }))
        );
        assert!(
            !body.contains_key("reasoning"),
            "strict OpenRouter provider routing rejects a reasoning parameter for flash models"
        );
    }

    #[test]
    fn sse_is_every_byte_fragment_safe_with_crlf_comments_multiline_and_usage_once() {
        let transcript = concat!(
            ": heartbeat\r\n\r\n",
            "data: {\"choices\":\r\n",
            "data: [{\"delta\":{\"content\":\"hé\"},\"finish_reason\":null}]}\r\n\r\n",
            "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call-provider\",\"function\":{\"name\":\"weather__lookup\",\"arguments\":\"{\"}}]}}]}\r\n\r\n",
            "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"\\\"city\\\":\\\"Paris\\\"}\"}}]},\"finish_reason\":\"tool_calls\"}]}\r\n\r\n",
            "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":7,\"total_tokens\":18}}\r\n\r\n",
            "data: [DONE]"
        );
        let attempt = AttemptId::parse("attempt-openrouter").expect("attempt");
        let mut parser = OpenRouterSseParser::new(&attempt);
        let mut fragmented = Vec::new();
        for byte in transcript.as_bytes() {
            fragmented.extend(parser.feed(std::slice::from_ref(byte)).expect("fragment"));
        }
        fragmented.extend(parser.finish().expect("complete"));
        assert!(matches!(
            &fragmented[0],
            OpenRouterStreamEvent::ContentDelta { text } if text == "hé"
        ));
        assert_eq!(
            fragmented
                .iter()
                .filter(|event| matches!(event, OpenRouterStreamEvent::Usage(_)))
                .count(),
            1
        );
        assert!(matches!(
            fragmented.last(),
            Some(OpenRouterStreamEvent::Done)
        ));

        let mut whole = OpenRouterSseParser::new(&attempt);
        let mut whole_events = whole.feed(transcript.as_bytes()).expect("whole stream");
        whole_events.extend(whole.finish().expect("complete"));
        assert_eq!(fragmented, whole_events);
    }

    #[test]
    fn post_visible_provider_error_denies_replay_without_echoing_canary() {
        let attempt = AttemptId::parse("attempt-error").expect("attempt");
        let mut parser = OpenRouterSseParser::new(&attempt);
        parser.mark_body_written();
        let error = parser
            .feed(
                concat!(
                    "data: {\"choices\":[{\"delta\":{\"content\":\"public\"}}]}\n\n",
                    "data: {\"error\":{\"message\":\"secret-canary\"}}\n\n"
                )
                .as_bytes(),
            )
            .expect_err("provider error");
        assert_eq!(error.kind, OpenRouterStreamErrorKind::ProviderError);
        assert_eq!(
            parser.replay_decision(),
            ReplayDecision::Denied(ReplayBoundary::PublicContent)
        );
        assert!(!format!("{error:?}").contains("secret-canary"));
    }

    #[test]
    fn usage_and_terminal_are_each_exactly_once() {
        let attempt = AttemptId::parse("attempt-terminal").expect("attempt");
        let usage = b"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n";
        let mut duplicate_usage = OpenRouterSseParser::new(&attempt);
        duplicate_usage.feed(usage).expect("first usage");
        assert_eq!(
            duplicate_usage
                .feed(usage)
                .expect_err("duplicate usage")
                .kind,
            OpenRouterStreamErrorKind::DuplicateUsage
        );

        let mut missing_usage = OpenRouterSseParser::new(&attempt);
        assert_eq!(
            missing_usage
                .feed(b"data: [DONE]\n\n")
                .expect_err("usage required")
                .kind,
            OpenRouterStreamErrorKind::MissingUsage
        );

        let mut duplicate_terminal = OpenRouterSseParser::new(&attempt);
        duplicate_terminal.feed(usage).expect("usage");
        duplicate_terminal
            .feed(b"data: [DONE]\n\n")
            .expect("terminal");
        assert_eq!(
            duplicate_terminal
                .feed(b"data: [DONE]\n\n")
                .expect_err("duplicate terminal")
                .kind,
            OpenRouterStreamErrorKind::DuplicateTerminal
        );
    }

    #[test]
    fn malformed_oversized_missing_terminal_and_finish_error_fail_closed() {
        let attempt = AttemptId::parse("attempt-invalid").expect("attempt");
        let cases: &[(&[u8], OpenRouterStreamErrorKind)] = &[
            (b"data: {bad}\n\n", OpenRouterStreamErrorKind::InvalidJson),
            (b"data: \xff\n\n", OpenRouterStreamErrorKind::InvalidUtf8),
            (
                b"data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"error\"}]}\n\n",
                OpenRouterStreamErrorKind::ProviderError,
            ),
        ];
        for (input, expected) in cases {
            let mut parser = OpenRouterSseParser::new(&attempt);
            assert_eq!(parser.feed(input).expect_err("invalid").kind, *expected);
        }
        let mut oversized = OpenRouterSseParser::with_max_event_bytes(&attempt, 8);
        assert_eq!(
            oversized
                .feed(b"data: 123456789")
                .expect_err("bounded")
                .kind,
            OpenRouterStreamErrorKind::EventTooLarge
        );
        let mut missing = OpenRouterSseParser::new(&attempt);
        missing
            .feed(b"data: {\"choices\":[]}\n\n")
            .expect("valid chunk");
        assert_eq!(
            missing.finish().expect_err("terminal required").kind,
            OpenRouterStreamErrorKind::MissingTerminal
        );
    }

    #[test]
    fn status_and_transport_failures_include_402_retry_bounds_and_ambiguity() {
        let matrix = [
            (400, OpenRouterFailureClass::InvalidRequest),
            (401, OpenRouterFailureClass::Authentication),
            (402, OpenRouterFailureClass::PaymentRequired),
            (403, OpenRouterFailureClass::Authorization),
            (404, OpenRouterFailureClass::NotFound),
            (408, OpenRouterFailureClass::Timeout),
            (422, OpenRouterFailureClass::Unprocessable),
            (429, OpenRouterFailureClass::RateLimited),
            (503, OpenRouterFailureClass::UpstreamServer),
        ];
        for (status, class) in matrix {
            let failure =
                normalize_http_response(status, None, Some("999999"), false).expect_err("failure");
            assert_eq!(failure.class, class);
            assert!(failure.retry_after.is_some_and(|retry| retry.capped));
        }
        assert!(
            normalize_http_response(200, Some("text/event-stream; charset=utf-8"), None, false)
                .is_ok()
        );
        assert_eq!(
            normalize_http_response(200, Some("application/json"), None, false)
                .expect_err("content type")
                .class,
            OpenRouterFailureClass::MismatchedContentType
        );
        assert_eq!(
            normalize_http_response(302, None, None, true)
                .expect_err("redirect")
                .class,
            OpenRouterFailureClass::RedirectRejected
        );
        assert_eq!(
            normalize_transport_failure(TransportFailureSignal::Timeout, BodyWriteState::Written,)
                .class,
            OpenRouterFailureClass::AmbiguousDispatch
        );
        let mut guard = OpenRouterReplayGuard::default();
        guard.mark_body_written();
        guard.mark_transport_outcome_unknown();
        assert_eq!(
            guard.decision(),
            ReplayDecision::Denied(ReplayBoundary::AmbiguousDispatch)
        );
    }

    #[test]
    fn rate_limit_retry_requires_one_short_explicit_delay_before_visibility() {
        let rate_limited =
            normalize_http_response(429, None, Some("2"), false).expect_err("rate limited");
        assert_eq!(
            authorized_rate_limit_retry(rate_limited, ReplayDecision::Allowed, false),
            OpenRouterRetryDecision::RetryOnce {
                retry_after_seconds: 2
            }
        );

        for (retry_after, replay, already_used) in [
            (None, ReplayDecision::Allowed, false),
            (Some("3"), ReplayDecision::Allowed, false),
            (Some("2"), ReplayDecision::Allowed, true),
            (
                Some("2"),
                ReplayDecision::Denied(ReplayBoundary::AmbiguousDispatch),
                false,
            ),
            (
                Some("2"),
                ReplayDecision::Denied(ReplayBoundary::PublicContent),
                false,
            ),
            (
                Some("2"),
                ReplayDecision::Denied(ReplayBoundary::VisibleToolCall),
                false,
            ),
            (
                Some("2"),
                ReplayDecision::Denied(ReplayBoundary::ForwardedToolResult),
                false,
            ),
        ] {
            let failure =
                normalize_http_response(429, None, retry_after, false).expect_err("rate limited");
            assert_eq!(
                authorized_rate_limit_retry(failure, replay, already_used),
                OpenRouterRetryDecision::NoRetry
            );
        }
    }
}