"""Runs resource (includes checkpoints and resume)."""

from __future__ import annotations

import requests

from .._base import BaseResource
from ..types import (
    CancelRunResponse,
    CheckpointResponse,
    ExportConnectorSpec,
    ImportConnectorSpec,
    Run,
    RunCreate,
    RunResumeRequest,
    RunStatusResponse,
)


class RunsResource(BaseResource):

    def __init__(self, session: requests.Session, base_url: str, *, openai_api_key: str | None = None) -> None:
        super().__init__(session, base_url)
        self._openai_api_key = openai_api_key

    def create(
        self,
        agent_id: str,
        prompt: str,
        *,
        model: str = "gpt-5.2",
        openai_api_key: str | None = None,
        input_files: list[str] | None = None,
        timeout_seconds: int = 3600,
        import_connectors: list[ImportConnectorSpec] | None = None,
        export_connectors: list[ExportConnectorSpec] | None = None,
    ) -> Run:
        body = RunCreate(
            agent_id=agent_id,
            prompt=prompt,
            model=model,
            openai_api_key=openai_api_key or self._openai_api_key,
            input_files=input_files or [],
            timeout_seconds=timeout_seconds,
            import_connectors=import_connectors or [],
            export_connectors=export_connectors or [],
        )
        resp = self._session.post(self._url("/runs"), data=body.model_dump_json())
        self._raise_for_status(resp)
        return Run.model_validate(resp.json())

    def list(self, agent_id: str | None = None) -> list[Run]:
        params = {}
        if agent_id is not None:
            params["agent_id"] = agent_id
        resp = self._session.get(self._url("/runs"), params=params)
        self._raise_for_status(resp)
        return [Run.model_validate(r) for r in resp.json()]

    def get(self, run_id: str) -> Run:
        resp = self._session.get(self._url(f"/runs/{run_id}"))
        self._raise_for_status(resp)
        return Run.model_validate(resp.json())

    def get_status(self, run_id: str) -> RunStatusResponse:
        resp = self._session.get(self._url(f"/runs/{run_id}/status"))
        self._raise_for_status(resp)
        return RunStatusResponse.model_validate(resp.json())

    def cancel(self, run_id: str) -> CancelRunResponse:
        resp = self._session.post(self._url(f"/runs/{run_id}/cancel"))
        self._raise_for_status(resp)
        return CancelRunResponse.model_validate(resp.json())

    def get_checkpoint_chain(self, run_id: str) -> list[CheckpointResponse]:
        resp = self._session.get(self._url(f"/runs/{run_id}/checkpoint-chain"))
        self._raise_for_status(resp)
        return [CheckpointResponse.model_validate(c) for c in resp.json()]

    def get_checkpoint(self, run_id: str, checkpoint_id: str) -> CheckpointResponse:
        resp = self._session.get(self._url(f"/runs/{run_id}/checkpoints/{checkpoint_id}"))
        self._raise_for_status(resp)
        return CheckpointResponse.model_validate(resp.json())

    def resume(
        self,
        run_id: str,
        prompt: str,
        *,
        model: str = "gpt-5.2",
        openai_api_key: str | None = None,
        timeout_seconds: int = 3600,
        import_connectors: list[ImportConnectorSpec] | None = None,
        export_connectors: list[ExportConnectorSpec] | None = None,
    ) -> Run:
        body = RunResumeRequest(
            prompt=prompt,
            model=model,
            openai_api_key=openai_api_key or self._openai_api_key,
            timeout_seconds=timeout_seconds,
            import_connectors=import_connectors or [],
            export_connectors=export_connectors or [],
        )
        resp = self._session.post(self._url(f"/runs/{run_id}/resume"), data=body.model_dump_json())
        self._raise_for_status(resp)
        return Run.model_validate(resp.json())
