from __future__ import annotations

from typing import Any

from device.manager import DeviceManager
from perception.ui_parser import parse_ui_tree
from tools._device import get_active_app_info, get_driver


def get_raw_ui_tree(
    device_id: str | None = None,
    *,
    manager: DeviceManager | None = None,
) -> Any:
    return get_driver(device_id, manager=manager).tree()


def get_ui_tree(
    device_id: str | None = None,
    *,
    manager: DeviceManager | None = None,
    include_app_info: bool = False,
) -> list[dict[str, object]] | dict[str, object]:
    driver = get_driver(device_id, manager=manager)
    nodes = [element.to_dict() for element in parse_ui_tree(driver.tree())]
    if not include_app_info:
        return nodes
    active_app = get_active_app_info(driver)
    return {
        "app": active_app.to_dict() if active_app is not None else None,
        "nodes": nodes,
    }