
import time
import requests
from requests.auth import HTTPBasicAuth
from typing import Any, Dict, Optional, Union

class SpectrometerClient:
    def __init__(self, base_url: str = "http://localhost:1818",
                 username: Optional[str] = None, password: Optional[str] = None,
                 timeout_s: float = 30.0):
        self.base = base_url.rstrip("/")
        # Auth removed: spectrometer API is open (local lab device)
        self.auth = None
        self.timeout = timeout_s

        # DRF root autodiscovery
        idx = self._get_json(f"{self.base}/api/")
        # коллекции с фоллбэками:
        self.url_users_coll   = (idx.get("users")   or f"{self.base}/api/users/").rstrip("/")   + "/"
        self.url_devices_coll = (idx.get("devices") or f"{self.base}/api/devices/").rstrip("/") + "/"
        self.url_mparams_coll = (idx.get("mparams") or f"{self.base}/api/mparams/").rstrip("/") + "/"
        self.url_mstate_coll  = (idx.get("mstate")  or f"{self.base}/api/mstate/").rstrip("/")  + "/"
        self.url_mdata_coll   = (idx.get("mdata")   or f"{self.base}/api/mdata/").rstrip("/")   + "/"
        self.url_measure_coll = (idx.get("measure") or f"{self.base}/api/measure/").rstrip("/") + "/"

        self.last_response: Optional[requests.Response] = None

    # --------------- high-level ---------------

    def start_measurement(self, body: Dict[str, Any]) -> int:
        """
        POST /api/measure/
        body: должен быть {"info": {...}}. Если пришли плоские поля - обернём.
        Возвращает числовой id измерения.
        """
        if "info" not in body:
            body = {"info": dict(body)}
        res = self._post_json(self.url_measure_coll, body)
        # Пытаемся извлечь id (разные реализации могут различаться)
        meas_id = (
            res.get("id")
            or res.get("pk")
            or res.get("uuid")
            or res.get("measurement_id")
            or (res.get("data") or {}).get("id")
        )
        if meas_id is None:
            raise RuntimeError(f"Cannot determine measurement id from response: {res}")
        try:
            return int(meas_id)
        except Exception:
            # если id не числовой, пусть будет строкой, но клиент явно ошибку не кидает
            # однако downstream ожидает число - лучше привести к ошибке здесь:
            raise RuntimeError(f"Measurement id is not an integer: {meas_id!r}")

    def get_state(self, meas_id: Union[int, str]) -> Dict[str, Any]:
        # 1) detail action: GET /api/measure/{id}/state/  (preferred)
        url1 = f"{self.url_measure_coll}{meas_id}/state/"
        r1 = self._get_any(url1)
        if r1 is not None:
            return r1
        # 2) fallback: /api/mstate/?id_measurement={id}
        #    filterset_fields = ['id_measurement'] — must use this param name
        result = self._get_json(self.url_mstate_coll, params={"id_measurement": meas_id})
        # DRF may return a list or paginated {"results": [...]}
        if isinstance(result, list) and result:
            return result[0]
        if isinstance(result, dict) and "results" in result:
            items = result["results"]
            return items[0] if items else {}
        return result

    def wait_data_ready(
        self,
        meas_id: Union[int, str],
        timeout_s: float = 120.0,
        poll_s: float = 1.0,
    ) -> Dict[str, Any]:
        """
        Poll GET /api/measure/{id}/state/ until data is ready.

        Completion criteria (in priority order):
          1. data_ready == True          ← primary signal (set by the engine)
          2. status in {done, failed, error}  ← secondary; engine may or may not set this

        Raises TimeoutError if neither condition is met within timeout_s.
        """
        t0 = time.time()
        while True:
            st = self.get_state(meas_id)
            if bool(st.get("data_ready")):
                return st
            hw_status = st.get("status", "")
            if hw_status in {"failed", "error"}:
                raise RuntimeError(
                    f"Measurement {meas_id} ended with status={hw_status!r}: {st}"
                )
            if time.time() - t0 > timeout_s:
                raise TimeoutError(
                    f"Timeout ({timeout_s}s) waiting for data_ready "
                    f"on measurement {meas_id}. Last state: {st}"
                )
            time.sleep(poll_s)

    def get_data(
        self,
        meas_id: Union[int, str],
        data_num: Optional[int] = None,
        averaging_num: Optional[int] = None,
    ) -> Any:
        """
        GET /api/measure/{id}/data/[?data_num=N&averaging_num=M]

        Returns a list of measurement_data records.
        Falls back to /api/mdata/?measurement__id={id} if the action endpoint
        returns 404.

        data_num / averaging_num are forwarded as query params when provided,
        matching the filter logic in measurement_ViewSet.data().
        """
        url1 = f"{self.url_measure_coll}{meas_id}/data/"
        params: Dict[str, Any] = {}
        if data_num is not None:
            params["data_num"] = data_num
        if averaging_num is not None:
            params["averaging_num"] = averaging_num

        r1 = self._get_any(url1, params=params or None)
        if r1 is not None:
            return r1
        # fallback: mdata collection filtered by measurement id
        # measurement_data_ViewSet filterset_fields = ['measurement__id', ...]
        fb_params: Dict[str, Any] = {"measurement__id": meas_id}
        if data_num is not None:
            fb_params["data_num"] = data_num
        if averaging_num is not None:
            fb_params["averaging_num"] = averaging_num
        return self._get_json(self.url_mdata_coll, params=fb_params)

    # --------------- low-level ---------------

    def _get_any(
        self, url: str, params: Optional[Dict[str, Any]] = None
    ) -> Optional[Dict[str, Any]]:
        try:
            r = requests.get(url, params=params, auth=self.auth, timeout=self.timeout)
            if r.status_code == 404:
                return None
            r.raise_for_status()
            self.last_response = r
            return r.json()
        except requests.HTTPError:
            raise

    def _get_json(self, url: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        r = requests.get(url, params=params, auth=self.auth, timeout=self.timeout)
        r.raise_for_status()
        self.last_response = r
        return r.json()

    def _post_json(self, url: str, body: Dict[str, Any]) -> Dict[str, Any]:
        r = requests.post(url, json=body, auth=self.auth, timeout=self.timeout)
        r.raise_for_status()
        self.last_response = r
        return r.json()
