1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
use std::sync::{
    Arc,
    Mutex,
};

use accesskit::{
    Action,
    Affine,
    Node,
    NodeBuilder,
    NodeId as AccessibilityId,
    Rect,
    Role,
    TextDirection,
    Tree,
    TreeUpdate,
};
use freya_common::AccessibilityDirtyNodes;
use freya_engine::prelude::{
    Color,
    Slant,
    TextAlign,
    TextDecoration,
    TextDecorationStyle,
};
use freya_native_core::{
    node::NodeType,
    prelude::NodeImmutable,
    tags::TagName,
    NodeId,
};
use freya_node_state::{
    AccessibilityNodeState,
    Fill,
    FontStyleState,
    OverflowMode,
    StyleState,
    TransformState,
};
use rustc_hash::{
    FxHashMap,
    FxHashSet,
};
use torin::{
    prelude::LayoutNode,
    torin::Torin,
};

use super::{
    AccessibilityFocusStrategy,
    NodeAccessibility,
};
use crate::dom::{
    DioxusDOM,
    DioxusNode,
};

pub const ACCESSIBILITY_ROOT_ID: AccessibilityId = AccessibilityId(0);

pub type SharedAccessibilityTree = Arc<Mutex<AccessibilityTree>>;

pub struct AccessibilityTree {
    pub map: FxHashMap<AccessibilityId, NodeId>,
    // Current focused Accessibility Node.
    pub focused_id: AccessibilityId,
}

impl AccessibilityTree {
    pub fn new(focused_id: AccessibilityId) -> Self {
        Self {
            focused_id,
            map: FxHashMap::default(),
        }
    }

    pub fn focused_node_id(&self) -> Option<NodeId> {
        self.map.get(&self.focused_id).cloned()
    }

    /// Initialize the Accessibility Tree
    pub fn init(
        &self,
        rdom: &DioxusDOM,
        layout: &Torin<NodeId>,
        dirty_nodes: &mut AccessibilityDirtyNodes,
    ) -> TreeUpdate {
        dirty_nodes.clear();

        let mut nodes = vec![];

        rdom.traverse_depth_first_advanced(|node_ref| {
            if !node_ref.node_type().is_element() {
                return false;
            }

            let accessibility_id = node_ref.get_accessibility_id();
            let layout_node = layout.get(node_ref.id());

            // Layout nodes might not exist yet when the app is lauched
            if let Some((accessibility_id, layout_node)) = accessibility_id.zip(layout_node) {
                let node_accessibility_state = node_ref.get::<AccessibilityNodeState>().unwrap();
                let accessibility_node =
                    Self::create_node(&node_ref, layout_node, &node_accessibility_state);
                nodes.push((accessibility_id, accessibility_node));
            }

            if let Some(tag) = node_ref.node_type().tag() {
                if *tag == TagName::Paragraph || *tag == TagName::Label {
                    return false;
                }
            }

            true
        });

        #[cfg(debug_assertions)]
        tracing::info!(
            "Initialized the Accessibility Tree with {} nodes.",
            nodes.len()
        );

        TreeUpdate {
            nodes,
            tree: Some(Tree::new(ACCESSIBILITY_ROOT_ID)),
            focus: ACCESSIBILITY_ROOT_ID,
        }
    }

    /// Process any pending Accessibility Tree update
    pub fn process_updates(
        &mut self,
        rdom: &DioxusDOM,
        layout: &Torin<NodeId>,
        dirty_nodes: &mut AccessibilityDirtyNodes,
    ) -> (TreeUpdate, NodeId) {
        let requested_focus_id = dirty_nodes.requested_focus.take();
        let removed_ids = dirty_nodes.removed.drain().collect::<FxHashMap<_, _>>();
        let mut added_or_updated_ids = dirty_nodes
            .added_or_updated
            .drain()
            .collect::<FxHashSet<_>>();

        #[cfg(debug_assertions)]
        if !removed_ids.is_empty() || !added_or_updated_ids.is_empty() {
            tracing::info!(
                "Updating the Accessibility Tree with {} removals and {} additions/modifications",
                removed_ids.len(),
                added_or_updated_ids.len()
            );
        }

        // Remove all the removed nodes from the update list
        for (node_id, _) in removed_ids.iter() {
            added_or_updated_ids.remove(node_id);
            self.map.retain(|_, id| id != node_id);
        }

        // Mark the parent of the removed nodes as updated
        for (_, parent_id) in removed_ids.iter() {
            if !removed_ids.contains_key(parent_id) {
                added_or_updated_ids.insert(*parent_id);
            }
        }

        // Mark the ancestors as modified
        for node_id in added_or_updated_ids.clone() {
            let node_ref = rdom.get(node_id).unwrap();
            let node_ref_parent = node_ref.parent_id().unwrap_or(rdom.root_id());
            added_or_updated_ids.insert(node_ref_parent);
            self.map
                .insert(node_ref.get_accessibility_id().unwrap(), node_id);
        }

        // Create the updated nodes
        let mut nodes = Vec::new();
        for node_id in added_or_updated_ids {
            let node_ref = rdom.get(node_id).unwrap();
            let node_accessibility_state = node_ref.get::<AccessibilityNodeState>();
            let layout_node = layout.get(node_id);

            if let Some((node_accessibility_state, layout_node)) =
                node_accessibility_state.as_ref().zip(layout_node)
            {
                let accessibility_node =
                    Self::create_node(&node_ref, layout_node, node_accessibility_state);
                let accessibility_id = node_ref.get_accessibility_id().unwrap();

                nodes.push((accessibility_id, accessibility_node));
            }
        }

        // Try to focus the requested node id
        if let Some(node_id) = requested_focus_id {
            let node_ref = rdom.get(node_id).unwrap();
            let node_accessibility_state = node_ref.get::<AccessibilityNodeState>();
            if let Some(focused_id) = node_accessibility_state
                .as_ref()
                .and_then(|state| state.a11y_id)
            {
                self.focused_id = focused_id;
            }
        }

        // Fallback the focused id to the root if the focused node no longer exists
        if !self.map.contains_key(&self.focused_id) {
            self.focused_id = ACCESSIBILITY_ROOT_ID;
        }

        let node_id = self.map.get(&self.focused_id).cloned().unwrap();

        (
            TreeUpdate {
                nodes,
                tree: Some(Tree::new(ACCESSIBILITY_ROOT_ID)),
                focus: self.focused_id,
            },
            node_id,
        )
    }

    /// Update the focused Node ID and generate a TreeUpdate if necessary.
    pub fn set_focus_with_update(
        &mut self,
        new_focus_id: AccessibilityId,
    ) -> Option<(TreeUpdate, NodeId)> {
        self.focused_id = new_focus_id;

        // Only focus the element if it exists
        if let Some(node_id) = self.map.get(&new_focus_id).copied() {
            #[cfg(debug_assertions)]
            tracing::info!("Focused {new_focus_id:?} node.");

            Some((
                TreeUpdate {
                    nodes: Vec::new(),
                    tree: Some(Tree::new(ACCESSIBILITY_ROOT_ID)),
                    focus: self.focused_id,
                },
                node_id,
            ))
        } else {
            None
        }
    }

    /// Focus a Node given the strategy.
    pub fn set_focus_on_next_node(
        &mut self,
        stragegy: AccessibilityFocusStrategy,
        rdom: &DioxusDOM,
    ) -> (TreeUpdate, NodeId) {
        let mut nodes = Vec::new();

        rdom.traverse_depth_first_advanced(|node_ref| {
            if !node_ref.node_type().is_element() {
                return false;
            }

            let accessibility_id = node_ref.get_accessibility_id();

            if let Some(accessibility_id) = accessibility_id {
                let accessibility_state = node_ref.get::<AccessibilityNodeState>().unwrap();
                if accessibility_state.a11y_focusable.is_enabled() {
                    nodes.push((accessibility_id, node_ref.id()))
                }
            }

            if let Some(tag) = node_ref.node_type().tag() {
                if *tag == TagName::Paragraph || *tag == TagName::Label {
                    return false;
                }
            }

            true
        });

        let node_index = nodes
            .iter()
            .enumerate()
            .find(|(_, (accessibility_id, _))| *accessibility_id == self.focused_id)
            .map(|(i, _)| i);

        let target_node = if stragegy == AccessibilityFocusStrategy::Forward {
            // Find the next Node
            if let Some(node_index) = node_index {
                if node_index == nodes.len() - 1 {
                    nodes.first()
                } else {
                    nodes.get(node_index + 1)
                }
            } else {
                nodes.first()
            }
        } else {
            // Find the previous Node
            if let Some(node_index) = node_index {
                if node_index == 0 {
                    nodes.last()
                } else {
                    nodes.get(node_index - 1)
                }
            } else {
                nodes.last()
            }
        };

        let (accessibility_id, node_id) = target_node
            .copied()
            .unwrap_or((ACCESSIBILITY_ROOT_ID, rdom.root_id()));

        self.focused_id = accessibility_id;

        #[cfg(debug_assertions)]
        tracing::info!("Focused {accessibility_id:?} node.");

        (
            TreeUpdate {
                nodes: Vec::new(),
                tree: Some(Tree::new(ACCESSIBILITY_ROOT_ID)),
                focus: self.focused_id,
            },
            node_id,
        )
    }

    /// Create an accessibility node
    pub fn create_node(
        node_ref: &DioxusNode,
        layout_node: &LayoutNode,
        node_accessibility: &AccessibilityNodeState,
    ) -> Node {
        let font_style_state = &*node_ref.get::<FontStyleState>().unwrap();
        let style_state = &*node_ref.get::<StyleState>().unwrap();
        let transform_state = &*node_ref.get::<TransformState>().unwrap();
        let node_type = node_ref.node_type();

        let mut builder = match node_type.tag() {
            // Make the root accessibility node.
            Some(&TagName::Root) => NodeBuilder::new(Role::Window),

            // All other node types will either don't have a builder (but don't support
            // accessibility attributes like with `text`) or have their builder made for
            // them already.
            Some(_) => node_accessibility.builder.clone().unwrap(),

            // Tag-less nodes can't have accessibility state
            None => unreachable!(),
        };

        // Set children
        let children = node_ref.get_accessibility_children();
        builder.set_children(children);

        // Set the area
        let area = layout_node.area.to_f64();
        builder.set_bounds(Rect {
            x0: area.min_x(),
            x1: area.max_x(),
            y0: area.min_y(),
            y1: area.max_y(),
        });

        if let NodeType::Element(node) = &*node_type {
            if matches!(node.tag, TagName::Label | TagName::Paragraph) && builder.name().is_none() {
                if let Some(inner_text) = node_ref.get_inner_texts() {
                    builder.set_name(inner_text);
                }
            }
        }

        // Set focusable action
        // This will cause assistive technology to offer the user an option
        // to focus the current element if it supports it.
        if node_accessibility.a11y_focusable.is_enabled() {
            builder.add_action(Action::Focus);
        }

        // Rotation transform
        if let Some((_, rotation)) = transform_state
            .rotations
            .iter()
            .find(|(id, _)| id == &node_ref.id())
        {
            builder.set_transform(Affine::rotate(rotation.to_radians() as _));
        }

        // Clipping overflow
        if style_state.overflow == OverflowMode::Clip {
            builder.set_clips_children();
        }

        // Foreground/Background color
        builder.set_foreground_color(skia_color_to_rgba_u32(font_style_state.color));
        if let Fill::Color(color) = style_state.background {
            builder.set_background_color(skia_color_to_rgba_u32(color));
        }

        // If the node is a block-level element in the layout, indicate that it will cause a linebreak.
        if !node_type.is_text() {
            if let NodeType::Element(node) = &*node_type {
                // This should be impossible currently but i'm checking for it just in case.
                // In the future, inline text spans should have their own own accessibility node,
                // but that's not a concern yet.
                if node.tag != TagName::Text {
                    builder.set_is_line_breaking_object();
                }
            }
        }

        // Font size
        builder.set_font_size(font_style_state.font_size as _);

        // If the font family has changed since the parent node, then we inform accesskit of this change.
        if let Some(parent_node) = node_ref.parent() {
            if parent_node.get::<FontStyleState>().unwrap().font_family
                != font_style_state.font_family
            {
                builder.set_font_family(font_style_state.font_family.join(", "));
            }
        } else {
            // Element has no parent elements, so we set the initial font style.
            builder.set_font_family(font_style_state.font_family.join(", "));
        }

        // Set bold flag for weights above 700
        if font_style_state.font_weight > 700.into() {
            builder.set_bold();
        }

        // Text alignment
        builder.set_text_align(match font_style_state.text_align {
            TextAlign::Center => accesskit::TextAlign::Center,
            TextAlign::Justify => accesskit::TextAlign::Justify,
            // TODO: change representation of `Start` and `End` once RTL text/writing modes are supported.
            TextAlign::Left | TextAlign::Start => accesskit::TextAlign::Left,
            TextAlign::Right | TextAlign::End => accesskit::TextAlign::Right,
        });

        // TODO: Adjust this once text direction support other than RTL is properly added
        builder.set_text_direction(TextDirection::LeftToRight);

        // Set italic property for italic/oblique font slants
        match font_style_state.font_slant {
            Slant::Italic | Slant::Oblique => builder.set_italic(),
            _ => {}
        }

        // Text decoration
        if font_style_state
            .decoration
            .ty
            .contains(TextDecoration::LINE_THROUGH)
        {
            builder.set_strikethrough(skia_decoration_style_to_accesskit(
                font_style_state.decoration.style,
            ));
        }
        if font_style_state
            .decoration
            .ty
            .contains(TextDecoration::UNDERLINE)
        {
            builder.set_underline(skia_decoration_style_to_accesskit(
                font_style_state.decoration.style,
            ));
        }
        if font_style_state
            .decoration
            .ty
            .contains(TextDecoration::OVERLINE)
        {
            builder.set_overline(skia_decoration_style_to_accesskit(
                font_style_state.decoration.style,
            ));
        }

        builder.build()
    }
}

fn skia_decoration_style_to_accesskit(style: TextDecorationStyle) -> accesskit::TextDecoration {
    match style {
        TextDecorationStyle::Solid => accesskit::TextDecoration::Solid,
        TextDecorationStyle::Dotted => accesskit::TextDecoration::Dotted,
        TextDecorationStyle::Dashed => accesskit::TextDecoration::Dashed,
        TextDecorationStyle::Double => accesskit::TextDecoration::Double,
        TextDecorationStyle::Wavy => accesskit::TextDecoration::Wavy,
    }
}

fn skia_color_to_rgba_u32(color: Color) -> u32 {
    ((color.a() as u32) << 24)
        | ((color.b() as u32) << 16)
        | (((color.g() as u32) << 8) + (color.r() as u32))
}