sedona_editor/device_session.py

"""
DeviceSession — bridges SoxClient to the editor's ET.Element-based UI.

Handles lazy loading, DOM synthesis from cached SoxComponents,
COV event routing, and subscription lifecycle.
"""

import queue
import xml.etree.ElementTree as ET

import app_state
import schema
from sox import SoxClient, SoxComponent, TREE, CONFIG, RUNTIME, LINKS, ALL


class DeviceSession:
    """High-level session manager for a live Sedona device connection."""

    def __init__(self):
        self.sox_client = None
        self._event_queue = queue.Queue()
        self._current_subs = set()   # comp_ids we subscribed for the open wiresheet

    # ── Connect / Disconnect ────────────────────────────────────

    def connect(self, host, port, user, password):
        """DASP connect → read schema → load root. Called from background thread."""
        self.sox_client = SoxClient(host, int(port), user, password)
        self.sox_client.on_event = self._on_sox_event
        self.sox_client.connect()

        parts = self.sox_client.read_schema()
        vi = self.sox_client.read_version()

        self._load_slot_registry(parts)

        root_comp = self.sox_client.load(0)
        return root_comp

    def disconnect(self):
        if self.sox_client is None:
            return
        try:
            if self._current_subs:
                comps = [self.sox_client.cache_get(cid)
                         for cid in self._current_subs]
                comps = [c for c in comps if c is not None]
                if comps:
                    self.sox_client.unsubscribe(comps, ALL)
        except Exception:
            pass
        self._current_subs.clear()
        try:
            self.sox_client.close()
        except Exception:
            pass
        self.sox_client = None

    @property
    def is_connected(self):
        return self.sox_client is not None and not self.sox_client.is_closed

    # ── Schema → Slot Registry ──────────────────────────────────

    def _load_slot_registry(self, parts):
        """Builds slot_registry from kit manifests for the device's schema."""
        registry = {}
        for kit_id, part in enumerate(parts):
            kit_name = part.name
            checksum = part.checksum
            manifest = schema.find_manifest(kit_name, checksum)
            if manifest is None:
                continue
            self._register_kit_types(registry, kit_id, kit_name, manifest)
        self.sox_client.slot_registry = registry

    def _register_kit_types(self, registry, kit_id, kit_name, manifest):
        root = manifest
        for type_elem in root.iter("type"):
            type_name = type_elem.get("name", "")
            type_id_val = int(type_elem.get("id", "-1"))
            if type_id_val < 0:
                continue
            slots = []
            for slot_elem in type_elem.iter("slot"):
                slot_name = slot_elem.get("name", "")
                slot_id = int(slot_elem.get("id", "-1"))
                slot_type = slot_elem.get("type", "int")
                flags_str = slot_elem.get("flags", "")
                is_config = "c" in flags_str or "p" in flags_str
                is_runtime = "r" in flags_str
                is_action = "a" in flags_str
                is_operator = "o" in flags_str
                is_as_str = slot_type == "Str" or slot_elem.get("asStr") == "true"
                slots.append({
                    "name": slot_name,
                    "id": slot_id,
                    "type_name": slot_type,
                    "is_config": is_config,
                    "is_runtime": is_runtime,
                    "is_action": is_action,
                    "is_operator": is_operator,
                    "is_as_str": is_as_str,
                    "is_prop": not is_action,
                })
            registry[(kit_id, type_id_val)] = slots
            full_name = f"{kit_name}::{type_name}"
            # update type_name on any cached component of this type
            for c in self.sox_client.cache:
                if c is not None and c.kit_id == kit_id and c.type_id == type_id_val:
                    c.type_name = full_name

    # ── Lazy Loading ────────────────────────────────────────────

    def load_children(self, comp_id):
        """Loads direct children of a component (cache-aware). Returns list of SoxComponent."""
        parent = self.sox_client.cache_get(comp_id)
        if parent is None:
            parent = self.sox_client.load(comp_id)
        if parent is None:
            return []
        uncached = [cid for cid in parent.children
                    if self.sox_client.cache_get(cid) is None]
        if uncached:
            self.sox_client.load_ids(uncached, checked=False)
        return [self.sox_client.cache_get(cid)
                for cid in parent.children
                if self.sox_client.cache_get(cid) is not None]

    # ── Wiresheet Open / Close ──────────────────────────────────

    def open_wiresheet(self, folder_comp_id):
        """Full wiresheet load: children, links, remote comps, config+runtime, subscribe.

        Returns list of SoxComponent for the folder's children.
        """
        children = self.load_children(folder_comp_id)
        if not children:
            return []

        all_link_remote_ids = set()
        for child in children:
            links = self.sox_client.read_links(child.id)
            child.links = links
            for lnk in links:
                all_link_remote_ids.add(lnk.from_comp_id)
                all_link_remote_ids.add(lnk.to_comp_id)

        child_ids = {c.id for c in children}
        child_ids.add(folder_comp_id)
        remote_ids = [rid for rid in all_link_remote_ids
                      if rid not in child_ids
                      and self.sox_client.cache_get(rid) is None]
        if remote_ids:
            self.sox_client.load_ids(remote_ids, checked=False)

        self.sox_client.update(children, CONFIG | RUNTIME)

        folder_comp = self.sox_client.cache_get(folder_comp_id)
        to_subscribe = [folder_comp] + children if folder_comp else children
        to_subscribe = [c for c in to_subscribe if c is not None]
        try:
            self.sox_client.subscribe(to_subscribe, ALL)
        except Exception:
            pass

        self._current_subs = {c.id for c in to_subscribe}
        return children

    def close_wiresheet(self):
        """Unsubscribe from previously subscribed components."""
        if not self._current_subs or not self.is_connected:
            self._current_subs.clear()
            return
        comps = [self.sox_client.cache_get(cid)
                 for cid in self._current_subs]
        comps = [c for c in comps if c is not None]
        if comps:
            try:
                self.sox_client.unsubscribe(comps, ALL)
            except Exception:
                pass
        self._current_subs.clear()

    # ── DOM Synthesis ───────────────────────────────────────────

    def build_dom_tree(self, root_comp):
        """Builds an ET.Element tree from cached SoxComponents, for the editor UI.

        Only builds from what is currently in the cache (lazy-loaded components).
        Returns a root <sedonaApp> element.
        """
        sax_root = ET.Element("sedonaApp")
        schema_elem = ET.SubElement(sax_root, "schema")
        if self.sox_client and self.sox_client.schema_parts:
            for part in self.sox_client.schema_parts:
                ET.SubElement(schema_elem, "kit",
                              name=part.name, checksum=str(part.checksum))

        app_elem = ET.SubElement(sax_root, "app")
        if root_comp is not None:
            self._build_comp_element(app_elem, root_comp)

        return sax_root

    def _build_comp_element(self, parent_elem, sox_comp):
        type_qname = sox_comp.type_name
        attribs = {
            "name": sox_comp.name,
            "id": str(sox_comp.id),
            "type": type_qname,
        }
        comp_elem = ET.SubElement(parent_elem, "comp", **attribs)

        slots = self.sox_client.slot_registry.get(
            (sox_comp.kit_id, sox_comp.type_id))
        if slots:
            for slot in slots:
                if slot.get("is_action"):
                    continue
                name = slot["name"]
                if name in sox_comp.values:
                    val = sox_comp.values[name]
                    ET.SubElement(comp_elem, "prop",
                                  name=name, val=str(val))

        for child_id in sox_comp.children:
            child_comp = self.sox_client.cache_get(child_id)
            if child_comp is not None:
                self._build_comp_element(comp_elem, child_comp)

        return comp_elem

    def build_links_elements(self, sax_root):
        """Adds <link> elements under the app element from cached link data."""
        app_elem = sax_root.find("app")
        if app_elem is None:
            return
        for c in self.sox_client.cache:
            if c is None:
                continue
            for lnk in c.links:
                from_comp = self.sox_client.cache_get(lnk.from_comp_id)
                to_comp = self.sox_client.cache_get(lnk.to_comp_id)
                if from_comp is None or to_comp is None:
                    continue

                from_slot_name = self._slot_name_by_id(
                    from_comp.kit_id, from_comp.type_id, lnk.from_slot_id)
                to_slot_name = self._slot_name_by_id(
                    to_comp.kit_id, to_comp.type_id, lnk.to_slot_id)

                from_path = self._comp_path(lnk.from_comp_id)
                to_path = self._comp_path(lnk.to_comp_id)

                if from_path and to_path and from_slot_name and to_slot_name:
                    ET.SubElement(app_elem, "link",
                                  **{"from": f"{from_path}.{from_slot_name}",
                                     "to": f"{to_path}.{to_slot_name}"})

    def _slot_name_by_id(self, kit_id, type_id_val, slot_id):
        slots = self.sox_client.slot_registry.get((kit_id, type_id_val))
        if slots:
            for s in slots:
                if s["id"] == slot_id:
                    return s["name"]
        return None

    def _comp_path(self, comp_id):
        parts = []
        cid = comp_id
        while cid >= 0:
            c = self.sox_client.cache_get(cid)
            if c is None:
                break
            if c.id == 0:
                break
            parts.append(c.name)
            cid = c.parent
        if not parts:
            return None
        parts.reverse()
        return "/".join(parts)

    # ── Update DOM from COV ─────────────────────────────────────

    def update_element_from_cache(self, comp_id, sax_root):
        """Updates a single comp element in the DOM from cached SoxComponent values."""
        sox_comp = self.sox_client.cache_get(comp_id)
        if sox_comp is None:
            return None

        comp_elem = self._find_element_by_id(sax_root, comp_id)
        if comp_elem is None:
            return None

        slots = self.sox_client.slot_registry.get(
            (sox_comp.kit_id, sox_comp.type_id))
        if slots:
            existing_props = {p.get("name"): p for p in comp_elem.findall("prop")}
            for slot in slots:
                if slot.get("is_action"):
                    continue
                name = slot["name"]
                if name in sox_comp.values:
                    val = str(sox_comp.values[name])
                    if name in existing_props:
                        existing_props[name].set("val", val)
                    else:
                        ET.SubElement(comp_elem, "prop", name=name, val=val)
        return comp_elem

    @staticmethod
    def _find_element_by_id(root, comp_id):
        target = str(comp_id)
        for elem in root.iter("comp"):
            if elem.get("id") == target:
                return elem
        return None

    # ── Event queue (thread-safe) ───────────────────────────────

    def _on_sox_event(self, comp_id, what_char):
        self._event_queue.put((comp_id, what_char))

    def poll_events(self):
        """Drains the event queue. Returns list of (comp_id, what_char) tuples."""
        events = []
        try:
            while True:
                events.append(self._event_queue.get_nowait())
        except queue.Empty:
            pass
        return events