from collections.abc import AsyncIterable, Iterable
from typing import Any, cast
from uuid import UUID

from ..models import (
    BatchStatusResponse,
    JobActionResponse,
    JobErrorsResponse,
    JobID,
    JobLogResponse,
    JobSpec,
    JobStatus,
    LogSource,
    TasksBaseUrl,
)
from ..util import event_stream, event_stream_async
from .base import BaseClient


class Jobs(BaseClient):
    def get_base_url(self) -> TasksBaseUrl:
        res = self.request("POST", endpoint="get-base-url", return_model=TasksBaseUrl)
        return cast(TasksBaseUrl, res)

    async def get_base_url_async(self) -> TasksBaseUrl:
        res = await self.request_async(
            "POST", endpoint="get-base-url", return_model=TasksBaseUrl
        )
        return cast(TasksBaseUrl, res)

    def start_job(self, body: JobID | JobSpec) -> JobActionResponse:
        res = self.request(
            "POST", endpoint="start-job", data=body, return_model=JobActionResponse
        )
        return cast(JobActionResponse, res)

    async def start_job_async(self, body: JobID | JobSpec) -> JobActionResponse:
        res = await self.request_async(
            "POST", endpoint="start-job", data=body, return_model=JobActionResponse
        )
        return cast(JobActionResponse, res)

    def stop_job(self, body: JobID) -> JobActionResponse:
        res = self.request(
            "POST", endpoint="stop-job", data=body, return_model=JobActionResponse
        )
        return cast(JobActionResponse, res)

    async def stop_job_async(self, body: JobID) -> JobActionResponse:
        res = await self.request_async(
            "POST", endpoint="stop-job", data=body, return_model=JobActionResponse
        )
        return cast(JobActionResponse, res)

    def delete_job(self, body: JobID) -> JobActionResponse:
        res = self.request(
            "POST", endpoint="delete-job", data=body, return_model=JobActionResponse
        )
        return cast(JobActionResponse, res)

    async def delete_job_async(self, body: JobID) -> JobActionResponse:
        res = await self.request_async(
            "POST", endpoint="delete-job", data=body, return_model=JobActionResponse
        )
        return cast(JobActionResponse, res)

    def get_status(self, job_id: UUID) -> JobStatus:
        res = self.request("GET", endpoint=f"{job_id}/status", return_model=JobStatus)
        return cast(JobStatus, res)

    async def get_status_async(self, job_id: UUID) -> JobStatus:
        res = await self.request_async(
            "GET", endpoint=f"{job_id}/status", return_model=JobStatus
        )
        return cast(JobStatus, res)

    def get_batch_status(self, job_ids: list[UUID]) -> BatchStatusResponse:
        res = self.request(
            "POST",
            endpoint="status",
            data={"job_ids": [str(jid) for jid in job_ids]},
            return_model=BatchStatusResponse,
        )
        return cast(BatchStatusResponse, res)

    async def get_batch_status_async(self, job_ids: list[UUID]) -> BatchStatusResponse:
        res = await self.request_async(
            "POST",
            endpoint="status",
            data={"job_ids": [str(jid) for jid in job_ids]},
            return_model=BatchStatusResponse,
        )
        return cast(BatchStatusResponse, res)

    def logs(
        self,
        job_id: UUID,
        *,
        tail_lines: int | None = 200,
        since_seconds: int | None = None,
        source: LogSource | None = None,
        query: str | None = None,
    ) -> str:
        params: dict[str, Any] = {"tail_lines": tail_lines}
        if since_seconds is not None:
            params["since_seconds"] = since_seconds
        if source is not None:
            params["source"] = source.value
        if query is not None:
            params["query"] = query
        res = self.request(
            "GET",
            endpoint=f"{job_id}/logs",
            params=params,
            return_model=JobLogResponse,
        )
        assert isinstance(res, JobLogResponse)
        return res.text or ""

    async def logs_async(
        self,
        job_id: UUID,
        *,
        tail_lines: int | None = 200,
        since_seconds: int | None = None,
        source: LogSource | None = None,
        query: str | None = None,
    ) -> str:
        params: dict[str, Any] = {"tail_lines": tail_lines}
        if since_seconds is not None:
            params["since_seconds"] = since_seconds
        if source is not None:
            params["source"] = source.value
        if query is not None:
            params["query"] = query
        res = await self.request_async(
            "GET",
            endpoint=f"{job_id}/logs",
            params=params,
            return_model=JobLogResponse,
        )
        assert isinstance(res, JobLogResponse)
        return res.text or ""

    def errors(self, job_id: UUID) -> JobErrorsResponse:
        res = self.request(
            "GET",
            endpoint=f"{job_id}/errors",
            return_model=JobErrorsResponse,
        )
        return cast(JobErrorsResponse, res)

    async def errors_async(self, job_id: UUID) -> JobErrorsResponse:
        res = await self.request_async(
            "GET",
            endpoint=f"{job_id}/errors",
            return_model=JobErrorsResponse,
        )
        return cast(JobErrorsResponse, res)

    def events(
        self,
        job_id: UUID,
        index: int = 0,
        timeout: int | None = None,
    ) -> Iterable[dict[str, Any]]:
        # No async 'yield from', so use the loop here for symmetry
        for item in event_stream(
            client=self._sync_client,
            method="GET",
            url=f"/api/v1/jobs/{job_id}/events",
            index=index,
            timeout=timeout,
        ):
            yield item

    async def events_async(
        self,
        job_id: UUID,
        index: int = 0,
        timeout: int | None = None,
    ) -> AsyncIterable[dict[str, Any]]:
        async for item in event_stream_async(
            client=self._async_client,
            method="GET",
            url=f"/api/v1/jobs/{job_id}/events",
            index=index,
            timeout=timeout,
        ):
            yield item

    def messages(
        self,
        job_id: UUID,
        index: int = 0,
        timeout: int | None = None,
    ) -> Iterable[dict[str, Any]]:
        # No async 'yield from', so use the loop here for symmetry
        sent_messages: set[str] = set()
        for batch in event_stream(
            client=self._sync_client,
            method="GET",
            url=f"/api/v1/jobs/{job_id}/messages",
            index=index,
            timeout=timeout,
        ):
            for msg in batch["Messages"]:
                if msg not in sent_messages:
                    yield msg
                    sent_messages.add(msg)

    async def messages_async(
        self,
        job_id: UUID,
        index: int = 0,
        timeout: int | None = None,
    ) -> AsyncIterable[dict[str, Any]]:
        sent_messages: set[str] = set()
        async for batch in event_stream_async(
            client=self._async_client,
            method="GET",
            url=f"/api/v1/jobs/{job_id}/messages",
            index=index,
            timeout=timeout,
        ):
            for msg in batch["Messages"]:
                if msg not in sent_messages:
                    yield msg
                    sent_messages.add(msg)
