"""
BADAS API Client
================
Python client for the BADAS Inference API (https://badas.api.corp.nexars.ai).

    from badas.client import BADASClient

    client = BADASClient(api_key="nxr_live_...")

    # Block until done — progress bar included
    preds = client.predict("video.mp4")

    # Stream predictions as each window completes — progress bar included
    for pred in client.stream("video.mp4", describe=True, threshold=0.8):
        print(f"[{pred['timestamp']:.1f}s] {pred['risk_level']} ({pred['probability']:.0%})")

    # Async — submit now, stream or poll later
    job_id = client.submit("video.mp4")
    preds  = list(client.stream_job(job_id))
    preds  = client.poll_job(job_id)
"""

from __future__ import annotations

import json
import time
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional, Union

_LARGE_FILE_THRESHOLD = 50 * 1024 * 1024  # 50 MB — use GCS signed URL above this


class BADASError(Exception):
    """Raised when the BADAS API returns an error response."""
    def __init__(self, status_code: int, message: str):
        self.status_code = status_code
        super().__init__(f"HTTP {status_code}: {message}")


class BADASClient:
    """Client for the BADAS Inference API.

    Args:
        base_url:  API base URL. Defaults to the production corp endpoint.
        api_key:   Bearer token (``nxr_live_...``). Optional for internal
                   corp access where IAP handles auth.
        timeout:   HTTP timeout in seconds (non-streaming requests).
        model:     Default model ID — overridable per call.
        stride:    Default stride in model-fps frames — overridable per call.
    """

    DEFAULT_MODEL = "nexar-ai/badas-2.0-flash"
    MODELS = {
        "nexar-ai/badas-2.0":            "330M ViT-L  — best accuracy",
        "nexar-ai/badas-2.0-flash":      " 86M ViT-B  — fast, default",
        "nexar-ai/badas-2.0-flash-lite": " 22M ViT-S  — fastest",
        "nexar-ai/badas-1.0":            "330M ViT-L  — previous version",
    }

    def __init__(
        self,
        base_url: str           = "https://badas.api.corp.nexars.ai",
        api_key:  Optional[str] = None,
        timeout:  float         = 120.0,
        model:    str           = DEFAULT_MODEL,
        stride:   int           = 1,
    ):
        self.base_url = base_url.rstrip("/")
        self.api_key  = api_key
        self.timeout  = timeout
        self.model    = model
        self.stride   = stride
        self._session = None

    def __repr__(self) -> str:
        return (f"BADASClient(base_url={self.base_url!r}, "
                f"model={self.model!r}, stride={self.stride})")

    # ── HTTP session ───────────────────────────────────────────────────────────

    @property
    def _sess(self):
        """Lazy-initialised requests.Session with auth header pre-set."""
        if self._session is None:
            import requests
            s = requests.Session()
            if self.api_key:
                s.headers["Authorization"] = f"Bearer {self.api_key}"
            self._session = s
        return self._session

    def _url(self, path: str) -> str:
        return f"{self.base_url}/{path.lstrip('/')}"

    def _check(self, resp) -> None:
        if resp.status_code >= 400:
            try:
                body = resp.json()
                msg = body.get("detail") or body.get("message") or resp.text
            except Exception:
                msg = resp.text
            raise BADASError(resp.status_code, msg)

    # ── Inference parameter helpers ────────────────────────────────────────────

    def _build_params(
        self,
        model:        Optional[str],
        stride:       Optional[int],
        heatmaps:     str,
        describe:     bool,
        threshold:    float,
        include_maps: bool,
    ) -> Dict[str, Any]:
        # describe=True implies heatmaps="threshold" unless caller was explicit
        if describe and heatmaps == "none":
            heatmaps = "threshold"
        return {
            "model":        model or self.model,
            "stride":       stride if stride is not None else self.stride,
            "heatmaps":     heatmaps,
            "describe":     str(describe).lower(),
            "threshold":    threshold,
            "include_maps": str(include_maps).lower(),
        }

    def _resolve_video(
        self,
        video:  Union[str, Path],
        params: Dict[str, Any],
        files:  Dict[str, Any],
    ) -> None:
        """Populate *params* / *files* based on the video source.

        - ``gs://`` URI   → ``gcs_uri`` param (pass-through)
        - ``http(s)://``  → ``video_url`` param (pass-through)
        - local < 50 MB   → multipart ``file`` upload
        - local ≥ 50 MB   → GCS signed-URL upload, then ``gcs_uri`` param
        """
        s = str(video)
        if s.startswith("gs://"):
            params["gcs_uri"] = s
            return
        if s.startswith(("http://", "https://")):
            params["video_url"] = s
            return

        path = Path(s)
        if not path.exists():
            raise FileNotFoundError(f"Video not found: {path}")
        if path.stat().st_size >= _LARGE_FILE_THRESHOLD:
            params["gcs_uri"] = self._upload_to_gcs(path)
        else:
            files["file"] = (path.name, open(path, "rb"), "video/mp4")

    def _upload_to_gcs(self, path: Path) -> str:
        """Upload a large file directly to GCS via signed URL, return gs:// URI."""
        resp = self._sess.get(
            self._url("/jobs/upload-url"),
            params={"filename": path.name},
            timeout=self.timeout,
        )
        self._check(resp)
        data = resp.json()

        size = path.stat().st_size
        print(f"Uploading {path.name} ({size / 1e6:.1f} MB) to GCS...")
        with open(path, "rb") as f:
            # GCS signed URL — must NOT carry the Authorization header
            import requests as _req
            r = _req.put(data["upload_url"], data=f, timeout=3600)
            r.raise_for_status()
        print("Upload complete.")
        return data["gcs_uri"]

    # ── Public API ─────────────────────────────────────────────────────────────

    def health(self) -> Dict[str, Any]:
        """Return API health — models loaded, VLM status."""
        resp = self._sess.get(self._url("/health"), timeout=self.timeout)
        self._check(resp)
        return resp.json()

    def predict(
        self,
        video: Union[str, Path],
        *,
        model:        Optional[str] = None,
        stride:       Optional[int] = None,
        heatmaps:     str           = "none",
        describe:     bool          = False,
        threshold:    float         = 0.75,
        include_maps: bool          = False,
        progress:     bool          = True,
    ) -> List[Dict[str, Any]]:
        """Run full-video inference and return all predictions.

        Submits an async job and polls until complete — works for any video
        length regardless of gateway timeouts.

        Args:
            video:        Local path, ``gs://`` URI, or ``http(s)://`` URL.
            model:        Model ID (default: ``badas-2.0-flash``).
            stride:       Windows per model-fps frame (default: client.stride).
            heatmaps:     ``"none"`` | ``"threshold"`` | ``"always"``.
            describe:     VLM hazard description on trigger windows.
                          Automatically sets ``heatmaps="threshold"`` if not
                          specified.
            threshold:    Risk score that triggers heatmaps / descriptions.
            include_maps: Include raw ``hazard_map`` / ``attention_map`` arrays.
            progress:     Show a progress bar (auto-adapts to notebook / CLI).

        Returns:
            List of prediction dicts, one per window.
        """
        job_id = self.submit(
            video, model=model, stride=stride, heatmaps=heatmaps,
            describe=describe, threshold=threshold, include_maps=include_maps,
        )
        return self.poll_job(job_id, progress=progress)

    def stream(
        self,
        video: Union[str, Path],
        *,
        model:        Optional[str] = None,
        stride:       Optional[int] = None,
        heatmaps:     str           = "none",
        describe:     bool          = False,
        threshold:    float         = 0.75,
        include_maps: bool          = False,
        progress:     bool          = True,
    ) -> Iterator[Dict[str, Any]]:
        """Stream predictions as they arrive.

        Submits an async job and streams results via SSE. Reconnects
        automatically on gateway timeout — transparent to the caller.

        A progress bar is shown by default (auto-adapts to notebook / CLI).
        Pass ``progress=False`` to suppress it, e.g. when driving your own UI.

        Yields one prediction dict per window. The iterator stops naturally
        when the job is done — no sentinel event is yielded.

        Args: same as :meth:`predict`.
        """
        job_id = self.submit(
            video, model=model, stride=stride, heatmaps=heatmaps,
            describe=describe, threshold=threshold, include_maps=include_maps,
        )
        yield from self.stream_job(job_id, progress=progress)

    def submit(
        self,
        video: Union[str, Path],
        *,
        model:        Optional[str] = None,
        stride:       Optional[int] = None,
        heatmaps:     str           = "none",
        describe:     bool          = False,
        threshold:    float         = 0.75,
        include_maps: bool          = False,
    ) -> str:
        """Submit an async inference job and return the job ID immediately.

        Use :meth:`stream_job` to stream results or :meth:`poll_job` to wait.

        Args: same as :meth:`predict` (minus ``progress``).

        Returns:
            Job ID string.
        """
        params: Dict[str, Any] = self._build_params(
            model, stride, heatmaps, describe, threshold, include_maps,
        )
        files: Dict[str, Any] = {}
        self._resolve_video(video, params, files)
        resp = self._sess.post(
            self._url("/jobs"),
            data=params,
            files=files or None,
            timeout=self.timeout,
        )
        self._check(resp)
        return resp.json()["job_id"]

    def stream_job(self, job_id: str, *, progress: bool = True) -> Iterator[Dict[str, Any]]:
        """Stream results for an existing job via SSE.

        Reconnects automatically if the gateway drops the connection.
        Yields one prediction dict per window. The iterator stops naturally
        when the job is done — no sentinel event is yielded.

        Args:
            job_id:   Job ID from :meth:`submit`.
            progress: Show a progress bar (auto-adapts to notebook / CLI).

        Raises:
            BADASError: If the job fails or returns an HTTP error.
        """
        gen = self._sse_stream(job_id)
        if progress:
            yield from _with_progress(gen, desc="Streaming")
        else:
            yield from gen

    def poll_job(
        self,
        job_id: str,
        *,
        poll_interval: float = 2.0,
        timeout:       float = 3600.0,
        progress:      bool  = True,
    ) -> List[Dict[str, Any]]:
        """Poll until a job is complete and return all predictions.

        Args:
            job_id:        Job ID from :meth:`submit`.
            poll_interval: Seconds between polls.
            timeout:       Max seconds to wait before raising ``TimeoutError``.
            progress:      Show a progress bar (auto-adapts to notebook / CLI).

        Returns:
            List of prediction dicts.

        Raises:
            TimeoutError: Job didn't finish within *timeout* seconds.
            BADASError:   Job failed.
        """
        try:
            from tqdm.auto import tqdm as _tqdm
        except ImportError:
            _tqdm = None

        pbar   = _tqdm(unit="win", desc="Processing") if (progress and _tqdm) else None
        last_n = 0
        try:
            deadline = time.time() + timeout
            while time.time() < deadline:
                status = self.job_status(job_id)
                n = status.get("num_results", 0)
                if pbar is not None and n > last_n:
                    pbar.update(n - last_n)
                    last_n = n
                if status["status"] == "done":
                    if pbar is not None:
                        remaining = status.get("num_results", last_n) - last_n
                        if remaining > 0:
                            pbar.update(remaining)
                        pbar.set_description("Done")
                    return status["results"]
                if status["status"] == "failed":
                    raise BADASError(500, f"Job {job_id} failed: {status.get('error')}")
                time.sleep(poll_interval)
        finally:
            if pbar is not None:
                pbar.close()
        raise TimeoutError(f"Job {job_id} did not complete within {timeout}s")

    def job_status(self, job_id: str) -> Dict[str, Any]:
        """Return current status dict and partial results for a job."""
        resp = self._sess.get(self._url(f"/jobs/{job_id}"), timeout=self.timeout)
        self._check(resp)
        return resp.json()

    # ── Internal ───────────────────────────────────────────────────────────────

    def _sse_stream(self, job_id: str) -> Iterator[Dict[str, Any]]:
        """Reconnect-safe raw SSE stream. Yields prediction dicts only."""
        last_id: Optional[str] = None
        while True:
            headers = {}
            if last_id is not None:
                headers["Last-Event-ID"] = last_id
            try:
                resp = self._sess.get(
                    self._url(f"/jobs/{job_id}/stream"),
                    headers=headers,
                    stream=True,
                    timeout=self.timeout,
                )
                self._check(resp)
                for event_id, event in _parse_sse(resp):
                    if event_id is not None:
                        last_id = event_id
                    if event.get("type") == "error":
                        raise BADASError(500, event.get("error", "Job failed"))
                    if event.get("type") == "done":
                        return
                    yield event
                # Stream closed without "done" — gateway timeout, reconnect.
            except BADASError:
                raise
            except Exception:
                pass
            time.sleep(1.0)


# ── Helpers ──────────────────────────────────────────────────────────────────

def _with_progress(gen: Iterator[Dict], desc: str = "Analyzing") -> Iterator[Dict]:
    """Wrap a prediction generator with a tqdm bar that adapts to notebook / CLI."""
    try:
        from tqdm.auto import tqdm
    except ImportError:
        yield from gen
        return
    with tqdm(gen, unit="win", desc=desc) as pbar:
        for pred in pbar:
            pbar.set_postfix(
                prob=f'{pred.get("probability", 0):.2f}',
                risk=pred.get("risk_level", ""),
            )
            yield pred


def _parse_sse(resp) -> Iterator[tuple]:
    """Parse an SSE response stream, yielding ``(event_id, data_dict)`` tuples."""
    current_id   = None
    current_data = None
    for line in resp.iter_lines(decode_unicode=True):
        if not line:
            if current_data is not None:
                try:
                    yield current_id, json.loads(current_data)
                except json.JSONDecodeError:
                    pass
                current_data = None
            continue
        if line.startswith("id:"):
            current_id = line[3:].strip()
        elif line.startswith("data:"):
            current_data = line[5:].strip()
