regy 0.1.0

use serde::{Deserialize, Deserializer, Serialize, de::Error as _, ser::SerializeMap};
use serde_json::{Map, Value};

// Pi supports image payloads even though Regy's current UI does not emit them.
#[allow(dead_code)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub(crate) enum PiImageType {
    #[serde(rename = "image")]
    Image,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct PiImageContent {
    #[serde(rename = "type")]
    pub(crate) kind: PiImageType,
    pub(crate) data: String,
    pub(crate) mime_type: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
pub(crate) enum StreamingBehavior {
    #[serde(rename = "steer")]
    // Pi's protocol retains this valid command mode for future UI support.
    #[allow(dead_code)]
    Steer,
    #[serde(rename = "followUp")]
    // Pi's protocol retains this valid command mode for future UI support.
    #[allow(dead_code)]
    FollowUp,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub(crate) enum ThinkingLevel {
    Off,
    Minimal,
    Low,
    Medium,
    High,
    XHigh,
    Max,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ExtensionUiResponsePayload {
    Value(String),
    Confirmed(bool),
    Cancelled,
}

impl Serialize for ExtensionUiResponsePayload {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        let mut map = serializer.serialize_map(Some(1))?;
        match self {
            Self::Value(value) => map.serialize_entry("value", value)?,
            Self::Confirmed(confirmed) => map.serialize_entry("confirmed", confirmed)?,
            Self::Cancelled => map.serialize_entry("cancelled", &true)?,
        }
        map.end()
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub(crate) enum PiRpcCommand {
    Prompt {
        id: String,
        message: String,
        #[serde(skip_serializing_if = "Option::is_none")]
        images: Option<Vec<PiImageContent>>,
        #[serde(rename = "streamingBehavior", skip_serializing_if = "Option::is_none")]
        streaming_behavior: Option<StreamingBehavior>,
    },
    Abort {
        id: String,
    },
    GetState {
        id: String,
    },
    GetMessages {
        id: String,
    },
    GetAvailableModels {
        id: String,
    },
    SetModel {
        id: String,
        provider: String,
        #[serde(rename = "modelId")]
        model_id: String,
    },
    SetThinkingLevel {
        id: String,
        level: ThinkingLevel,
    },
    GetCommands {
        id: String,
    },
    GetSessionStats {
        id: String,
    },
    ExtensionUiResponse {
        id: String,
        #[serde(flatten)]
        response: ExtensionUiResponsePayload,
    },
}

impl PiRpcCommand {
    pub(crate) fn id(&self) -> &str {
        match self {
            Self::Prompt { id, .. }
            | Self::Abort { id }
            | Self::GetState { id }
            | Self::GetMessages { id }
            | Self::GetAvailableModels { id }
            | Self::SetModel { id, .. }
            | Self::SetThinkingLevel { id, .. }
            | Self::GetCommands { id }
            | Self::GetSessionStats { id }
            | Self::ExtensionUiResponse { id, .. } => id,
        }
    }

    pub(crate) fn expects_response(&self) -> bool {
        !matches!(self, Self::ExtensionUiResponse { .. })
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct PiSessionState {
    pub(crate) model: Option<Value>,
    pub(crate) thinking_level: ThinkingLevel,
    pub(crate) is_streaming: bool,
    pub(crate) is_compacting: bool,
    pub(crate) steering_mode: String,
    pub(crate) follow_up_mode: String,
    pub(crate) session_file: Option<String>,
    pub(crate) session_id: String,
    pub(crate) session_name: Option<String>,
    pub(crate) auto_compaction_enabled: bool,
    pub(crate) message_count: u64,
    pub(crate) pending_message_count: u64,
}

#[derive(Debug, Clone, PartialEq, Deserialize)]
pub(crate) struct PiMessagesData {
    pub(crate) messages: Vec<Value>,
}

#[derive(Debug, Clone, PartialEq, Deserialize)]
pub(crate) struct PiAvailableModelsData {
    pub(crate) models: Vec<Value>,
}

#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct PiCommandInfo {
    pub(crate) name: String,
    pub(crate) description: Option<String>,
    pub(crate) source: String,
    pub(crate) location: Option<String>,
    pub(crate) path: Option<String>,
    pub(crate) source_info: Option<Value>,
}

#[derive(Debug, Clone, PartialEq, Deserialize)]
pub(crate) struct PiCommandsData {
    pub(crate) commands: Vec<PiCommandInfo>,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct PiSessionTokens {
    pub(crate) input: u64,
    pub(crate) output: u64,
    pub(crate) cache_read: u64,
    pub(crate) cache_write: u64,
    pub(crate) total: u64,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct PiContextUsage {
    pub(crate) tokens: Option<u64>,
    pub(crate) context_window: u64,
    pub(crate) percent: Option<f64>,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct PiSessionStats {
    pub(crate) session_file: Option<String>,
    pub(crate) session_id: String,
    pub(crate) user_messages: u64,
    pub(crate) assistant_messages: u64,
    pub(crate) tool_calls: u64,
    pub(crate) tool_results: u64,
    pub(crate) total_messages: u64,
    pub(crate) tokens: PiSessionTokens,
    pub(crate) cost: f64,
    pub(crate) context_usage: Option<PiContextUsage>,
}

#[derive(Debug, Clone, PartialEq)]
pub(crate) enum PiRpcResponse {
    Prompt {
        id: Option<String>,
    },
    Abort {
        id: Option<String>,
    },
    GetState {
        id: Option<String>,
        data: PiSessionState,
    },
    GetMessages {
        id: Option<String>,
        data: PiMessagesData,
    },
    GetAvailableModels {
        id: Option<String>,
        data: PiAvailableModelsData,
    },
    SetModel {
        id: Option<String>,
        data: Value,
    },
    SetThinkingLevel {
        id: Option<String>,
    },
    GetCommands {
        id: Option<String>,
        data: PiCommandsData,
    },
    GetSessionStats {
        id: Option<String>,
        data: PiSessionStats,
    },
    Failure {
        id: Option<String>,
        command: String,
        error: String,
    },
}

impl PiRpcResponse {
    pub(crate) fn id(&self) -> Option<&str> {
        match self {
            Self::Prompt { id }
            | Self::Abort { id }
            | Self::GetState { id, .. }
            | Self::GetMessages { id, .. }
            | Self::GetAvailableModels { id, .. }
            | Self::SetModel { id, .. }
            | Self::SetThinkingLevel { id }
            | Self::GetCommands { id, .. }
            | Self::GetSessionStats { id, .. }
            | Self::Failure { id, .. } => id.as_deref(),
        }
    }

    pub(crate) fn command(&self) -> &str {
        match self {
            Self::Prompt { .. } => "prompt",
            Self::Abort { .. } => "abort",
            Self::GetState { .. } => "get_state",
            Self::GetMessages { .. } => "get_messages",
            Self::GetAvailableModels { .. } => "get_available_models",
            Self::SetModel { .. } => "set_model",
            Self::SetThinkingLevel { .. } => "set_thinking_level",
            Self::GetCommands { .. } => "get_commands",
            Self::GetSessionStats { .. } => "get_session_stats",
            Self::Failure { command, .. } => command,
        }
    }

    pub(crate) fn success(&self) -> bool {
        !matches!(self, Self::Failure { .. })
    }

    pub(crate) fn error(&self) -> Option<&str> {
        match self {
            Self::Failure { error, .. } => Some(error),
            _ => None,
        }
    }
}

impl<'de> Deserialize<'de> for PiRpcResponse {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let mut object = Map::<String, Value>::deserialize(deserializer)?;
        require_string::<D::Error>(&mut object, "type").and_then(|kind| {
            (kind == "response")
                .then_some(())
                .ok_or_else(|| D::Error::custom("Pi RPC response type must be response"))
        })?;
        let id = optional_string::<D::Error>(&mut object, "id")?;
        let command = require_string::<D::Error>(&mut object, "command")?;
        let success = require_bool::<D::Error>(&mut object, "success")?;

        let response = if success {
            if object.contains_key("error") {
                return Err(D::Error::custom(
                    "successful Pi RPC response cannot contain error",
                ));
            }
            let data = object.remove("data");
            match command.as_str() {
                "prompt" => no_data::<D::Error>(data, Self::Prompt { id })?,
                "abort" => no_data::<D::Error>(data, Self::Abort { id })?,
                "get_state" => Self::GetState {
                    id,
                    data: required_data::<D::Error, _>(data, "get_state")?,
                },
                "get_messages" => Self::GetMessages {
                    id,
                    data: required_data::<D::Error, _>(data, "get_messages")?,
                },
                "get_available_models" => Self::GetAvailableModels {
                    id,
                    data: required_data::<D::Error, _>(data, "get_available_models")?,
                },
                "set_model" => Self::SetModel {
                    id,
                    data: required_value::<D::Error>(data, "set_model")?,
                },
                "set_thinking_level" => no_data::<D::Error>(data, Self::SetThinkingLevel { id })?,
                "get_commands" => Self::GetCommands {
                    id,
                    data: required_data::<D::Error, _>(data, "get_commands")?,
                },
                "get_session_stats" => Self::GetSessionStats {
                    id,
                    data: required_data::<D::Error, _>(data, "get_session_stats")?,
                },
                _ => {
                    return Err(D::Error::custom(format!(
                        "unsupported successful Pi RPC command: {command}"
                    )));
                }
            }
        } else {
            if object.contains_key("data") {
                return Err(D::Error::custom(
                    "failed Pi RPC response cannot contain data",
                ));
            }
            Self::Failure {
                id,
                command,
                error: require_string::<D::Error>(&mut object, "error")?,
            }
        };

        if let Some(field) = object.keys().next() {
            return Err(D::Error::custom(format!(
                "unexpected Pi RPC response field: {field}"
            )));
        }
        Ok(response)
    }
}

fn require_string<E>(object: &mut Map<String, Value>, field: &str) -> Result<String, E>
where
    E: serde::de::Error,
{
    object
        .remove(field)
        .and_then(|value| value.as_str().map(str::to_owned))
        .ok_or_else(|| E::custom(format!("Pi RPC response requires string {field}")))
}

fn optional_string<E>(object: &mut Map<String, Value>, field: &str) -> Result<Option<String>, E>
where
    E: serde::de::Error,
{
    match object.remove(field) {
        None => Ok(None),
        Some(Value::String(value)) => Ok(Some(value)),
        Some(_) => Err(E::custom(format!(
            "Pi RPC response {field} must be a string"
        ))),
    }
}

fn require_bool<E>(object: &mut Map<String, Value>, field: &str) -> Result<bool, E>
where
    E: serde::de::Error,
{
    object
        .remove(field)
        .and_then(|value| value.as_bool())
        .ok_or_else(|| E::custom(format!("Pi RPC response requires boolean {field}")))
}

fn no_data<E>(data: Option<Value>, response: PiRpcResponse) -> Result<PiRpcResponse, E>
where
    E: serde::de::Error,
{
    if data.is_some() {
        return Err(E::custom(format!(
            "successful {} response must not contain data",
            response.command()
        )));
    }
    Ok(response)
}

fn required_value<E>(data: Option<Value>, command: &str) -> Result<Value, E>
where
    E: serde::de::Error,
{
    match data {
        Some(Value::Null) | None => Err(E::custom(format!(
            "successful {command} response requires data"
        ))),
        Some(value) => Ok(value),
    }
}

fn required_data<E, T>(data: Option<Value>, command: &str) -> Result<T, E>
where
    E: serde::de::Error,
    T: serde::de::DeserializeOwned,
{
    let value = required_value::<E>(data, command)?;
    serde_json::from_value(value).map_err(E::custom)
}

#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct PiCompactionResult {
    pub(crate) summary: String,
    pub(crate) first_kept_entry_id: String,
    pub(crate) tokens_before: u64,
    pub(crate) estimated_tokens_after: Option<u64>,
    pub(crate) details: Option<Value>,
}

#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub(crate) enum PiRpcEvent {
    AgentStart,
    AgentEnd {
        messages: Vec<Value>,
        #[serde(rename = "willRetry")]
        will_retry: bool,
    },
    AgentSettled,
    TurnStart,
    TurnEnd {
        message: Value,
        #[serde(rename = "toolResults")]
        tool_results: Vec<Value>,
    },
    MessageStart {
        message: Value,
    },
    MessageUpdate {
        message: Value,
        #[serde(rename = "assistantMessageEvent")]
        assistant_message_event: Value,
    },
    MessageEnd {
        message: Value,
    },
    ToolExecutionStart {
        #[serde(rename = "toolCallId")]
        tool_call_id: String,
        #[serde(rename = "toolName")]
        tool_name: String,
        args: Value,
    },
    ToolExecutionUpdate {
        #[serde(rename = "toolCallId")]
        tool_call_id: String,
        #[serde(rename = "toolName")]
        tool_name: String,
        args: Value,
        #[serde(rename = "partialResult")]
        partial_result: Value,
    },
    ToolExecutionEnd {
        #[serde(rename = "toolCallId")]
        tool_call_id: String,
        #[serde(rename = "toolName")]
        tool_name: String,
        result: Value,
        #[serde(rename = "isError")]
        is_error: bool,
    },
    QueueUpdate {
        steering: Vec<String>,
        #[serde(rename = "followUp")]
        follow_up: Vec<String>,
    },
    CompactionStart {
        reason: String,
    },
    CompactionEnd {
        reason: String,
        result: Option<PiCompactionResult>,
        aborted: bool,
        #[serde(rename = "willRetry")]
        will_retry: bool,
        #[serde(rename = "errorMessage")]
        error_message: Option<String>,
    },
    AutoRetryStart {
        attempt: u64,
        #[serde(rename = "maxAttempts")]
        max_attempts: u64,
        #[serde(rename = "delayMs")]
        delay_ms: u64,
        #[serde(rename = "errorMessage")]
        error_message: String,
    },
    AutoRetryEnd {
        success: bool,
        attempt: u64,
        #[serde(rename = "finalError")]
        final_error: Option<String>,
    },
    ExtensionError {
        #[serde(rename = "extensionPath")]
        extension_path: String,
        event: String,
        error: String,
    },
    ExtensionUiRequest {
        id: String,
        method: String,
        #[serde(flatten)]
        payload: Map<String, Value>,
    },
    ThinkingLevelChanged {
        level: ThinkingLevel,
    },
    SessionInfoChanged {
        name: Option<String>,
    },
}

#[derive(Debug, Clone, PartialEq)]
pub(crate) enum PiRpcOutput {
    Response(PiRpcResponse),
    Event(PiRpcEvent),
}

impl<'de> Deserialize<'de> for PiRpcOutput {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let value = Value::deserialize(deserializer)?;
        let record_type = value
            .get("type")
            .and_then(Value::as_str)
            .ok_or_else(|| D::Error::custom("Pi RPC record must contain a string type"))?;

        if record_type == "response" {
            serde_json::from_value(value)
                .map(Self::Response)
                .map_err(D::Error::custom)
        } else {
            serde_json::from_value(value)
                .map(Self::Event)
                .map_err(D::Error::custom)
        }
    }
}