"""Slices AI Infrastructure Service Job API."""

from collections.abc import Collection
from datetime import datetime, timedelta
from enum import Enum
from functools import cached_property
from typing import Any
from uuid import UUID

from httpx import Client, HTTPError
from msgspec.json import Decoder as JsonDecoder
from slices_clientlib_core.base.config import ApiEntry
from slices_clientlib_core.base.exceptions import from_httpx_exception
from slices_clientlib_core.client import (
    SlicesClient as SlicesCoreClient,
)
from slices_clientlib_core.client import (
    SlicesExperimentClient as SlicesCoreExperimentClient,
)
from slices_clientlib_core.client import (
    SlicesProjectClient as SlicesCoreProjectClient,
)
from slices_clientlib_core.util.pagination import PagedCollection

from .aiis_job_debug_logs import AIJobDebugLogsClient
from .aiis_job_output import AIJobOutputClient
from .schema.job import Job, JobStatus
from .schema.new_job import NewJob, NewJobReq
from .util.msgspec import json_encoder

_job_decoder = JsonDecoder(Job)


class AIJobClient:
    """AI Infrastructure Service Job API Client."""

    API_BASE_PATH = "/apis/ai.slices-ri.eu/v1/"

    def __init__(
        self,
        site_id: str,
        api_entry: ApiEntry,
        client: Client,
        slices_client: SlicesCoreProjectClient | SlicesCoreExperimentClient,
        slices_root_client: SlicesCoreClient,
    ) -> None:
        """Create AI Infrastructure Service Jobs Client."""
        self._site_id = site_id
        self._api_entry = api_entry
        self._client = client
        self._slices_client = slices_client
        self._slices_root_client = slices_root_client

        self._base_path = f"{self.API_BASE_PATH}jobs/"

        assert client.base_url is not None, "Expected base_url to be set"

    @cached_property
    def output(self) -> AIJobOutputClient:
        """Job Output Client."""
        return AIJobOutputClient(self._client)

    @cached_property
    def logs(self) -> AIJobDebugLogsClient:
        """Job Debug Logs Client."""
        return AIJobDebugLogsClient(self._client)

    def get(self, job: Job | str | UUID) -> Job:
        """Get job by ID."""
        job_id = job.uuid if isinstance(job, Job) else str(job)
        response = self._client.get(f"{self._base_path}{job_id}")

        try:
            response.raise_for_status()
            return _job_decoder.decode(response.content)
        except HTTPError as err:
            raise from_httpx_exception(err) from err

    def update_min_duration(self, job: Job | str | UUID, new_min_duration: str | timedelta):
        """Update job minDuration. Note that not all changes will be allowed."""
        job_id = job.uuid if isinstance(job, Job) else str(job)
        dur = (
            f"{new_min_duration.total_seconds()} s"
            if isinstance(new_min_duration, timedelta)
            else str(new_min_duration)
        )

        response = self._client.put(
            f"{self._base_path}{job_id}/request/scheduling/minDuration",
            headers={"Content-Type": "text/plain"},
            content=dur.encode("utf-8"),
        )

        try:
            response.raise_for_status()
        except HTTPError as err:
            raise from_httpx_exception(err) from err

    def update_max_duration(self, job: Job | str | UUID, new_max_duration: str | timedelta):
        """Update job maxDuration. Note that not all changes will be allowed."""

        if not isinstance(self._slices_client, SlicesCoreExperimentClient):
            raise RuntimeError("Can only change Job maxDuration within an experiment context.")  # noqa: TRY004

        job_id = job.uuid if isinstance(job, Job) else str(job)
        dur = (
            f"{new_max_duration.total_seconds()} s"
            if isinstance(new_max_duration, timedelta)
            else str(new_max_duration)
        )

        response = self._client.put(
            f"{self._base_path}{job_id}/request/scheduling/maxDuration",
            headers={"Content-Type": "text/plain"},
            content=dur.encode("utf-8"),
        )

        try:
            response.raise_for_status()
        except HTTPError as err:
            raise from_httpx_exception(err) from err

    def update_not_after(self, job: Job | str | UUID, new_not_after: datetime):
        """Update job notAfter."""

        if not isinstance(self._slices_client, SlicesCoreExperimentClient):
            raise RuntimeError("Can only change Job notAfter within an experiment context.")  # noqa: TRY004

        job_id = job.uuid if isinstance(job, Job) else str(job)

        response = self._client.put(
            f"{self._base_path}{job_id}/request/scheduling/notAfter",
            headers={"Content-Type": "text/plain"},
            content=new_not_after.isoformat(sep="T", timespec="seconds").encode("utf-8"),
        )

        try:
            response.raise_for_status()
        except HTTPError as err:
            raise from_httpx_exception(err) from err

    def _update_status(self, job: Job | str | UUID, new_status: JobStatus):
        """Update job status. Only certain values allowed."""
        job_id = job.uuid if isinstance(job, Job) else str(job)
        response = self._client.put(
            f"{self._base_path}{job_id}/state/status",
            headers={"Content-Type": "text/plain"},
            content=str(new_status.name).encode("utf-8"),
        )

        try:
            response.raise_for_status()
        except HTTPError as err:
            raise from_httpx_exception(err) from err

    def cancel(self, job: Job | str | UUID):
        """Cancel a job."""
        job_id = job.uuid if isinstance(job, Job) else str(job)
        self._update_status(job_id, JobStatus.CANCELLED)

    def halt(self, job: Job | str | UUID):
        """Manually HALT a job."""
        job_id = job.uuid if isinstance(job, Job) else str(job)
        self._update_status(job_id, JobStatus.MUSTHALT)

    def hold(self, job: Job | str | UUID):
        """Hold a QUEUED job so it cannot run."""
        job_id = job.uuid if isinstance(job, Job) else str(job)
        self._update_status(job_id, JobStatus.ONHOLD)

    def requeue(self, job: Job | str | UUID):
        """Requeue a job that is ONHOLD."""
        job_id = job.uuid if isinstance(job, Job) else str(job)
        self._update_status(job_id, JobStatus.QUEUED)

    def delete(self, job: Job | str | UUID) -> Job:
        """Delete compute job."""
        job_id = job.uuid if isinstance(job, Job) else str(job)
        response = self._client.delete(f"{self._base_path}{job_id}")

        try:
            response.raise_for_status()
            return _job_decoder.decode(response.content)
        except HTTPError as err:
            raise from_httpx_exception(err) from err

    def list(  # noqa: PLR0912, PLR0913, C901 we support all search options, and there are a lot.
        self,
        *,
        cluster_id: None | list[int] | int = None,
        allowed_states: None | list[JobStatus | str] | JobStatus | str = None,
        user_urn: None | list[str] | str = None,
        user_id: None | list[str] | str = None,
        user_name: None | list[str] | str = None,
        project_urn: None | list[str] | str = None,
        project_id: None | list[str] | str = None,
        project_name: None | list[str] | str = None,
        assigned_slave_name: None | list[str] | str = None,
        assigned_slave_instance_id: None | list[str] | str = None,
        interactive: bool | None = None,
        waste: bool | None = None,
        reservation_id: None | str | UUID = None,
    ) -> Collection[Job]:
        """List jobs.

        Be careful, this returns ALL jobs, there's no limit!
        They are returned as a "PagedCollection", which means they will only fetch jobs as required.
        You can use len() to get the total number of jobs, without requesting all of them.
        """

        def _list_or_single(v: list[Any] | Any) -> str:
            if isinstance(v, list):
                return ",".join([str(i) for i in v])
            return str(v)

        def _list_or_single_enum(v: list[Enum] | Enum) -> str:
            if isinstance(v, list):
                return ",".join([i.name for i in v])
            return v.name

        def _list_or_single_jobstatus_or_str(v: list[JobStatus | str] | JobStatus | str) -> str:
            if isinstance(v, str):
                return JobStatus(v).name
            if isinstance(v, JobStatus):
                return v.name
            if isinstance(v, list):
                return ",".join([i.name if isinstance(i, JobStatus) else JobStatus(str(i)).name for i in v])
            raise TypeError(f"unsupported type {type(v)}")

        params = {}
        if cluster_id is not None:
            params["cluster_id"] = _list_or_single(cluster_id)
        if allowed_states is not None:
            params["allowed_states"] = _list_or_single_jobstatus_or_str(allowed_states)
        if user_urn is not None:
            params["user_urn"] = _list_or_single(user_urn)
        if user_id is not None:
            params["user_id"] = _list_or_single(user_id)
        if user_name is not None:
            params["user_name"] = _list_or_single(user_name)
        if project_urn is not None:
            params["project_urn"] = _list_or_single(project_urn)
        if project_id is not None:
            params["project_id"] = _list_or_single(project_id)
        if project_name is not None:
            params["project_name"] = _list_or_single(project_name)
        if assigned_slave_name is not None:
            params["assigned_slave_name"] = _list_or_single(assigned_slave_name)
        if assigned_slave_instance_id is not None:
            params["assigned_slave_instance_id"] = _list_or_single(assigned_slave_instance_id)
        if interactive is not None:
            params["interactive"] = "true" if interactive else "false"
        if waste is not None:
            params["waste"] = "true" if waste else "false"
        if reservation_id is not None:
            params["reservation_id"] = str(reservation_id)

        return PagedCollection("job", Job, self._base_path, self._client, params)

    def create(self, new_job: NewJob) -> Job:
        """Create new jobs."""

        if not isinstance(self._slices_client, SlicesCoreExperimentClient):
            raise RuntimeError("Can only create AI Jobs within an experiment context.")  # noqa: TRY004

        id_token = self._slices_root_client.get_id_token(self._api_entry.audience, self._slices_client.project_id)

        # note: self._client has authentication with experiment token setup
        response = self._client.post(
            self._base_path,
            headers={"Content-Type": "application/json"},
            content=json_encoder.encode(NewJobReq(job=new_job, idtoken=id_token)),
        )

        try:
            response.raise_for_status()
        except HTTPError as err:
            raise from_httpx_exception(err) from err

        return _job_decoder.decode(response.content)
