DS4Server/src/a2ui.rs

mod evaluation;
mod export;
mod validation;

pub(crate) use export::surface_svg;

use evaluation::{evaluate, resolve, resolve_at};
use validation::{validate_component, validate_function};

#[cfg(test)]
use evaluation::format_date;
#[cfg(test)]
use validation::ds4_catalog;

use regex::Regex;
use serde_json::{Map, Value, json};
use std::cell::Cell;
use std::collections::{BTreeMap, BTreeSet};
use std::sync::OnceLock;

pub(crate) const VERSION: &str = "v1.0";
pub(crate) const CATALOG_ID: &str = "https://ds4server.local/a2ui/v1_0/catalog.json";
pub(crate) const BASIC_CATALOG_ID: &str =
    "https://a2ui.org/specification/v1_0/catalogs/basic/catalog.json";
pub(crate) const CATALOG_JSON: &str = include_str!("../assets/a2ui/catalog.json");

thread_local! {
    static TEMPLATE_INDEX: Cell<Option<usize>> = const { Cell::new(None) };
}
pub(crate) const BASIC_COMPONENTS: &[&str] = &[
    "Text",
    "Image",
    "Icon",
    "Video",
    "AudioPlayer",
    "Divider",
    "Row",
    "Column",
    "List",
    "Card",
    "Modal",
    "Tabs",
    "Button",
    "TextField",
    "CheckBox",
    "Slider",
    "DateTimeInput",
    "ChoicePicker",
];
pub(crate) const BASIC_NON_MEDIA_COMPONENTS: &[&str] = &[
    "Text",
    "Image",
    "Icon",
    "Divider",
    "Row",
    "Column",
    "List",
    "Card",
    "Modal",
    "Tabs",
    "Button",
    "TextField",
    "CheckBox",
    "Slider",
    "DateTimeInput",
    "ChoicePicker",
];
const FUNCTIONS: &[&str] = &[
    "required",
    "regex",
    "length",
    "numeric",
    "email",
    "formatString",
    "formatNumber",
    "formatCurrency",
    "formatDate",
    "pluralize",
    "openUrl",
    "and",
    "or",
    "not",
    "@index",
];

pub(crate) const SYSTEM_PROMPT: &str = r#"A2UI is an escalation beyond Markdown, not the default response format. Use it only when the user explicitly asks for an interactive or graphical UI, or when the task materially requires controls, dynamic state, or a visualization that Markdown cannot carry. Use Markdown for web and file search results, research summaries, prose, lists, and tables, even when they contain substantive data. Do not create a surface for progress or status updates, acknowledgements, plans, tool-call narration, or ordinary prose. If unsure, use Markdown. When using A2UI, use the newest A2UI v1.0 protocol. Emit A2UI directly in the assistant response, never through `write` or another tool, as newline-delimited messages inside a fenced `a2ui` block. Keep ordinary prose outside the block. Use catalogId `https://ds4server.local/a2ui/v1_0/catalog.json`.

Every line must be one JSON object with `version":"v1.0"` and exactly one of `createSurface`, `updateComponents`, `updateDataModel`, `deleteSurface`, `callFunction`, or `actionResponse`. Create a surface before updating it. Treat the current client metadata as authoritative: if its surfaces object is empty, prior surfaces in chat history were dismissed and you must create a new surface instead of updating them. Components are a flat adjacency list and the root component has id `root`. Compose complete UIs by combining basic components through container child ids; include every referenced child, tab child, and list template component. Reuse the active surfaceId to update it incrementally; never recreate an existing surface. `createSurface` may include initial `components`, `dataModel`, and `surfaceProperties`. For `actionResponse`, put `actionId` beside `version` and put only `value` or `error` inside `actionResponse`.

Catalog components:
- Basic: Text(text Markdown,variant caption|body), Image(url,description,fit contain|cover|fill|none|scaleDown,variant icon|avatar|smallFeature|mediumFeature|largeFeature|header), Icon(name; safe examples: info|error|check|close|search|settings), Video(url,posterUrl), AudioPlayer(url,description), Divider(axis horizontal|vertical), Row/Column(children,justify start|center|end|spaceBetween|spaceAround|spaceEvenly|stretch,align start|center|end|stretch), List(children,direction vertical|horizontal,align start|center|end|stretch), Card(child), Modal(trigger,content), Tabs(tabs[{title,child}]), Button(child,variant default|primary|borderless,action:{event:{name,context,wantResponse}}), TextField(label,value:{path},variant shortText|longText|number|obscured,placeholder), CheckBox(label,value:{path}), Slider(value:{path},max,min,steps positive integer), DateTimeInput(label,value:{path},enableDate,enableTime,min,max), ChoicePicker(label,options[{label,value}],value:{path},variant multipleSelection|mutuallyExclusive,displayStyle checkbox|chips,filterable). Put numeric weight on direct Row/Column children to distribute available space.
- Research: Chart(title,chartType bar|line|area|stackedBar|pie|donut|heatmap,series[{label,value,segments}]), Table(title,columns,rows), Metric(label,value,detail,trend), Timeline(title,events[{time,title,description,status}]), Map(title,locations[{label,latitude,longitude,detail}]), MindMap(title,nodes[{id,label,children}]), Form(title,children,submitLabel,action). Use pie or donut for proportional breakdowns; donut displays the total in its center. Use heatmap for a matrix: every series entry is a row and its segments are the labeled columns whose numeric values determine color intensity. A heatmap without segments is empty.
- Shared fields: id, accessibility, weight, checks. Checks use {"condition":{"call":"required","args":{...}},"message":"..."}. Bind dynamic values with {"path":"/json/pointer"}. The renderer supports every function in the v1.0 Basic Catalog, including validation, formatting, logic, pluralize, openUrl, and @index. Input edits are local and synchronous. Agent events receive their resolved context and current data model. Server-initiated `callFunction` messages are supported.

Example:
```a2ui
{"version":"v1.0","createSurface":{"surfaceId":"answer","catalogId":"https://ds4server.local/a2ui/v1_0/catalog.json","sendDataModel":true,"surfaceProperties":{},"dataModel":{"name":""},"components":[{"id":"root","component":"Card","child":"body"},{"id":"body","component":"Column","children":["title","name","submit"]},{"id":"title","component":"Text","text":"Research result","variant":"body"},{"id":"name","component":"TextField","label":"Name","value":{"path":"/name"}},{"id":"submit-label","component":"Text","text":"Continue"},{"id":"submit","component":"Button","child":"submit-label","checks":[{"condition":{"call":"required","args":{"value":{"path":"/name"}}},"message":"Name is required"}],"action":{"event":{"name":"continue","wantResponse":true,"responsePath":"/result","context":{"name":{"path":"/name"}}}}}]}}
```
If the renderer reports a `VALIDATION_FAILED` A2UI client error, correct the named message or component with another valid A2UI message."#;

#[derive(Clone, Debug)]
pub(crate) struct Surface {
    pub(crate) id: String,
    pub(crate) catalog_id: String,
    pub(crate) surface_properties: Value,
    pub(crate) send_data_model: bool,
    pub(crate) components: BTreeMap<String, Value>,
    pub(crate) data: Value,
    pub(crate) owner_message_id: i32,
}

#[derive(Clone, Default)]
pub(crate) struct Store {
    surfaces: BTreeMap<String, Surface>,
    surface_order: Vec<String>,
    pending_actions: BTreeMap<String, (String, Option<String>)>,
    next_action_id: u64,
}

#[derive(Debug)]
pub(crate) struct ExtractedLine {
    pub(crate) raw: String,
    pub(crate) value: Result<Value, String>,
}

#[derive(Debug)]
pub(crate) struct Applied {
    pub(crate) raws: Vec<String>,
    pub(crate) reply: Option<Value>,
    pub(crate) open_url: Option<String>,
}

pub(crate) fn replay_epochs<'a>(
    records: impl IntoIterator<Item = (i32, i32, bool, &'a str)>,
) -> (Vec<Store>, Store, Vec<String>) {
    let mut history = Vec::new();
    let mut active = Store::default();
    let mut errors = Vec::new();
    for (record_id, message_id, dismissed, raw) in records {
        if dismissed {
            if active.active_surface().is_some() {
                history.push(active.clone());
            }
            active.clear();
        } else if let Err(error) = active.apply_raw(raw, message_id) {
            errors.push(format!("A2UI message {record_id}: {error}"));
        }
    }
    (history, active, errors)
}

pub(crate) fn validate_surface_composition(surface: &Surface) -> Result<BTreeSet<String>, String> {
    fn visit(
        surface: &Surface,
        id: &str,
        ancestors: &mut BTreeSet<String>,
        visited: &mut BTreeSet<String>,
        kinds: &mut BTreeSet<String>,
    ) -> Result<(), String> {
        if visited.contains(id) {
            return Ok(());
        }
        if !ancestors.insert(id.to_owned()) {
            return Err(format!("cyclic component reference at `{id}`"));
        }
        let component = surface
            .components
            .get(id)
            .and_then(Value::as_object)
            .ok_or_else(|| format!("referenced component `{id}` is missing"))?;
        let kind = required_string(component, "component")?;
        kinds.insert(kind.to_owned());

        let mut children = Vec::new();
        match kind {
            "Row" | "Column" | "List" | "Form" => {
                if let Some(ids) = component.get("children").and_then(Value::as_array) {
                    children.extend(ids.iter().filter_map(Value::as_str));
                } else if let Some(id) = component
                    .get("children")
                    .and_then(Value::as_object)
                    .and_then(|template| template.get("componentId"))
                    .and_then(Value::as_str)
                {
                    children.push(id);
                }
            }
            "Card" | "Button" => {
                children.extend(component.get("child").and_then(Value::as_str));
            }
            "Modal" => {
                children.extend(component.get("trigger").and_then(Value::as_str));
                children.extend(component.get("content").and_then(Value::as_str));
            }
            "Tabs" => {
                children.extend(
                    component
                        .get("tabs")
                        .and_then(Value::as_array)
                        .into_iter()
                        .flatten()
                        .filter_map(|tab| tab.get("child").and_then(Value::as_str)),
                );
            }
            _ => {}
        }
        for child in children {
            visit(surface, child, ancestors, visited, kinds)?;
        }
        ancestors.remove(id);
        visited.insert(id.to_owned());
        Ok(())
    }

    let mut ancestors = BTreeSet::new();
    let mut visited = BTreeSet::new();
    let mut kinds = BTreeSet::new();
    visit(surface, "root", &mut ancestors, &mut visited, &mut kinds)?;
    Ok(kinds)
}

impl Store {
    pub(crate) fn clear(&mut self) {
        self.surfaces.clear();
        self.surface_order.clear();
        self.pending_actions.clear();
        self.next_action_id = 0;
    }

    pub(crate) fn surfaces(&self) -> impl Iterator<Item = &Surface> {
        self.surfaces.values()
    }

    pub(crate) fn surface(&self, id: &str) -> Option<&Surface> {
        self.surfaces.get(id)
    }

    pub(crate) fn active_surface(&self) -> Option<&Surface> {
        self.surface_order
            .last()
            .and_then(|id| self.surfaces.get(id))
    }

    pub(crate) fn image_urls(&self) -> impl Iterator<Item = String> + '_ {
        self.active_surface().into_iter().flat_map(|surface| {
            surface.components.values().filter_map(|component| {
                let component = component.as_object()?;
                let field = match component.get("component").and_then(Value::as_str) {
                    Some("Image") => "url",
                    Some("Video") => "posterUrl",
                    _ => return None,
                };
                component.get(field).map(|value| {
                    display_value(&bound_value_at(Some(value), &surface.data, &surface.data))
                })
            })
        })
    }

    pub(crate) fn apply_raw(
        &mut self,
        raw: &str,
        owner_message_id: i32,
    ) -> Result<Applied, String> {
        let value: Value = serde_json::from_str(raw).map_err(|error| error.to_string())?;
        self.apply(value, raw.to_owned(), owner_message_id)
    }

    pub(crate) fn apply(
        &mut self,
        value: Value,
        raw: String,
        owner_message_id: i32,
    ) -> Result<Applied, String> {
        let mut next = self.clone();
        let applied = next.apply_inner(value, raw, owner_message_id)?;
        *self = next;
        Ok(applied)
    }

    fn apply_inner(
        &mut self,
        value: Value,
        raw: String,
        owner_message_id: i32,
    ) -> Result<Applied, String> {
        let envelope = value
            .as_object()
            .ok_or_else(|| "A2UI message must be a JSON object".to_owned())?;
        let version = envelope
            .get("version")
            .and_then(Value::as_str)
            .ok_or_else(|| "A2UI message requires a string version".to_owned())?;
        if version != VERSION {
            return Err(format!("unsupported A2UI version `{version}`"));
        }
        let kinds = [
            "createSurface",
            "updateComponents",
            "updateDataModel",
            "deleteSurface",
            "callFunction",
            "actionResponse",
        ];
        let present = kinds
            .iter()
            .filter(|key| envelope.contains_key(**key))
            .copied()
            .collect::<Vec<_>>();
        if present.len() != 1 {
            return Err("A2UI envelope must contain version and exactly one message type".into());
        }
        let kind = present[0];
        let allowed_envelope: &[&str] = match kind {
            "callFunction" => &["version", "callFunction", "functionCallId", "wantResponse"],
            "actionResponse" => &["version", "actionResponse", "actionId"],
            _ => &["version", kind],
        };
        reject_unknown(envelope, allowed_envelope, "A2UI envelope")?;
        let payload = object(envelope.get(kind), kind)?;
        let mut raws = vec![raw];
        let mut reply = None;
        let mut open_url = None;
        match kind {
            "createSurface" => {
                reject_unknown(
                    payload,
                    &[
                        "surfaceId",
                        "catalogId",
                        "surfaceProperties",
                        "sendDataModel",
                        "components",
                        "dataModel",
                    ],
                    "createSurface",
                )?;
                let id = required_string(payload, "surfaceId")?;
                if self.surfaces.contains_key(id) {
                    return Err(format!("surface `{id}` already exists"));
                }
                let catalog_id = required_string(payload, "catalogId")?;
                if !matches!(catalog_id, CATALOG_ID | BASIC_CATALOG_ID) {
                    return Err(format!("unsupported catalog `{catalog_id}`"));
                }
                let surface_properties = payload
                    .get("surfaceProperties")
                    .cloned()
                    .unwrap_or_else(|| json!({}));
                if !surface_properties.is_object() {
                    return Err("createSurface surface properties must be an object".into());
                }
                let properties = surface_properties.as_object().unwrap();
                if properties
                    .get("agentDisplayName")
                    .is_some_and(|value| !value.is_string())
                {
                    return Err("surfaceProperties.agentDisplayName must be a string".into());
                }
                if let Some(icon) = properties.get("iconUrl") {
                    let icon = icon.as_str().ok_or_else(|| {
                        "surfaceProperties.iconUrl must be a URI string".to_owned()
                    })?;
                    url::Url::parse(icon)
                        .map_err(|error| format!("invalid surface iconUrl: {error}"))?;
                }
                let send_data_model = optional_bool(payload, "sendDataModel")?.unwrap_or(false);
                self.surfaces.insert(
                    id.to_owned(),
                    Surface {
                        id: id.to_owned(),
                        catalog_id: catalog_id.to_owned(),
                        surface_properties,
                        send_data_model,
                        components: BTreeMap::new(),
                        data: json!({}),
                        owner_message_id,
                    },
                );
                self.surface_order.push(id.to_owned());
                let surface = self.surfaces.get_mut(id).unwrap();
                if let Some(data) = payload.get("dataModel") {
                    if !data.is_object() {
                        return Err("createSurface.dataModel must be an object".into());
                    }
                    surface.data = data.clone();
                }
                if let Some(components) = payload.get("components") {
                    let components = components
                        .as_array()
                        .filter(|components| !components.is_empty())
                        .ok_or_else(|| {
                            "createSurface.components must be a non-empty array".to_owned()
                        })?;
                    for component in components {
                        let component = component
                            .as_object()
                            .ok_or_else(|| "every component must be an object".to_owned())?;
                        validate_component(component, &surface.catalog_id)?;
                        surface.components.insert(
                            required_string(component, "id")?.to_owned(),
                            Value::Object(component.clone()),
                        );
                    }
                }
            }
            "updateComponents" => {
                reject_unknown(payload, &["surfaceId", "components"], "updateComponents")?;
                let id = required_string(payload, "surfaceId")?;
                let surface = self
                    .surfaces
                    .get_mut(id)
                    .ok_or_else(|| format!("surface `{id}` has not been created"))?;
                let components = payload
                    .get("components")
                    .and_then(Value::as_array)
                    .ok_or_else(|| "updateComponents.components must be an array".to_owned())?;
                if components.is_empty() {
                    return Err("updateComponents.components cannot be empty".into());
                }
                let mut ids = BTreeSet::new();
                for component in components {
                    let component = component
                        .as_object()
                        .ok_or_else(|| "every component must be an object".to_owned())?;
                    validate_component(component, &surface.catalog_id)?;
                    let component_id = required_string(component, "id")?;
                    if !ids.insert(component_id) {
                        return Err(format!(
                            "component `{component_id}` occurs twice in one update"
                        ));
                    }
                }
                for component in components {
                    let component_id = component["id"].as_str().unwrap().to_owned();
                    surface.components.insert(component_id, component.clone());
                }
            }
            "updateDataModel" => {
                reject_unknown(payload, &["surfaceId", "path", "value"], "updateDataModel")?;
                let id = required_string(payload, "surfaceId")?;
                let surface = self
                    .surfaces
                    .get_mut(id)
                    .ok_or_else(|| format!("surface `{id}` has not been created"))?;
                let path = payload.get("path").and_then(Value::as_str).unwrap_or("/");
                if path != "/" && !path.starts_with('/') {
                    return Err("updateDataModel.path must be a JSON Pointer".into());
                }
                set_pointer(&mut surface.data, path, payload.get("value").cloned())?;
            }
            "deleteSurface" => {
                reject_unknown(payload, &["surfaceId"], "deleteSurface")?;
                let id = required_string(payload, "surfaceId")?;
                if self.surfaces.remove(id).is_none() {
                    return Err(format!("surface `{id}` has not been created"));
                }
                self.surface_order.retain(|surface_id| surface_id != id);
            }
            "callFunction" => {
                let call_id = required_string(envelope, "functionCallId")?;
                reject_unknown(payload, &["call", "args"], "callFunction")?;
                let call = required_string(payload, "call")?;
                validate_function(&Value::Object(payload.clone()))?;
                let value = if call == "openUrl" {
                    let url = payload
                        .get("args")
                        .and_then(|args| args.get("url"))
                        .and_then(Value::as_str)
                        .ok_or_else(|| "openUrl requires string `args.url`".to_owned())?;
                    let parsed = url::Url::parse(url).map_err(|error| error.to_string())?;
                    if !matches!(parsed.scheme(), "http" | "https") {
                        return Err("openUrl only supports HTTP(S) URLs".into());
                    }
                    open_url = Some(parsed.to_string());
                    Value::Null
                } else {
                    evaluate(
                        call,
                        payload.get("args").unwrap_or(&json!({})),
                        &json!({}),
                        &json!({}),
                    )?
                };
                if optional_bool(envelope, "wantResponse")?.unwrap_or(false) {
                    reply = Some(json!({
                        "version": VERSION,
                        "functionResponse": {
                            "functionCallId": call_id,
                            "call": call,
                            "value": value
                        }
                    }));
                }
            }
            "actionResponse" => {
                let action_id = required_string(envelope, "actionId")?;
                let response = payload
                    .get("value")
                    .or_else(|| payload.get("error"))
                    .cloned()
                    .ok_or_else(|| "actionResponse requires value or error".to_owned())?;
                reject_unknown(payload, &["value", "error"], "actionResponse")?;
                if payload.contains_key("value") == payload.contains_key("error") {
                    return Err("actionResponse requires exactly one of value or error".into());
                }
                if let Some(error) = payload.get("error") {
                    let error = object(Some(error), "actionResponse.error")?;
                    reject_unknown(error, &["code", "message"], "actionResponse.error")?;
                    required_string(error, "code")?;
                    required_string(error, "message")?;
                }
                if let Some((surface_id, response_path)) = self.pending_actions.remove(action_id)
                    && payload.contains_key("value")
                    && let Some(path) = response_path
                {
                    let surface = self.surfaces.get_mut(&surface_id).ok_or_else(|| {
                        format!("surface `{surface_id}` for action `{action_id}` no longer exists")
                    })?;
                    set_pointer(&mut surface.data, &path, Some(response.clone()))?;
                    raws.push(
                        serde_json::to_string(&json!({
                            "version": VERSION,
                            "updateDataModel": {
                                "surfaceId": surface_id,
                                "path": path,
                                "value": response
                            }
                        }))
                        .map_err(|error| error.to_string())?,
                    );
                }
            }
            _ => unreachable!(),
        }
        Ok(Applied {
            raws,
            reply,
            open_url,
        })
    }

    pub(crate) fn local_update(
        &mut self,
        surface_id: &str,
        path: &str,
        value: Value,
    ) -> Result<String, String> {
        let surface = self
            .surfaces
            .get_mut(surface_id)
            .ok_or_else(|| format!("surface `{surface_id}` does not exist"))?;
        set_pointer(&mut surface.data, path, Some(value.clone()))?;
        serde_json::to_string(&json!({
            "version": VERSION,
            "updateDataModel": {"surfaceId": surface_id, "path": path, "value": value}
        }))
        .map_err(|error| error.to_string())
    }

    pub(crate) fn action(
        &mut self,
        surface_id: &str,
        component_id: &str,
        context_path: Option<&str>,
    ) -> Result<Value, String> {
        let surface = self
            .surfaces
            .get(surface_id)
            .ok_or_else(|| format!("surface `{surface_id}` does not exist"))?;
        let component = surface
            .components
            .get(component_id)
            .and_then(Value::as_object)
            .ok_or_else(|| format!("component `{component_id}` does not exist"))?;
        if let Some(error) = first_failed_check(component, &surface.data) {
            return Err(error);
        }
        let event = component
            .get("action")
            .and_then(Value::as_object)
            .and_then(|action| action.get("event"))
            .and_then(Value::as_object)
            .ok_or_else(|| format!("component `{component_id}` has no agent event"))?;
        let name = required_string(event, "name")?;
        let context_data = context_path
            .and_then(|path| surface.data.pointer(path))
            .unwrap_or(&surface.data);
        let context = resolve_at(
            event.get("context").unwrap_or(&Value::Object(Map::new())),
            &surface.data,
            context_data,
        )?;
        self.next_action_id += 1;
        let action_id = format!("ds4-action-{}", self.next_action_id);
        let want_response = event
            .get("wantResponse")
            .and_then(Value::as_bool)
            .unwrap_or(false);
        let response_path = event
            .get("responsePath")
            .map(|value| {
                value
                    .as_str()
                    .filter(|path| path.starts_with('/'))
                    .map(str::to_owned)
                    .ok_or_else(|| "action responsePath must be a JSON Pointer".to_owned())
            })
            .transpose()?;
        let timestamp = time::OffsetDateTime::now_utc()
            .format(&time::format_description::well_known::Rfc3339)
            .map_err(|error| error.to_string())?;
        let mut event_payload = json!({
            "name": name,
            "surfaceId": surface_id,
            "sourceComponentId": component_id,
            "timestamp": timestamp,
            "context": context,
            "wantResponse": want_response
        });
        if want_response {
            event_payload["actionId"] = Value::String(action_id.clone());
            self.pending_actions
                .insert(action_id, (surface_id.to_owned(), response_path));
        }
        let action = json!({
            "version": VERSION,
            "action": event_payload
        });
        Ok(action)
    }

    pub(crate) fn client_metadata(&self) -> Value {
        let surfaces = self
            .active_surface()
            .into_iter()
            .filter(|surface| surface.send_data_model)
            .map(|surface| (surface.id.clone(), surface.data.clone()))
            .collect::<Map<_, _>>();
        let mut metadata = json!({
            "a2uiClientCapabilities": {
                (VERSION): {
                    "supportedCatalogIds": [CATALOG_ID, BASIC_CATALOG_ID],
                    "inlineCatalogs": []
                }
            }
        });
        metadata["a2uiClientDataModel"] = json!({
            "version": VERSION,
            "surfaces": surfaces,
        });
        metadata
    }
}

pub(crate) fn extract_lines(content: &str) -> Vec<ExtractedLine> {
    let mut lines = Vec::new();
    let mut in_block = false;
    for part in content.split_inclusive('\n') {
        let complete = part.ends_with('\n');
        let line = part.trim();
        if !in_block && line.eq_ignore_ascii_case("```a2ui") {
            in_block = true;
        } else if in_block && line.starts_with("```") {
            in_block = false;
        } else if in_block && complete && !line.is_empty() {
            lines.push(ExtractedLine {
                raw: line.to_owned(),
                value: serde_json::from_str(line).map_err(|error| error.to_string()),
            });
        }
    }
    lines
}

pub(crate) fn message_surface_id(value: &Value) -> Option<&str> {
    value
        .as_object()?
        .values()
        .filter_map(Value::as_object)
        .find_map(|payload| payload.get("surfaceId").and_then(Value::as_str))
}

pub(crate) fn transcript_fallback(content: &str) -> String {
    let mut output = String::new();
    let mut in_block = false;
    let mut fallback = Vec::new();
    for part in content.split_inclusive('\n') {
        let line = part.trim();
        if !in_block && line.eq_ignore_ascii_case("```a2ui") {
            in_block = true;
            continue;
        }
        if in_block && line.starts_with("```") {
            in_block = false;
            if !fallback.is_empty() {
                output.push_str("\n> ");
                output.push_str(&fallback.join("  \n> "));
                output.push('\n');
                fallback.clear();
            }
            continue;
        }
        if in_block {
            if line.is_empty() {
                continue;
            }
            fallback.push(match serde_json::from_str::<Value>(line) {
                Ok(value) => message_summary(&value),
                Err(error) => format!("A2UI message could not be read: {error}"),
            });
        } else {
            output.push_str(part);
        }
    }
    if !fallback.is_empty() {
        output.push_str("\n> ");
        output.push_str(&fallback.join("  \n> "));
    }
    output
}

pub(crate) fn bound_value_at(value: Option<&Value>, data: &Value, context: &Value) -> Value {
    value
        .map(|value| resolve_at(value, data, context).unwrap_or(Value::Null))
        .unwrap_or(Value::Null)
}

pub(crate) fn binding_path(value: Option<&Value>) -> Option<&str> {
    value?.as_object()?.get("path")?.as_str()
}

pub(crate) fn with_template_index<T>(index: usize, render: impl FnOnce() -> T) -> T {
    TEMPLATE_INDEX.with(|current| {
        let previous = current.replace(Some(index));
        let rendered = render();
        current.set(previous);
        rendered
    })
}

pub(crate) fn first_failed_check(component: &Map<String, Value>, data: &Value) -> Option<String> {
    component
        .get("checks")
        .and_then(Value::as_array)
        .and_then(|checks| {
            checks.iter().find_map(|check| {
                let check = check.as_object()?;
                let fallback = Value::Object(check.clone());
                let condition = check.get("condition").unwrap_or(&fallback);
                match resolve(condition, data) {
                    Ok(Value::Bool(true)) => None,
                    Ok(_) => Some(
                        check
                            .get("message")
                            .and_then(Value::as_str)
                            .unwrap_or("This value is invalid")
                            .to_owned(),
                    ),
                    Err(error) => Some(error),
                }
            })
        })
}

fn set_pointer(root: &mut Value, path: &str, value: Option<Value>) -> Result<(), String> {
    if path.is_empty() || path == "/" {
        *root = value.unwrap_or(Value::Null);
        return Ok(());
    }
    let segments = path[1..]
        .split('/')
        .map(|segment| segment.replace("~1", "/").replace("~0", "~"))
        .collect::<Vec<_>>();
    let (last, parents) = segments.split_last().unwrap();
    let mut current = root;
    for segment in parents {
        match current {
            Value::Object(object) => {
                current = object
                    .entry(segment.clone())
                    .or_insert_with(|| Value::Object(Map::new()));
            }
            Value::Array(array) => {
                let index = segment
                    .parse::<usize>()
                    .map_err(|_| format!("`{segment}` is not an array index"))?;
                while array.len() <= index {
                    array.push(Value::Object(Map::new()));
                }
                current = &mut array[index];
            }
            _ => return Err(format!("cannot traverse through `{segment}`")),
        }
    }
    match current {
        Value::Object(object) => match value {
            Some(value) => {
                object.insert(last.clone(), value);
            }
            None => {
                object.remove(last);
            }
        },
        Value::Array(array) => {
            let index = last
                .parse::<usize>()
                .map_err(|_| format!("`{last}` is not an array index"))?;
            if let Some(value) = value {
                while array.len() <= index {
                    array.push(Value::Null);
                }
                array[index] = value;
            } else if index < array.len() {
                array.remove(index);
            }
        }
        _ => return Err(format!("cannot update `{path}`")),
    }
    Ok(())
}

fn message_summary(value: &Value) -> String {
    for (key, verb) in [
        ("createSurface", "Created"),
        ("updateComponents", "Updated components on"),
        ("updateDataModel", "Updated data on"),
        ("deleteSurface", "Deleted"),
        ("actionResponse", "Responded to action on"),
    ] {
        if let Some(id) = value
            .get(key)
            .and_then(|value| value.get("surfaceId"))
            .and_then(Value::as_str)
        {
            return format!("{verb} interactive surface `{id}`.");
        }
    }
    "A2UI interactive update.".into()
}

pub(crate) fn display_value(value: &Value) -> String {
    match value {
        Value::Null => String::new(),
        Value::String(value) => value.clone(),
        Value::Bool(value) => value.to_string(),
        Value::Number(value) => value.to_string(),
        value => serde_json::to_string(value).unwrap_or_default(),
    }
}

fn normalize_pointer(path: &str) -> &str {
    if path == "/" { "" } else { path }
}

fn object<'a>(value: Option<&'a Value>, name: &str) -> Result<&'a Map<String, Value>, String> {
    value
        .and_then(Value::as_object)
        .ok_or_else(|| format!("{name} must be an object"))
}

fn reject_unknown(object: &Map<String, Value>, allowed: &[&str], name: &str) -> Result<(), String> {
    if let Some(key) = object.keys().find(|key| !allowed.contains(&key.as_str())) {
        return Err(format!("{name} contains unknown property `{key}`"));
    }
    Ok(())
}

fn required_string<'a>(object: &'a Map<String, Value>, key: &str) -> Result<&'a str, String> {
    object
        .get(key)
        .and_then(Value::as_str)
        .filter(|value| !value.is_empty())
        .ok_or_else(|| format!("`{key}` must be a non-empty string"))
}

fn optional_bool(object: &Map<String, Value>, key: &str) -> Result<Option<bool>, String> {
    object
        .get(key)
        .map(|value| {
            value
                .as_bool()
                .ok_or_else(|| format!("`{key}` must be a boolean"))
        })
        .transpose()
}

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

    fn apply(store: &mut Store, value: Value) -> Result<Applied, String> {
        store.apply(value.clone(), value.to_string(), 7)
    }

    #[test]
    fn lifecycle_updates_incrementally_and_json_pointer_upserts() {
        let mut store = Store::default();
        apply(
            &mut store,
            json!({"version":VERSION,"createSurface":{"surfaceId":"s","catalogId":CATALOG_ID}}),
        )
        .unwrap();
        apply(&mut store, json!({"version":VERSION,"updateComponents":{"surfaceId":"s","components":[{"id":"root","component":"Text","text":{"path":"/user/name"}}]}})).unwrap();
        apply(&mut store, json!({"version":VERSION,"updateDataModel":{"surfaceId":"s","path":"/user/name","value":"Ada"}})).unwrap();
        assert_eq!(
            store.surface("s").unwrap().data.pointer("/user/name"),
            Some(&json!("Ada"))
        );
        apply(&mut store, json!({"version":VERSION,"updateComponents":{"surfaceId":"s","components":[{"id":"root","component":"Metric","label":"Name","value":{"path":"/user/name"}}]}})).unwrap();
        assert_eq!(
            store.surface("s").unwrap().components["root"]["component"],
            "Metric"
        );
        apply(&mut store, json!({"version":VERSION,"updateDataModel":{"surfaceId":"s","path":"/items","value":["first","second"]}})).unwrap();
        apply(
            &mut store,
            json!({"version":VERSION,"updateDataModel":{"surfaceId":"s","path":"/items/0"}}),
        )
        .unwrap();
        assert_eq!(store.surface("s").unwrap().data["items"], json!(["second"]));
        apply(
            &mut store,
            json!({"version":VERSION,"deleteSurface":{"surfaceId":"s"}}),
        )
        .unwrap();
        assert!(store.surface("s").is_none());
    }

    #[test]
    fn latest_created_surface_is_the_active_surface() {
        let mut store = Store::default();
        for id in ["first", "second"] {
            apply(
                &mut store,
                json!({"version":VERSION,"createSurface":{"surfaceId":id,"catalogId":CATALOG_ID,"sendDataModel":true,"dataModel":{"id":id}}}),
            )
            .unwrap();
        }
        assert_eq!(store.active_surface().unwrap().id, "second");
        let metadata = store.client_metadata();
        assert_eq!(
            metadata.pointer("/a2uiClientDataModel/surfaces/second/id"),
            Some(&json!("second"))
        );
        assert!(
            metadata
                .pointer("/a2uiClientDataModel/surfaces/first")
                .is_none()
        );
        apply(
            &mut store,
            json!({"version":VERSION,"deleteSurface":{"surfaceId":"second"}}),
        )
        .unwrap();
        assert_eq!(store.active_surface().unwrap().id, "first");
    }

    #[test]
    fn dismissal_boundaries_replay_history_and_a_fresh_active_epoch() {
        let first = json!({"version":VERSION,"createSurface":{"surfaceId":"first","catalogId":CATALOG_ID,"components":[{"id":"root","component":"Text","text":"First"}]}}).to_string();
        let update = json!({"version":VERSION,"updateComponents":{"surfaceId":"first","components":[{"id":"root","component":"Text","text":"First final"}]}}).to_string();
        let second = json!({"version":VERSION,"createSurface":{"surfaceId":"second","catalogId":CATALOG_ID,"components":[{"id":"root","component":"Text","text":"Second"}]}}).to_string();
        let records = [
            (1, 10, false, first.as_str()),
            (2, 11, false, update.as_str()),
            (3, 12, true, "{}"),
            (4, 13, false, second.as_str()),
        ];
        let (history, active, errors) = replay_epochs(records);
        assert!(errors.is_empty());
        assert_eq!(history.len(), 1);
        assert_eq!(history[0].active_surface().unwrap().id, "first");
        assert_eq!(
            history[0].active_surface().unwrap().components["root"]["text"],
            "First final"
        );
        assert_eq!(active.active_surface().unwrap().id, "second");
    }

    #[test]
    fn catalog_validation_and_actions_use_current_local_data() {
        let mut store = Store::default();
        apply(&mut store, json!({"version":VERSION,"createSurface":{"surfaceId":"s","catalogId":CATALOG_ID,"sendDataModel":true}})).unwrap();
        apply(&mut store, json!({"version":VERSION,"updateComponents":{"surfaceId":"s","components":[{"id":"label","component":"Text","text":"Send"},{"id":"root","component":"Button","child":"label","checks":[{"condition":{"call":"required","args":{"value":{"path":"/name"}}},"message":"Name required"}],"action":{"event":{"name":"submit","context":{"name":{"path":"/name"}}}}}]}})).unwrap();
        assert_eq!(
            store.action("s", "root", None).unwrap_err(),
            "Name required"
        );
        let raw = store.local_update("s", "/name", json!("Ada")).unwrap();
        assert!(raw.contains("updateDataModel"));
        let action = store.action("s", "root", None).unwrap();
        assert_eq!(action.pointer("/action/context/name"), Some(&json!("Ada")));
        let metadata = store.client_metadata();
        assert_eq!(
            metadata.pointer("/a2uiClientDataModel/surfaces/s/name"),
            Some(&json!("Ada"))
        );
    }

    #[test]
    fn streaming_parser_only_accepts_complete_a2ui_jsonl_lines() {
        let content = "Before\n```a2ui\n{\"version\":\"v1.0\",\"deleteSurface\":{\"surfaceId\":\"s\"}}\n{\"version\":";
        let lines = extract_lines(content);
        assert_eq!(lines.len(), 1);
        assert!(lines[0].value.is_ok());
        assert!(transcript_fallback(content).contains("Deleted interactive surface `s`"));
    }

    #[test]
    fn v1_action_and_function_round_trips_follow_current_schema() {
        let mut store = Store::default();
        apply(
            &mut store,
            json!({
                "version": VERSION,
                "createSurface": {
                    "surfaceId": "s",
                    "catalogId": CATALOG_ID,
                    "dataModel": {"result": null},
                    "components": [
                        {"id":"label","component":"Text","text":"Run"},
                        {"id":"root","component":"Button","child":"label","action":{"event":{
                            "name":"run","wantResponse":true,"responsePath":"/result"
                        }}}
                    ]
                }
            }),
        )
        .unwrap();
        let action = store.action("s", "root", None).unwrap();
        assert!(action.pointer("/action/timestamp").is_some());
        let action_id = action["action"]["actionId"].as_str().unwrap();
        let response = apply(
            &mut store,
            json!({
                "version": VERSION,
                "actionId": action_id,
                "actionResponse": {"value": "done"}
            }),
        )
        .unwrap();
        assert_eq!(response.raws.len(), 2);
        assert_eq!(store.surface("s").unwrap().data["result"], "done");

        let applied = apply(
            &mut store,
            json!({
                "version": VERSION,
                "functionCallId": "f1",
                "wantResponse": true,
                "callFunction": {"call":"formatNumber","args":{"value":1234.5,"decimals":1}}
            }),
        )
        .unwrap();
        assert_eq!(
            applied.reply.unwrap().pointer("/functionResponse/value"),
            Some(&json!("1,234.5"))
        );
    }

    #[test]
    fn embedded_catalog_and_validator_cover_the_same_components() {
        let catalog = ds4_catalog();
        assert_eq!(catalog["catalogId"], CATALOG_ID);
        let components = catalog["components"].as_object().unwrap();
        for component in BASIC_COMPONENTS {
            assert!(components.contains_key(*component));
        }
        for component in [
            "Chart", "Table", "Metric", "Timeline", "Map", "MindMap", "Form",
        ] {
            assert!(components.contains_key(component));
        }
        let functions = catalog["functions"].as_object().unwrap();
        for function in FUNCTIONS.iter().filter(|function| **function != "@index") {
            assert!(functions.contains_key(*function));
        }
    }

    #[test]
    fn all_basics_compose_into_one_surface() {
        let mut store = Store::default();
        apply(
            &mut store,
            json!({
                "version": VERSION,
                "createSurface": {
                    "surfaceId": "composed",
                    "catalogId": CATALOG_ID,
                    "sendDataModel": true,
                    "dataModel": {
                        "name": "Ada",
                        "active": true,
                        "priority": 3,
                        "due": "2026-08-01",
                        "language": ["rust"],
                        "items": [{"label": "Renderer"}, {"label": "Validator"}]
                    },
                    "components": [
                        {"id":"root","component":"Card","child":"layout"},
                        {"id":"layout","component":"Column","children":["title","identity","divider","controls","tabs","items","media","modal"]},
                        {"id":"title","component":"Text","text":"**Project dashboard**"},
                        {"id":"identity","component":"Row","children":["icon","avatar"],"align":"center"},
                        {"id":"icon","component":"Icon","name":"accountCircle"},
                        {"id":"avatar","component":"Image","url":"https://example.com/avatar.png","description":"Project owner","variant":"avatar"},
                        {"id":"divider","component":"Divider","axis":"horizontal"},
                        {"id":"controls","component":"Row","children":["name","active","priority","due","language"]},
                        {"id":"name","component":"TextField","label":"Name","value":{"path":"/name"},"weight":2},
                        {"id":"active","component":"CheckBox","label":"Active","value":{"path":"/active"}},
                        {"id":"priority","component":"Slider","label":"Priority","value":{"path":"/priority"},"min":1,"max":5},
                        {"id":"due","component":"DateTimeInput","label":"Due","value":{"path":"/due"},"enableDate":true},
                        {"id":"language","component":"ChoicePicker","label":"Language","options":[{"label":"Rust","value":"rust"},{"label":"Python","value":"python"}],"value":{"path":"/language"},"variant":"mutuallyExclusive","displayStyle":"chips","filterable":true},
                        {"id":"tabs","component":"Tabs","tabs":[{"title":"Overview","child":"chart"},{"title":"Details","child":"details"}]},
                        {"id":"chart","component":"Chart","chartType":"donut","series":[{"label":"Done","value":2},{"label":"Open","value":1}]},
                        {"id":"details","component":"Card","child":"details-text"},
                        {"id":"details-text","component":"Text","text":"All systems operational."},
                        {"id":"items","component":"List","children":{"path":"/items","componentId":"item"}},
                        {"id":"item","component":"Text","text":{"path":"label"}},
                        {"id":"media","component":"Row","children":["video","audio"]},
                        {"id":"video","component":"Video","url":"https://example.com/demo.mp4","posterUrl":"https://example.com/poster.jpg"},
                        {"id":"audio","component":"AudioPlayer","url":"https://example.com/demo.mp3","description":"Project update"},
                        {"id":"modal","component":"Modal","trigger":"open","content":"dialog"},
                        {"id":"open","component":"Button","child":"open-label","variant":"primary","action":{"event":{"name":"openDetails","context":{"name":{"path":"/name"}}}}},
                        {"id":"open-label","component":"Text","text":"Open details"},
                        {"id":"dialog","component":"Card","child":"dialog-body"},
                        {"id":"dialog-body","component":"Text","text":"Composed modal content"}
                    ]
                }
            }),
        )
        .unwrap();

        let surface = store.surface("composed").unwrap();
        let kinds = validate_surface_composition(surface).unwrap();
        for component in BASIC_COMPONENTS {
            assert!(kinds.contains(*component), "missing {component}");
        }
        assert!(kinds.contains("Chart"));
        assert!(
            store
                .image_urls()
                .any(|url| url == "https://example.com/poster.jpg")
        );

        let mut broken = surface.clone();
        broken.components.insert(
            "root".into(),
            json!({"id":"root","component":"Card","child":"missing"}),
        );
        assert_eq!(
            validate_surface_composition(&broken).unwrap_err(),
            "referenced component `missing` is missing"
        );
        broken.components.insert(
            "root".into(),
            json!({"id":"root","component":"Card","child":"root"}),
        );
        assert_eq!(
            validate_surface_composition(&broken).unwrap_err(),
            "cyclic component reference at `root`"
        );
    }

    #[test]
    fn catalog_validation_rejects_invalid_variants_and_function_arguments() {
        let mut store = Store::default();
        apply(
            &mut store,
            json!({"version":VERSION,"createSurface":{"surfaceId":"s","catalogId":CATALOG_ID}}),
        )
        .unwrap();
        assert_eq!(
            apply(&mut store, json!({"version":VERSION,"updateComponents":{"surfaceId":"s","components":[{"id":"root","component":"Text","text":"x","variant":"h1"}]}})).unwrap_err(),
            "component `root` has invalid variant `h1`; expected one of: caption, body"
        );
        assert!(apply(&mut store, json!({"version":VERSION,"updateComponents":{"surfaceId":"s","components":[{"id":"root","component":"Text","text":{"call":"formatDate","args":{"value":"2026-01-01T00:00:00Z"}}}]}})).is_err());
    }

    #[test]
    fn catalog_accepts_pie_donut_and_heatmap_charts() {
        let mut store = Store::default();
        apply(
            &mut store,
            json!({"version":VERSION,"createSurface":{"surfaceId":"s","catalogId":CATALOG_ID}}),
        )
        .unwrap();
        for chart_type in ["pie", "donut", "heatmap"] {
            apply(
                &mut store,
                json!({"version":VERSION,"updateComponents":{"surfaceId":"s","components":[{"id":"root","component":"Chart","chartType":chart_type,"series":[{"label":"rs","value":41}]}]}}),
            )
            .unwrap();
        }
    }

    #[test]
    fn basic_formatters_follow_catalog_examples() {
        let date = time::OffsetDateTime::parse(
            "2026-03-16T14:30:00Z",
            &time::format_description::well_known::Rfc3339,
        )
        .unwrap();
        assert_eq!(format_date(date, "EEEE, d MMMM"), "Monday, 16 March");
        let data = json!({"name":"Ada","currentDate":"2026-03-16T14:30:00Z"});
        assert_eq!(
            evaluate(
                "formatString",
                &json!({"value":"Hello ${/name}: ${formatDate(value:${/currentDate}, format:'MMM dd, yyyy')} \\${literal}"}),
                &data,
                &data,
            )
            .unwrap(),
            json!("Hello Ada: Mar 16, 2026 ${literal}")
        );
    }
}