from threading import Lock

from app.models.schemas import DeviceState


class DeviceStore:
    def __init__(self) -> None:
        self._devices: dict[str, DeviceState] = {}
        self._acks: dict[str, list[dict]] = {}
        self._events: dict[str, list[dict]] = {}
        self._lock = Lock()

    def update_state(self, device_id: str, payload: dict, timestamp: int) -> DeviceState:
        with self._lock:
            state = DeviceState(
                device_id=device_id,
                cpu=payload.get("cpu"),
                mem=payload.get("mem"),
                battery=payload.get("battery"),
                net=payload.get("net"),
                mode=payload.get("mode"),
                last_seen=timestamp,
            )
            self._devices[device_id] = state
            return state

    def record_ack(self, device_id: str, payload: dict) -> None:
        with self._lock:
            self._acks.setdefault(device_id, []).append(payload)

    def record_event(self, device_id: str, payload: dict) -> None:
        with self._lock:
            self._events.setdefault(device_id, []).append(payload)

    def get_device(self, device_id: str) -> DeviceState | None:
        return self._devices.get(device_id)

    def list_devices(self) -> list[DeviceState]:
        return list(self._devices.values())

    def get_acks(self, device_id: str) -> list[dict]:
        return list(self._acks.get(device_id, []))

    def get_events(self, device_id: str) -> list[dict]:
        return list(self._events.get(device_id, []))