codex-mobile-bridge 0.2.6

use std::collections::HashMap;

use serde_json::{json, Value};

use super::metadata::{TimelineSemanticInfo, TimelineSummaryInfo};
use crate::bridge_protocol::TimelineEntry;

pub(super) fn finalize_timeline_entries(entries: &mut [TimelineEntry]) {
    let mut semantic_by_call: HashMap<
        (String, String, String),
        (TimelineSemanticInfo, TimelineSummaryInfo),
    > = HashMap::new();

    for entry in entries.iter() {
        let Some(call_id) = metadata_call_id(&entry.metadata) else {
            continue;
        };
        let Some(semantic) = metadata_semantic(&entry.metadata) else {
            continue;
        };
        if semantic.role == "output" {
            continue;
        }
        let summary = metadata_summary(&entry.metadata).unwrap_or_default();
        semantic_by_call.insert(
            (
                entry.thread_id.clone(),
                turn_lookup_key(entry.turn_id.as_deref()),
                call_id,
            ),
            (semantic, summary),
        );
    }

    for entry in entries.iter_mut() {
        if !is_output_entry_type(&entry.entry_type) {
            continue;
        }
        if metadata_semantic(&entry.metadata).is_some() {
            continue;
        }
        let Some(call_id) = metadata_call_id(&entry.metadata) else {
            continue;
        };
        let lookup_key = (
            entry.thread_id.clone(),
            turn_lookup_key(entry.turn_id.as_deref()),
            call_id,
        );
        let Some((semantic, summary)) = semantic_by_call.get(&lookup_key) else {
            continue;
        };
        let mut inherited_semantic = semantic.clone();
        inherited_semantic.role = "output".to_string();
        set_metadata_semantic(entry, &inherited_semantic, Some(summary));
    }
}

pub(super) fn merge_timeline_metadata(existing: &mut Value, incoming: Value) {
    if existing.is_null() || !existing.is_object() || !incoming.is_object() {
        *existing = incoming;
        return;
    }

    let Some(existing_object) = existing.as_object_mut() else {
        *existing = incoming;
        return;
    };
    let Some(incoming_object) = incoming.as_object() else {
        *existing = incoming;
        return;
    };

    for (key, value) in incoming_object {
        match key.as_str() {
            "payload" => merge_timeline_payload(existing_object, value),
            "stream" | "lifecycle" | "semantic" | "summary" => {
                merge_timeline_object(existing_object, key, value)
            }
            _ => {
                existing_object.insert(key.clone(), value.clone());
            }
        }
    }
}

fn is_output_entry_type(entry_type: &str) -> bool {
    matches!(
        entry_type,
        "function_call_output"
            | "custom_tool_call_output"
            | "shell_call_output"
            | "local_shell_call_output"
            | "apply_patch_call_output"
            | "computer_call_output"
    )
}

fn turn_lookup_key(turn_id: Option<&str>) -> String {
    turn_id.unwrap_or_default().to_string()
}

fn metadata_call_id(metadata: &Value) -> Option<String> {
    metadata
        .get("payload")
        .and_then(|payload| payload.get("callId"))
        .and_then(Value::as_str)
        .map(ToOwned::to_owned)
}

fn metadata_semantic(metadata: &Value) -> Option<TimelineSemanticInfo> {
    TimelineSemanticInfo::from_value(metadata.get("semantic")?)
}

fn metadata_summary(metadata: &Value) -> Option<TimelineSummaryInfo> {
    TimelineSummaryInfo::from_value(metadata.get("summary")?)
}

fn set_metadata_semantic(
    entry: &mut TimelineEntry,
    semantic: &TimelineSemanticInfo,
    summary: Option<&TimelineSummaryInfo>,
) {
    if entry.metadata.is_null() || !entry.metadata.is_object() {
        entry.metadata = json!({});
    }
    let Some(metadata_object) = entry.metadata.as_object_mut() else {
        return;
    };
    metadata_object.insert("semantic".to_string(), semantic.to_value());
    if let Some(summary) = summary {
        metadata_object.insert("summary".to_string(), summary.to_value());
    }
}

fn merge_timeline_payload(existing_object: &mut serde_json::Map<String, Value>, incoming: &Value) {
    let payload = existing_object
        .entry("payload".to_string())
        .or_insert_with(|| json!({}));
    if payload.is_null() || !payload.is_object() || !incoming.is_object() {
        *payload = incoming.clone();
        return;
    }

    let Some(existing_payload) = payload.as_object_mut() else {
        *payload = incoming.clone();
        return;
    };
    let Some(incoming_payload) = incoming.as_object() else {
        *payload = incoming.clone();
        return;
    };

    for (key, value) in incoming_payload {
        if matches!(key.as_str(), "summary" | "content")
            && value.as_array().is_some_and(|array| array.is_empty())
            && existing_payload
                .get(key)
                .and_then(Value::as_array)
                .is_some_and(|array| !array.is_empty())
        {
            continue;
        }
        existing_payload.insert(key.clone(), value.clone());
    }
}

fn merge_timeline_object(
    existing_object: &mut serde_json::Map<String, Value>,
    key: &str,
    incoming: &Value,
) {
    if incoming.is_null() {
        return;
    }
    let target = existing_object
        .entry(key.to_string())
        .or_insert_with(|| json!({}));
    if target.is_null() || !target.is_object() || !incoming.is_object() {
        *target = incoming.clone();
        return;
    }

    let Some(target_object) = target.as_object_mut() else {
        *target = incoming.clone();
        return;
    };
    let Some(incoming_object) = incoming.as_object() else {
        *target = incoming.clone();
        return;
    };

    for (child_key, child_value) in incoming_object {
        target_object.insert(child_key.clone(), child_value.clone());
    }
}