"""
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


def _fmt_float(val, precision=2):
    """Format a float for display with *precision* decimal places (default 2)."""
    return f"{val:.{precision}f}"


def _slot_precision(type_name, slot_name):
    """Looks up the precision facet for a slot, returning 2 as the default."""
    reg = app_state.final_components_registry.get(type_name)
    if reg:
        for s in reg:
            if s["name"] == slot_name:
                p = s.get("facets", {}).get("precision")
                if isinstance(p, int) and 0 <= p <= 12:
                    return p
                break
    return 2


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.checksum_resolver = schema.resolve_truncated_checksum
        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)
        if root_comp and root_comp.children:
            children = self.load_children(root_comp.id)
            if children:
                self.sox_client.update(children, CONFIG | RUNTIME)
        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 = {}
        own_slots = {}   # "kit::Type" -> declared slots
        bases = {}       # "kit::Type" -> "kit::Base" or None
        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,
                                     own_slots, bases)
        # Flatten each type's slots through its whole base chain (e.g.
        # Component -> Service -> ontrolOx::Widget -> PushButtonImage), then
        # renumber IDs to match the wire slot index (positional in the
        # flattened inherited + own list — see sedona Type.java resolveSlots).
        for key in registry:
            qname = self.sox_client.type_name_map[key]
            chain = []
            seen = set()
            t = qname
            while t and t not in seen and t in own_slots:
                seen.add(t)
                chain.append(t)
                t = bases.get(t)
            flat = []
            for t in reversed(chain):
                flat.extend(dict(s) for s in own_slots[t])
            for i, slot in enumerate(flat):
                slot["id"] = i
            registry[key] = flat
        self.sox_client.slot_registry = registry

    def _register_kit_types(self, registry, kit_id, kit_name, manifest,
                            own_slots, bases):
        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.findall("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}"
            own_slots[full_name] = slots
            bases[full_name] = type_elem.get("base") or None
            self.sox_client.type_name_map[(kit_id, type_id_val)] = full_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]

    def find_component_of_type(self, type_name, slot_name):
        """A component of `type_name` with `slot_name` read, loading where services live.

        Nothing is loaded until a folder is opened, so a platform service the user never
        expanded to is not cached yet. It sits one or two levels down (/service/plat), and
        those are loaded to find it rather than the whole app.
        """
        client = self.sox_client
        if client is None:
            return None

        def cached():
            return next((c for c in client.cache
                         if c is not None and c.type_name == type_name), None)

        comp = cached()
        if comp is None:
            for top in self.load_children(0):
                if top.name == "service" or top.type_name == type_name:
                    for child in self.load_children(top.id):
                        if child.children and child.type_name != type_name:
                            self.load_children(child.id)
            comp = cached()
        if comp is not None and slot_name not in comp.values:
            client.update([comp], CONFIG)
        return comp

    # ── 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()

    # ── Post-add / Post-delete bookkeeping ─────────────────────

    def after_add(self, comp_id):
        """Subscribes a newly added component and reads its initial values."""
        comp = self.sox_client.cache_get(comp_id)
        if comp is None or not self.is_connected:
            return
        try:
            self.sox_client.update([comp], CONFIG | RUNTIME)
        except Exception:
            pass
        try:
            self.sox_client.subscribe([comp], ALL)
            self._current_subs.add(comp_id)
        except Exception:
            pass

    def after_delete(self, comp_id, descendant_ids=None):
        """Removes a deleted component (and its descendants) from the subscription set."""
        self._current_subs.discard(comp_id)
        if descendant_ids:
            for did in descendant_ids:
                self._current_subs.discard(did)

    # ── 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:
            app_elem.set("id", str(root_comp.id))
            # The app is a component like any other, so the Properties view needs its
            # type and slot values too. A .sax file leaves both implied; here they are not.
            app_elem.set("type", root_comp.type_name)
            self._add_prop_elements(app_elem, root_comp)
            for child_id in root_comp.children:
                child_comp = self.sox_client.cache_get(child_id)
                if child_comp is not None:
                    self._build_comp_element(app_elem, child_comp)
                else:
                    ET.SubElement(app_elem, "comp",
                                  name="...", id=str(child_id), type="")

        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)
        self._add_prop_elements(comp_elem, sox_comp)

        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)
            else:
                # Placeholder so the tree shows a + expand arrow for uncached children
                ET.SubElement(comp_elem, "comp",
                              name="...", id=str(child_id), type="")

        return comp_elem

    def _add_prop_elements(self, comp_elem, sox_comp):
        """Writes a <prop> for every cached non-action slot value of `sox_comp`."""
        slots = self.sox_client.slot_registry.get(
            (sox_comp.kit_id, sox_comp.type_id))
        if not slots:
            return
        for slot in slots:
            if slot.get("is_action"):
                continue
            name = slot["name"]
            if name in sox_comp.values:
                val = sox_comp.values[name]
                if isinstance(val, float):
                    val_str = _fmt_float(val, _slot_precision(sox_comp.type_name, name))
                else:
                    val_str = str(val)
                ET.SubElement(comp_elem, "prop", name=name, val=val_str)

    def build_links_elements(self, sax_root):
        """Adds <link> elements under a <links> container, matching .sax structure."""
        links_elem = sax_root.find("links")
        if links_elem is None:
            links_elem = ET.SubElement(sax_root, "links")

        seen = set()
        for c in self.sox_client.cache:
            if c is None:
                continue
            for lnk in c.links:
                link_key = (lnk.from_comp_id, lnk.from_slot_id,
                            lnk.to_comp_id, lnk.to_slot_id)
                if link_key in seen:
                    continue
                seen.add(link_key)

                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(links_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):
        # SOX links use absolute slot IDs across the full inheritance chain.
        type_name = self.sox_client.type_name_map.get((kit_id, type_id_val))
        if type_name and type_name in app_state.final_components_registry:
            for s in app_state.final_components_registry[type_name]:
                if s["slot_id"] == slot_id:
                    return s["name"]
        # Fallback to slot_registry (flattened through the base chain, so
        # its IDs are absolute too)
        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

        type_name = sox_comp.type_name
        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:
                    raw = sox_comp.values[name]
                    if isinstance(raw, float):
                        val = _fmt_float(raw, _slot_precision(type_name, name))
                    else:
                        val = str(raw)
                    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)
        app_elem = root.find("app")
        if app_elem is not None and app_elem.get("id") == target:
            return app_elem
        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