"""
Lightweight session context for message and artifact management.

Aligned with the refactored backend session and agent message APIs.
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from .client import PopingClient


class Context:
    """
    Manage session messages and artifacts via backend endpoints.

    This mirrors the reference SDK's Context surface but routes to the
    new backend session APIs (and agent message append endpoint).
    """

    def __init__(
        self,
        *,
        client: Optional[PopingClient] = None,
        agent_id: Optional[str] = None,
        session_id: Optional[str] = None,
        project_id: Optional[str] = None,
        end_user_id: Optional[str] = None,
    ):
        self.client = client or self._get_default_client()
        self.agent_id = agent_id
        self.session_id = session_id
        self.project_id = project_id
        self.end_user_id = end_user_id

    def _get_default_client(self) -> PopingClient:
        try:
            from . import get_client  # type: ignore

            return get_client()
        except Exception as exc:  # pragma: no cover - defensive fallback
            raise RuntimeError("Poping client not configured; call poping.set(...) first.") from exc

    def _ensure_session(self) -> str:
        if not self.session_id:
            meta = self.client.create_session(project_id=self.project_id, end_user_id=self.end_user_id)
            self.session_id = meta.get("session_id")
        if not self.session_id:
            raise ValueError("Session ID is required but could not be created.")
        return self.session_id

    # ------------------------------------------------------------------
    # Message CRUD
    # ------------------------------------------------------------------
    def add(self, role: str, content: str) -> int:
        """
        Append a message to the session using the agent message append API.

        Returns:
            The zero-based message index after append (best effort).
        """
        if not self.agent_id:
            raise ValueError("agent_id is required to append messages.")
        sid = self._ensure_session()
        self.client.append_session_message(
            self.agent_id,
            sid,
            role=role,
            content=content,
            project_id=self.project_id,
            end_user_id=self.end_user_id,
        )
        messages = self.list()
        return len(messages) - 1 if messages else 0

    def list(self) -> List[Dict[str, Any]]:
        sid = self._ensure_session()
        return self.client.list_messages(sid, project_id=self.project_id, end_user_id=self.end_user_id)

    def get(self, key: int) -> Dict[str, Any]:
        sid = self._ensure_session()
        return self.client.get_message(sid, key, project_id=self.project_id, end_user_id=self.end_user_id)

    def delete(self, key: int) -> Dict[str, Any]:
        sid = self._ensure_session()
        return self.client.delete_message(sid, key, project_id=self.project_id, end_user_id=self.end_user_id)

    def token_usage(self) -> Dict[str, Any]:
        sid = self._ensure_session()
        return self.client.token_usage(sid, project_id=self.project_id, end_user_id=self.end_user_id)

    # ------------------------------------------------------------------
    # Artifacts
    # ------------------------------------------------------------------
    def list_artifacts(self) -> Dict[str, Any]:
        sid = self._ensure_session()
        return self.client.list_artifacts(sid, project_id=self.project_id, end_user_id=self.end_user_id)

    def upload_artifact(self, file_path: str) -> Dict[str, Any]:
        sid = self._ensure_session()
        return self.client.upload_artifact(sid, file_path, project_id=self.project_id, end_user_id=self.end_user_id)

    def delete_artifacts(self, filename: Optional[str] = None) -> Dict[str, Any]:
        sid = self._ensure_session()
        params = {"filename": filename, "project_id": self.project_id, "end_user_id": self.end_user_id}
        return self.client.delete_artifacts(sid, **{k: v for k, v in params.items() if v is not None})

    def grouped_artifacts(self) -> Dict[str, Any]:
        sid = self._ensure_session()
        return self.client.grouped_artifacts(sid, project_id=self.project_id, end_user_id=self.end_user_id)

    def download_artifact(self, filename: str, dest_path: str) -> str:
        sid = self._ensure_session()
        data = self.client.download_artifact(sid, filename, project_id=self.project_id, end_user_id=self.end_user_id)
        with open(dest_path, "wb") as f:
            f.write(data)
        return dest_path


__all__ = ["Context"]
