from typing import Any, Dict, Optional
from urllib.parse import urljoin

import httpx

from owslib import constants, context, services
from owslib.enums import Environment


class BaseOwsClient:
    def __init__(
        self, environment: Environment, service_name: str, *, retries: int = 5
    ) -> None:
        self.environment = environment
        self.service_name = service_name
        self.retries = retries

    def prepare_url(self, service_name: str, *, path: str) -> str:
        environment, service_url = services.discover_service_url(
            environment=self.environment,
            service_name=service_name,
        )
        return urljoin(service_url, path)

    @staticmethod
    def prepare_headers(
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
    ) -> Dict[str, str]:
        headers = headers or {}

        # Pass Correlation-Id header
        correlation_id = correlation_id or context.correlation_id.get()
        if correlation_id:
            headers[constants.HEADER_CORRELATION_ID] = correlation_id

        # Pass Authorization header
        request_context = context.request_context.get()
        if request_context and request_context.authorization:
            headers.setdefault("Authorization", request_context.authorization)

        return headers


class OwsClient(BaseOwsClient):
    def request(
        self,
        service_name: str,
        method: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        with httpx.Client(
            transport=httpx.HTTPTransport(retries=self.retries)
        ) as client:
            return client.request(
                method=method,
                url=self.prepare_url(service_name, path=path),
                headers=self.prepare_headers(
                    headers=headers, correlation_id=correlation_id
                ),
                **kwargs,
            )

    def head(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return self.request(
            service_name,
            method="HEAD",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    def get(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return self.request(
            service_name,
            method="GET",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    def post(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return self.request(
            service_name,
            method="POST",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    def put(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return self.request(
            service_name,
            method="PUT",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    def patch(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return self.request(
            service_name,
            method="PATCH",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    def delete(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return self.request(
            service_name,
            method="DELETE",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )


class AsyncOwsClient(BaseOwsClient):
    def __init__(
        self, environment: Environment, service_name: str, *, retries: int = 5
    ) -> None:
        super().__init__(environment, service_name, retries=retries)
        self._client = httpx.AsyncClient(
            transport=httpx.AsyncHTTPTransport(retries=self.retries)
        )

    async def close(self) -> None:
        await self._client.aclose()

    async def request(
        self,
        service_name: str,
        method: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return await self._client.request(
            method=method,
            url=self.prepare_url(service_name, path=path),
            headers=self.prepare_headers(
                headers=headers,
                correlation_id=correlation_id,
            ),
            **kwargs,
        )

    async def head(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return await self.request(
            service_name,
            method="HEAD",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    async def get(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return await self.request(
            service_name,
            method="GET",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    async def post(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return await self.request(
            service_name,
            method="POST",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    async def put(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return await self.request(
            service_name,
            method="PUT",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    async def patch(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return await self.request(
            service_name,
            method="PATCH",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    async def delete(
        self,
        service_name: str,
        path: str,
        *,
        headers: Optional[Dict[str, str]] = None,
        correlation_id: Optional[str] = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return await self.request(
            service_name,
            method="DELETE",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )