from attrs import field, frozen
from datetime import datetime
from metafold.api import asdatetime, asdict
from metafold.assets import Asset
from metafold.client import Client
from metafold.exceptions import PollTimeout
from requests import Response
from typing import Any, Optional
import time


def _assets(v: list[dict[str, Any] | Asset]) -> list[Asset]:
    return [a if isinstance(a, Asset) else Asset(**a) for a in v]


@frozen(kw_only=True)
class Job:
    """Job resource.

    Attributes:
        id: Job ID.
        name: Job name.
        type: Job type.
        parameters: Job parameters.
        created: Job creation datetime.
        state: Job state. May be one of: pending, started, success, or failure.
        assets: List of generated asset resources.
        meta: Additional metadata generated by the job.
    """
    id: str
    name: Optional[str] = None
    type: str
    parameters: dict[str, Any]
    created: datetime = field(converter=asdatetime)
    state: str
    assets: list[Asset] = field(converter=_assets)
    meta: dict[str, Any]


class JobsEndpoint:
    """Metafold jobs endpoint."""

    def __init__(self, client: Client) -> None:
        self._client = client

    def list(self, sort: Optional[str] = None, q: Optional[str] = None) -> list[Job]:
        """List jobs.

        Args:
            sort: Sort string. For details on syntax see the Metafold API docs.
                Supported sorting fields are: "id", "name", or "created".
            q: Query string. For details on syntax see the Metafold API docs.
                Supported search fields are: "id", "name", "type", and "state".

        Returns:
            List of job resources.
        """
        url = f"/projects/{self._client.project_id}/jobs"
        payload = asdict(sort=sort, q=q)
        r: Response = self._client.get(url, params=payload)
        return [Job(**j) for j in r.json()]

    def get(self, id: str) -> Job:
        """Get a job.

        Args:
            id: ID of job to get.

        Returns:
            Job resource.
        """
        url = f"/projects/{self._client.project_id}/jobs/{id}"
        r: Response = self._client.get(url)
        return Job(**r.json())

    def run(
        self, type: str, params: dict[str, Any],
        name: Optional[str] = None,
        timeout: int | float = 120,
    ) -> Job:
        """Dispatch a new job and wait for a result.

        See Metafold API docs for the full list of jobs.

        Args:
            type: Job type.
            params: Job parameters.
            name: Optional job name.
            timeout: Time in seconds to wait for a result.

        Returns:
            Completed job resource.
        """
        payload = asdict(type=type, parameters=params, name=name)
        r: Response = self._client.post(f"/projects/{self._client.project_id}/jobs", json=payload)
        r = self._poll(r.json()["link"], timeout)
        return Job(**r.json())

    def update(self, id: str, name: Optional[str] = None) -> Job:
        """Update a job.

        Args:
            id: ID of job to update.
            name: New job name. The existing name remains unchanged if None.

        Returns:
            Updated job resource.
        """
        url = f"/projects/{self._client.project_id}/jobs/{id}"
        payload = asdict(name=name)
        r: Response = self._client.patch(url, data=payload)
        return Job(**r.json())

    def _poll(self, url: str, timeout: int | float) -> Response:
        t0 = time.monotonic()
        r = self._client.get(url)
        while r.status_code == 202:
            elapsed = time.monotonic() - t0
            if elapsed >= timeout:
                raise PollTimeout("Job timed out")
            time.sleep(1)
            r = self._client.get(url)
        return r
