"""Mocks for testing services using the Safir UWS support."""

from __future__ import annotations

import asyncio
import json
from collections import defaultdict
from datetime import UTC, datetime
from urllib.parse import parse_qs

import respx
from httpx import AsyncClient, Request, Response
from vo_models.uws.types import ExecutionPhase

from safir.arq import JobMetadata, JobResult, MockArqQueue
from safir.arq.uws import WorkerResult
from safir.datetime import current_datetime, parse_isodatetime
from safir.uws import (
    Job,
    JobCreate,
    JobUpdateAborted,
    JobUpdateCompleted,
    JobUpdateError,
    JobUpdateExecuting,
    JobUpdateMetadata,
    JobUpdateQueued,
    SerializedJob,
    UWSConfig,
)
from safir.uws._storage import JobStore

__all__ = [
    "MockUWSJobRunner",
    "MockWobbly",
    "patch_wobbly",
]


class MockWobbly:
    """Mock the Wobbly web service, which stores UWS job information.

    Use of this mock web service requires presentation of a token generated by
    the `make_token` class method, which encodes username and service
    information that would normally be taken from HTTP headers after
    Gafaelfawr processing.

    Attributes
    ----------
    jobs
        Stored jobs, organized by service, username, and then job ID.
    """

    @staticmethod
    def make_token(service: str, username: str) -> str:
        """Create a fake internal token for Wobbly calls.

        Parameters
        ----------
        service
            Service name encoded in fake internal token.
        username
            Username encoded in fake internal token.

        Returns
        -------
        str
            Fake internal token.
        """
        return f"gt-{service}.{username}"

    def __init__(self) -> None:
        # Next available job ID.
        self._job_id = 1

        # Maps service to username to job ID to a job record.
        self.jobs: defaultdict[str, defaultdict[str, dict[str, SerializedJob]]]
        self.jobs = defaultdict(lambda: defaultdict(dict))

    def create_job(self, request: Request) -> Response:
        """Create a new job record."""
        service, username = self._get_auth(request)
        job_create = JobCreate.model_validate_json(request.content)
        job_id = str(self._job_id)
        self._job_id += 1
        job = SerializedJob(
            id=job_id,
            service=service,
            owner=username,
            phase=ExecutionPhase.PENDING,
            creation_time=current_datetime(),
            **job_create.model_dump(),
        )
        self.jobs[service][username][job_id] = job
        return Response(
            201,
            json=job.model_dump(mode="json"),
            headers={"Location": str(request.url) + f"/{job_id}"},
        )

    def delete_job(self, request: Request, *, job_id: str) -> Response:
        """Delete a job."""
        service, username = self._get_auth(request)
        if job_id not in self.jobs[service][username]:
            return Response(404)
        del self.jobs[service][username][job_id]
        return Response(204)

    def get_job(self, request: Request, *, job_id: str) -> Response:
        """Retrieve a job."""
        service, username = self._get_auth(request)
        job = self.jobs[service][username].get(job_id)
        if not job:
            return Response(404)
        return Response(200, json=job.model_dump(mode="json"))

    def list_jobs(self, request: Request) -> Response:
        """List jobs matching the search parameters.

        Cursors are not supported. ``limit`` is, but does not result in a
        paginated response.
        """
        service, username = self._get_auth(request)

        # Parse query.
        query = parse_qs(request.url.query.decode())
        phases = set()
        if "phase" in query:
            phases = set(query["phase"])
        since = None
        if "since" in query:
            since = parse_isodatetime(query["since"][0])

        # Perform the search.
        results = []
        for job in self.jobs[service][username].values():
            if phases and job.phase not in phases:
                continue
            if since and job.creation_time <= since:
                continue
            results.append(job)

        # Sort the results and limit them if needed.
        json_results = [
            j.model_dump(mode="json")
            for j in sorted(
                results, key=lambda j: (j.creation_time, j.id), reverse=True
            )
        ]
        if "limit" in query:
            limit = int(query["limit"][0])
            if len(json_results) > limit:
                json_results = json_results[:limit]

        # Return the response.
        return Response(200, json=json_results)

    def update_job(self, request: Request, *, job_id: str) -> Response:
        """Make an update to a job record."""
        service, username = self._get_auth(request)
        job = self.jobs[service][username].get(job_id)
        if not job:
            return Response(404)
        body = json.loads(request.content)

        # First handle the only case without a phase, which is a metadata
        # update.
        if "phase" not in body or not body["phase"]:
            update = JobUpdateMetadata.model_validate(body)
            job.destruction_time = update.destruction_time
            job.execution_duration = update.execution_duration
            return Response(200, json=job.model_dump(mode="json"))

        # Otherwise, handle all the phase modification cases.
        match body["phase"]:
            case ExecutionPhase.ABORTED:
                _ = JobUpdateAborted.model_validate(body)
                if job.start_time:
                    job.end_time = current_datetime()
            case ExecutionPhase.COMPLETED:
                completed_update = JobUpdateCompleted.model_validate(body)
                job.end_time = current_datetime()
                job.results = completed_update.results
            case ExecutionPhase.EXECUTING:
                executing_update = JobUpdateExecuting.model_validate(body)
                start_time = executing_update.start_time.replace(microsecond=0)
                job.start_time = start_time
            case ExecutionPhase.ERROR:
                error_update = JobUpdateError.model_validate(body)
                job.end_time = current_datetime()
                job.errors = error_update.errors
            case ExecutionPhase.QUEUED:
                queued_update = JobUpdateQueued.model_validate(body)
                job.message_id = queued_update.message_id
        job.phase = ExecutionPhase[body["phase"]]
        return Response(200, json=job.model_dump(mode="json"))

    def _get_auth(self, request: Request) -> tuple[str, str]:
        """Parse the fake internal token into service and username."""
        auth_type, token = request.headers["Authorization"].split()
        assert auth_type.lower() == "bearer"
        assert token.startswith("gt-")
        service, username = token[3:].split(".", 1)
        return (service, username)


def patch_wobbly(respx_mock: respx.Router, wobbly_url: str) -> MockWobbly:
    """Set up the mock for a Wobbly server.

    Parameters
    ----------
    respx_mock
        Mock router.
    wobby_url
        Base URL on which Wobbly is "listening."

    Returns
    -------
    MockWobbly
        Mock Wobbly service.
    """
    wobbly_url = wobbly_url.rstrip("/")
    mock = MockWobbly()
    respx_mock.get(wobbly_url + "/jobs").mock(side_effect=mock.list_jobs)
    respx_mock.post(wobbly_url + "/jobs").mock(side_effect=mock.create_job)
    job_url = rf"{wobbly_url}/jobs/(?P<job_id>[^/]+)"
    respx_mock.get(url__regex=job_url).mock(side_effect=mock.get_job)
    respx_mock.delete(url__regex=job_url).mock(side_effect=mock.delete_job)
    respx_mock.patch(url__regex=job_url).mock(side_effect=mock.update_job)
    return mock


class MockUWSJobRunner:
    """Simulate execution of jobs with a mock queue.

    When running the test suite, the arq queue is replaced with a mock queue
    that doesn't execute workers. That execution has to be simulated by
    manually updating state in the mock queue and running the UWS database
    worker functions that normally would be run automatically by the queue.

    Parameters
    ----------
    config
        UWS configuration.
    arq_queue
        Mock arq queue for testing.
    """

    def __init__(self, config: UWSConfig, arq_queue: MockArqQueue) -> None:
        self._arq = arq_queue
        self._store = JobStore(config, AsyncClient())

    async def get_job_metadata(self, token: str, job_id: str) -> JobMetadata:
        """Get the arq job metadata for a job.

        Parameters
        ----------
        token
            Token for the user.
        job_id
            UWS job ID.

        Returns
        -------
        JobMetadata
            arq job metadata.
        """
        job = await self._store.get(token, job_id)
        assert job.message_id
        return await self._arq.get_job_metadata(job.message_id)

    async def get_job_result(self, token: str, job_id: str) -> JobResult:
        """Get the arq job result for a job.

        Parameters
        ----------
        token
            Token for the user.
        job_id
            UWS job ID.

        Returns
        -------
        JobMetadata
            arq job metadata.
        """
        job = await self._store.get(token, job_id)
        assert job.message_id
        return await self._arq.get_job_result(job.message_id)

    async def mark_in_progress(
        self, token: str, job_id: str, *, delay: float | None = None
    ) -> Job:
        """Mark a queued job in progress.

        Parameters
        ----------
        token
            Token for the user.
        job_id
            Job ID.
        delay
            How long to delay in seconds before marking the job as complete.

        Returns
        -------
        Job
            Record of the job.
        """
        if delay:
            await asyncio.sleep(delay)
        job = await self._store.get(token, job_id)
        assert job.message_id
        await self._arq.set_in_progress(job.message_id)
        await self._store.mark_executing(token, job_id, datetime.now(tz=UTC))
        return await self._store.get(token, job_id)

    async def mark_complete(
        self,
        token: str,
        job_id: str,
        results: list[WorkerResult] | Exception,
        *,
        delay: float | None = None,
    ) -> Job:
        """Mark an in progress job as complete.

        Parameters
        ----------
        token
            Token for the user.
        job_id
            Job ID.
        results
            Results to return. May be an exception to simulate a job failure.
        delay
            How long to delay in seconds before marking the job as complete.

        Returns
        -------
        Job
            Record of the job.
        """
        if delay:
            await asyncio.sleep(delay)
        job = await self._store.get(token, job_id)
        assert job.message_id
        await self._arq.set_complete(job.message_id, result=results)
        job_result = await self._arq.get_job_result(job.message_id)
        await self._store.mark_completed(token, job_id, job_result)
        return await self._store.get(token, job_id)
