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


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


AssetDict: TypeAlias = dict[str, Union[dict[str, Any], Asset]]


def _assets_dict(v: AssetDict) -> dict[str, Asset]:
    return dict(
        (k, a if isinstance(a, Asset) else Asset(**a))
        for k, a in v.items()
    )


class IODict(TypedDict):
    params: Optional[dict[str, Any]]
    assets: Optional[dict[str, AssetDict]]


@frozen(kw_only=True)
class IO:
    """Job input/output.

    Attributes:
        params: JSON-encoded parameter values.
        assets: Related assets.
    """
    params: Optional[dict[str, Any]] = None
    assets: Optional[dict[str, Asset]] = field(
        converter=lambda v: optional(_assets_dict)(v), default=None)

    @staticmethod
    def from_dict(d: IODict) -> "IO":
        return IO(params=d.get("params"), assets=d.get("assets"))


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

    Attributes:
        id: Job ID.
        name: Job name.
        type: Job type.
        state: Job state. May be one of: pending, started, success, failure, or
            canceled.
        created: Job creation datetime.
        started: Job started datetime.
        finished: Job finished datetime.
        error: Error message for failed jobs.
        inputs: Input assets and parameters.
        outputs: Output assets and parameters.
        needs: List of upstream job IDs in a workflow graph.
        project_id: Project ID.
        workflow_id: Workflow ID.
        assets: (Deprecated) List of generated asset resources.
        parameters: (Deprecated) Job parameters.
        meta: (Deprecated) Additional metadata generated by the job.
    """
    id: str
    name: Optional[str] = None
    type: str
    state: str
    created: datetime = field(converter=asdatetime)
    started: Optional[datetime] = field(
        converter=lambda v: optional_datetime(v), default=None)
    finished: Optional[datetime] = field(
        converter=lambda v: optional_datetime(v), default=None)
    error: Optional[str] = None
    inputs: IO = field(converter=lambda v: v if isinstance(v, IO) else IO.from_dict(v))
    outputs: IO = field(converter=lambda v: v if isinstance(v, IO) else IO.from_dict(v))
    needs: list[str]
    project_id: Optional[str] = None
    workflow_id: Optional[str] = None
    # NOTE(ryan): Deprecated
    assets: Optional[list[Asset]] = field(
        converter=lambda v: optional(_assets)(v), default=None)
    parameters: dict[str, Any]
    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", "created", "started", or
                "finished".
            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, *args, **kwargs):
        return self._client.poll(*args, **kwargs)

    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())

    def delete(self, job_id: str, project_id: Optional[str] = None):
        """Delete a job.

        Args:
            job_id: ID of job to delete.
            project_id: Job project ID.
        """
        project_id = self._client.project_id(project_id)
        url = f"/projects/{project_id}/jobs/{job_id}"
        self._client.delete(url)
