from __future__ import annotations

from collections.abc import AsyncIterator
from typing import Any

import httpx

from .downstream import DownstreamProtocolError, make_timeout
from .models import MediaJob, MediaJobRequest, MediaLease
from .sse import SSEEvent, parse_sse_lines


class MediaJobAdapter:
    """On-demand manager + media job contract. Results are artifact references only."""

    def __init__(
        self,
        manager_url: str,
        media_url: str,
        *,
        connect_timeout_s: float = 0.75,
        request_timeout_s: float = 20.0,
        stream_timeout_s: float = 300.0,
        client: httpx.AsyncClient | None = None,
    ) -> None:
        self.manager_url = manager_url.rstrip("/")
        self.media_url = media_url.rstrip("/")
        self._owns_client = client is None
        self.client = client or httpx.AsyncClient(
            timeout=make_timeout(connect_timeout_s, request_timeout_s, stream_timeout_s)
        )

    async def close(self) -> None:
        if self._owns_client:
            await self.client.aclose()

    async def acquire(self, kind: str, *, idempotency_key: str) -> MediaLease:
        response = await self.client.post(
            f"{self.manager_url}/v1/leases",
            headers={"Idempotency-Key": idempotency_key},
            json={"capability": kind},
        )
        response.raise_for_status()
        payload = response.json()
        lease_id = payload.get("lease_id") or payload.get("id")
        if not isinstance(lease_id, str) or not lease_id:
            raise DownstreamProtocolError("Model manager response has no lease_id")
        endpoint = payload.get("endpoint")
        return MediaLease(lease_id=lease_id, endpoint=endpoint if isinstance(endpoint, str) else None)

    async def release(self, lease_id: str) -> None:
        response = await self.client.delete(f"{self.manager_url}/v1/leases/{lease_id}")
        if response.status_code not in {200, 202, 204, 404}:
            response.raise_for_status()

    async def create_job(
        self,
        request: MediaJobRequest,
        *,
        lease: MediaLease,
        idempotency_key: str,
    ) -> MediaJob:
        base_url = (lease.endpoint or self.media_url).rstrip("/")
        response = await self.client.post(
            f"{base_url}/v1/jobs",
            headers={"Idempotency-Key": idempotency_key, "X-Model-Lease": lease.lease_id},
            json=request.model_dump(mode="json"),
        )
        response.raise_for_status()
        payload = response.json()
        job_id = payload.get("job_id") or payload.get("id")
        if not isinstance(job_id, str) or not job_id:
            raise DownstreamProtocolError("Media create-job response has no job_id")
        return MediaJob(job_id=job_id)

    async def events(
        self,
        job_id: str,
        *,
        endpoint: str | None = None,
        after_event_id: str | None = None,
    ) -> AsyncIterator[SSEEvent]:
        headers = {"Accept": "text/event-stream"}
        if after_event_id:
            headers["Last-Event-ID"] = after_event_id
        base_url = (endpoint or self.media_url).rstrip("/")
        async with self.client.stream(
            "GET", f"{base_url}/v1/jobs/{job_id}/events", headers=headers
        ) as response:
            response.raise_for_status()
            async for event in parse_sse_lines(response.aiter_lines()):
                yield event

    async def cancel(self, job_id: str, *, endpoint: str | None = None) -> None:
        base_url = (endpoint or self.media_url).rstrip("/")
        response = await self.client.post(f"{base_url}/v1/jobs/{job_id}/cancel")
        if response.status_code not in {200, 202, 204, 404, 409}:
            response.raise_for_status()

    async def health(self) -> dict[str, Any]:
        result: dict[str, Any] = {}
        for name, url in (("manager", self.manager_url), ("media", self.media_url)):
            try:
                response = await self.client.get(f"{url}/health")
                response.raise_for_status()
                result[name] = {"status": "ready", "detail": response.json()}
            except (httpx.HTTPError, ValueError) as exc:
                result[name] = {"status": "unavailable", "detail": type(exc).__name__}
        return result