from __future__ import annotations
from typing import Optional, Dict, Any
from .jobs.backend import Backend


class Session:
    """Session object to manage IonQ sessions."""

    def __init__(
        self,
        backend: Backend,
        *,
        max_time: Optional[int] = None,
        max_cost: Optional[int] = None,
        max_jobs: Optional[int] = None,
        create_new: bool = True,
        session_id: Optional[str] = None,
    ) -> None:
        self._backend = backend
        self._client = backend.client
        self._orig_run = None
        self._session_id: str

        # Create or attach
        if create_new and not session_id:
            self._create_session(max_time, max_cost, max_jobs)
        elif session_id:
            self._session_id = session_id
            # lazily validate that the session exists - will raise if unknown
            self.details()
        else:
            raise ValueError("Either create_new must be True or a session_id supplied.")

    # Accessors
    @property
    def session_id(self) -> str:
        """Return the IonQ session UUID."""
        return self._session_id

    def details(self) -> Dict[str, Any]:
        """Return JSON for this session."""
        return self._client.request(f"/sessions/{self._session_id}", "GET")

    def status(self) -> Optional[str]:
        """Return the status of this session."""
        return self.details().get("state")

    def usage(self) -> Optional[float]:
        """Return the usage time of this session in seconds."""
        return self.details().get("usage_time")

    # Lifecycle
    def cancel(self) -> None:
        """Cancel all queued jobs inside this session."""
        jobs_resp = self._client.request(
            "/jobs",
            "GET",
            params={"session_id": self._session_id, "status": "queued"},
        )
        jobs = jobs_resp.get("jobs", []) or []
        job_ids = [j["id"] for j in jobs if "id" in j]

        if not job_ids:
            return

        # Prefer to reuse SDK logic
        jobs_api = getattr(self._client, "jobs", None)
        if jobs_api is not None and hasattr(jobs_api, "cancel_job"):
            for job_id in job_ids:
                jobs_api.cancel_job(job_id)
        else:
            # Fallback: hit the REST API directly
            for job_id in job_ids:
                self._client.request(
                    f"/jobs/{job_id}/status/cancel",
                    "PUT",
                )

    def close(self) -> None:
        """End this session."""
        self._client.request(f"/sessions/{self._session_id}/end", "POST")

    # Context manager
    def __enter__(self) -> "Session":
        # inject the session_id into any backend.run call
        assert self._orig_run is None  # only set once per context
        self._orig_run = self._backend.run

        def _run_with_session(*args, **kwargs):
            # ensure session_id is set
            kwargs.setdefault("session_id", self._session_id)
            return self._orig_run(*args, **kwargs)

        # monkey-patch for the life of the context
        self._backend.run = _run_with_session
        return self

    def __exit__(self, exc_type, *_):
        # On exception try to cancel queued jobs before closing the session.
        if exc_type is not None:
            self.cancel()
        self.close()

        # restore the backend.run we overwrote in __enter__
        if self._orig_run is not None:  # defensive
            self._backend.run = self._orig_run
            self._orig_run = None

        # propagate exceptions
        return False

    # Helpers
    def _create_session(
        self,
        max_time: Optional[int],
        max_cost: Optional[int],
        max_jobs: Optional[int],
    ) -> None:
        """Create a new session via POST /sessions."""
        payload: Dict[str, Any] = {
            # python-ionq Backend names are already IonQ backend ids,
            # e.g. "simulator", "qpu.aria-1", etc.
            "backend": self._backend.name,
            "settings": {},
        }

        if max_jobs is not None:
            payload["settings"]["job_count_limit"] = max_jobs
        if max_time is not None:
            payload["settings"]["duration_limit_min"] = max_time
        if max_cost is not None:
            payload["settings"]["cost_limit"] = {"unit": "usd", "value": max_cost}

        resp = self._client.request("/sessions", "POST", json=payload)
        self._session_id = resp["id"]

    # class-method shortcut
    @classmethod
    def from_id(cls, session_id: str, *, backend: Backend) -> "Session":
        """Create a Session object from an existing session ID."""
        return cls(backend=backend, session_id=session_id, create_new=False)

    def run(self, *args, **kwargs):
        """Run a job using the session."""
        kwargs.setdefault("session_id", self._session_id)
        return self._backend.run(*args, **kwargs)


__all__ = ["Session"]
