# apps/aroflo_connector_app/agent/cheap/context_bridge.py
from __future__ import annotations

from typing import Any, Dict, List, Optional

from .. import context_store


def _extract_list(result_wrapper: Any, key: str) -> Optional[List[Dict[str, Any]]]:
    # Case A: {"zoneresponse": {key: [...]}}
    if (
        isinstance(result_wrapper, dict)
        and isinstance(result_wrapper.get("zoneresponse"), dict)
        and isinstance(result_wrapper["zoneresponse"].get(key), list)
    ):
        return result_wrapper["zoneresponse"][key]

    # Case B: wrapper: {"data": <resp>}  where <resp> may contain zoneresponse or direct key
    if isinstance(result_wrapper, dict) and "data" in result_wrapper:
        data = result_wrapper.get("data")
        if isinstance(data, dict):
            zr = data.get("zoneresponse")
            if isinstance(zr, dict) and isinstance(zr.get(key), list):
                return zr[key]
            if isinstance(data.get(key), list):
                return data[key]

    return None


def remember_any_if_present(zone_code: str, result_wrapper: Any) -> None:
    if zone_code == "tasks":
        items = _extract_list(result_wrapper, "tasks")
        if items is not None:
            context_store.remember_tasks_from_result({"data": {"tasks": items}})
        return

    if zone_code == "users":
        items = _extract_list(result_wrapper, "users")
        if items is not None:
            context_store.remember_users_from_result({"data": {"users": items}})
        return

    if zone_code == "permissiongroups":
        items = _extract_list(result_wrapper, "permissiongroups")
        if items is not None:
            context_store.remember_permissiongroups_from_result({"data": {"permissiongroups": items}})
        return

    if zone_code == "userpositions":
        items = _extract_list(result_wrapper, "userpositions")
        if items is not None:
            context_store.remember_userpositions_from_result({"data": {"userpositions": items}})
        return
    
    if zone_code == "businessunits":
        items = _extract_list(result_wrapper, "businessunits")
        if items is not None:
            context_store.remember_businessunits_from_result({"data": {"businessunits": items}})
        return




def get_task_ref_by_index(idx: int) -> Optional[Dict[str, Any]]:
    return context_store.get_task_by_index(idx)


def get_user_ref_by_index(idx: int) -> Optional[Dict[str, Any]]:
    return context_store.get_user_by_index(idx)

def get_businessunit_ref_by_index(idx: int) -> Optional[Dict[str, Any]]:
    return context_store.get_businessunit_by_index(idx)