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


def _assets(v: list[Union[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.
        finished: Job finished 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)
    finished: Optional[datetime] = field(
        converter=lambda v: optional(asdatetime)(v),
        default=None,
    )
    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,
        project_id: 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".
            project_id: Job project ID.

        Returns:
            List of job resources.
        """
        project_id = self._client.project_id(project_id)
        url = f"/projects/{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, job_id: str, project_id: Optional[str] = None) -> Job:
        """Get a job.

        Args:
            job_id: ID of job to get.
            project_id: Job project ID.

        Returns:
            Job resource.
        """
        project_id = self._client.project_id(project_id)
        url = f"/projects/{project_id}/jobs/{job_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: Union[int, float] = 120,
        project_id: Optional[str] = None,
    ) -> 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.
            project_id: Job project ID.

        Returns:
            Completed job resource.
        """
        url = self.run_status(type, params, name=name, project_id=project_id)
        try:
            r: Response = self.poll(url, timeout)
        except PollTimeout as e:
            raise RuntimeError(
                f"Job '{name or type}' failed to complete within {timeout} seconds"
            ) from e
        return Job(**r.json())

    def run_status(
        self, type: str, params: dict[str, Any],
        name: Optional[str] = None,
        project_id: Optional[str] = None,
    ) -> str:
        """Dispatch a new job and return immediately without waiting for result.

        See Metafold API docs for the full list of jobs.

        Args:
            type: Job type.
            params: Job parameters.
            name: Optional job name.
            project_id: Job project ID.

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

    def poll(
        self, url: str,
        timeout: Union[int, float] = 120,
        every: Union[int, float] = 1,
    ) -> Response:
        """Poll the given URL in regular intervals.

        Helpful for waiting on job results given a status URL.

        Args:
            timeout: Time in seconds to wait for a result.
            every: Frequency in seconds.

        Returns:
            HTTP 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

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

        Args:
            job_id: ID of job to update.
            name: New job name. The existing name remains unchanged if None.
            project_id: Job project ID.

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