"""OwsClient implementations."""
from __future__ import annotations

import weakref
from typing import Any, Callable
from urllib.parse import urljoin, urlparse
from uuid import uuid1

import httpx

from owsclient import constants
from owsclient.logging_adapter import get_logger
from owsclient.protocols import AsyncM2MTokenManager, M2MTokenManager, RequestContext
from owsclient.services import discover_service_url

# Registry for test discovery
_OWS_CLIENT_REGISTRY: weakref.WeakSet[OwsClient] = weakref.WeakSet()
_ASYNC_OWS_CLIENT_REGISTRY: weakref.WeakSet[AsyncOwsClient] = weakref.WeakSet()

logger = get_logger(__name__)


def make_new_correlation_id() -> str:
    """Make a new correlation id."""
    return str(uuid1())


class BaseOwsClient:
    """Common client functionality."""

    def __init__(
        self,
        environment: str,
        service_name: str,
        *,
        timeout: httpx.Timeout = httpx.Timeout(5.0),
        retries: int = 5,
        correlation_id_getter: Callable[[], str | None] | None = None,
        request_context_getter: Callable[[], RequestContext | None] | None = None,
    ) -> None:
        """Create base client."""
        self.environment = environment
        self.service_name = service_name
        self.timeout = timeout
        self.retries = retries
        self.correlation_id_getter = correlation_id_getter
        self.request_context_getter = request_context_getter

    def _get_request_context(self) -> RequestContext | None:
        """Get the request context if available."""
        return (
            self.request_context_getter()
            if self.request_context_getter is not None
            else None
        )

    def prepare_url(self, service_name: str, *, path: str) -> str:
        """Get the request URL."""
        environment, service_url = discover_service_url(
            environment=self.environment,
            service_name=service_name,
        )
        parse_url = urlparse(service_url)
        if parse_url.scheme not in ("http", "https"):
            # urljoin has troubling handling certain urls. Default to http protocol.
            service_url = f"http://{service_url}"
        url = urljoin(service_url, path)
        return url

    def pass_correlation_id_header(
        self,
        headers: dict[str, str],
        correlation_id: str | None,
    ) -> dict[str, str]:
        """Get and return the correlation id as part of the dict of headers.

        If a correlation_id is passed, use it.
        Else, if a get_correlation_id_func is passed, call it and return its value.
            If that value happens to be empty, make a new correlation id.
        Else, make a new correlation id.
        """
        correlation_id_header: dict[str, str] = {}
        correlation_id = (
            correlation_id
            or (
                self.correlation_id_getter()
                if self.correlation_id_getter
                else make_new_correlation_id()
            )
            or make_new_correlation_id()
        )
        if correlation_id:
            headers[constants.HEADER_CORRELATION_ID] = correlation_id
        return {**headers, **correlation_id_header}

    def pass_graphql_headers(
        self,
        headers: dict[str, str],
        identity_id: str | None,
        profile_id: int | None,
        profile_type: str | None,
    ) -> dict[str, str]:
        """Get headers required for making a graphql request."""
        request_context = self._get_request_context()
        if request_context:
            if not identity_id:
                identity_id = request_context.identity_id
            if not profile_id:
                profile_id = request_context.profile_id
            if not profile_type:
                profile_type = request_context.profile_type

        if not (identity_id and profile_id and profile_type):
            raise ValueError("identity_id, profile_id and profile_type should passed.")

        graphql_headers = {
            constants.HEADER_ORCHARD_IDENTITY_ID: identity_id,
            constants.HEADER_ORCHARD_PROFILE_TYPE: profile_type,
            constants.HEADER_ORCHARD_PROFILE_ID: str(profile_id),
            constants.HEADER_APOLLO_GRAPHQL_CLIENT_NAME: self.service_name,
        }

        return {**headers, **graphql_headers}

    @staticmethod
    def pass_graphql_body(
        query: str,
        operation_name: str | None = None,
        variables: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Get body required for making a graphql request."""
        return {
            "operationName": operation_name,
            "query": query.strip(),
            "variables": variables or {},
        }


class OwsClient(BaseOwsClient):
    """Client for synchronous operations."""

    def __init__(
        self,
        environment: str,
        service_name: str,
        *,
        timeout: httpx.Timeout = httpx.Timeout(5.0),
        retries: int = 5,
        m2m_token_manager: M2MTokenManager | None = None,
        correlation_id_getter: Callable[[], str | None] | None = None,
        request_context_getter: Callable[[], RequestContext | None] | None = None,
    ) -> None:
        """Construct a synchronous client."""
        super().__init__(
            environment,
            service_name,
            timeout=timeout,
            retries=retries,
            correlation_id_getter=correlation_id_getter,
            request_context_getter=request_context_getter,
        )
        self.m2m_token_manager = m2m_token_manager
        # Store OwsClient instances only for testing purposes
        _OWS_CLIENT_REGISTRY.add(self)

    def prepare_headers(
        self,
        *,
        headers: dict[str, str] | None = None,
        correlation_id: str | None = None,
    ) -> dict[str, str]:
        """Get all headers."""
        headers = headers or {}
        headers = self.pass_correlation_id_header(
            headers,
            correlation_id=correlation_id,
        )
        if "authorization" not in headers:
            headers = self.pass_authorization_header(
                headers,
            )
        return headers

    def pass_authorization_header(
        self,
        headers: dict[str, str],
    ) -> dict[str, str]:
        """Get Authorization header."""
        authorization_header = {}
        request_context = self._get_request_context()
        if request_context and request_context.authorization:
            authorization_header = {"authorization": request_context.authorization}
        elif self.m2m_token_manager:
            try:
                token = self.m2m_token_manager.get_token_string()
            except Exception as exc:
                logger.warning(
                    "Failed to get M2M token",
                    exc_info=exc,
                    extra={
                        "library": {
                            "name": "python-owsclient",
                            "language": "python",
                        },
                    },
                )
            else:
                authorization_header = {"authorization": f"Bearer {token}"}
        return {**headers, **authorization_header}

    def request(
        self,
        service_name: str,
        method: str,
        path: str,
        *,
        headers: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make a request to an ows service."""
        with httpx.Client(
            transport=httpx.HTTPTransport(retries=self.retries),
            timeout=self.timeout,
        ) 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: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make a HEAD request to an ows service."""
        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: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make a GET request to an ows service."""
        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: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make a POST request to an ows service."""
        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: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make a PUT request to an ows service."""
        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: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make a PATCH request to an ows service."""
        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: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make a DELETE request to an ows service."""
        return self.request(
            service_name,
            method="DELETE",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    def graphql_query(
        self,
        service_name: str,
        *,
        query: str,
        operation_name: str | None = None,
        variables: dict[str, Any] | None = None,
        headers: dict[str, str] | None = None,
        identity_id: str | None = None,
        profile_id: int | None = None,
        profile_type: str | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make a request to graphql."""
        headers = headers or {}

        kwargs["json"] = self.pass_graphql_body(
            query,
            operation_name=operation_name,
            variables=variables,
        )

        return self.post(
            service_name,
            path="/graphql",
            headers=self.pass_graphql_headers(
                headers,
                identity_id=identity_id,
                profile_id=profile_id,
                profile_type=profile_type,
            ),
            correlation_id=correlation_id,
            **kwargs,
        )


class AsyncOwsClient(BaseOwsClient):
    """Client for asynchronous operations."""

    def __init__(
        self,
        environment: str,
        service_name: str,
        *,
        timeout: httpx.Timeout = httpx.Timeout(5.0),
        retries: int = 5,
        m2m_token_manager: AsyncM2MTokenManager | None = None,
        correlation_id_getter: Callable[[], str | None] | None = None,
        request_context_getter: Callable[[], RequestContext | None] | None = None,
    ) -> None:
        """Construct an asynchronous client."""
        super().__init__(
            environment,
            service_name,
            timeout=timeout,
            retries=retries,
            correlation_id_getter=correlation_id_getter,
            request_context_getter=request_context_getter,
        )
        self.m2m_token_manager = m2m_token_manager
        self._client: httpx.AsyncClient | None = None
        _ASYNC_OWS_CLIENT_REGISTRY.add(self)

    @property
    def client(self) -> httpx.AsyncClient:
        """Get/set the underlying http client."""
        if self._client is None:
            self._client = httpx.AsyncClient(
                transport=httpx.AsyncHTTPTransport(retries=self.retries),
                timeout=self.timeout,
            )
        return self._client

    async def close(self) -> None:
        """Close the client connection."""
        await self.client.aclose()
        self._client = None

    async def prepare_headers(
        self,
        *,
        headers: dict[str, str] | None = None,
        correlation_id: str | None = None,
    ) -> dict[str, str]:
        """Get all headers."""
        headers = headers or {}
        headers = self.pass_correlation_id_header(
            headers,
            correlation_id=correlation_id,
        )
        if "authorization" not in headers:
            headers = await self.pass_authorization_header(
                headers,
            )
        return headers

    async def pass_authorization_header(
        self,
        headers: dict[str, str],
    ) -> dict[str, str]:
        """Get Authorization header."""
        authorization_header = {}
        request_context = self._get_request_context()
        if request_context and request_context.authorization:
            authorization_header = {"authorization": request_context.authorization}
        elif self.m2m_token_manager:
            try:
                token = await self.m2m_token_manager.get_token_string()
            except Exception as exc:
                logger.warning(
                    "Failed to get M2M token",
                    exc_info=exc,
                    extra={
                        "library": {
                            "name": "python-owsclient",
                            "language": "python",
                        },
                    },
                )
            else:
                authorization_header = {"authorization": f"Bearer {token}"}
        return {**headers, **authorization_header}

    async def request(
        self,
        service_name: str,
        method: str,
        path: str,
        *,
        headers: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make an async request to an ows service."""
        return await self.client.request(
            method=method,
            url=self.prepare_url(service_name, path=path),
            headers=await self.prepare_headers(
                headers=headers,
                correlation_id=correlation_id,
            ),
            **kwargs,
        )

    async def head(
        self,
        service_name: str,
        path: str,
        *,
        headers: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make an async HEAD request to an ows service."""
        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: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make an async GET request to an ows service."""
        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: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make an async POST request to an ows service."""
        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: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make an async PUT request to an ows service."""
        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: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make an async PATCH request to an ows service."""
        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: dict[str, str] | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make an async DELETE request to an ows service."""
        return await self.request(
            service_name,
            method="DELETE",
            path=path,
            headers=headers,
            correlation_id=correlation_id,
            **kwargs,
        )

    async def graphql_query(
        self,
        service_name: str,
        *,
        query: str,
        operation_name: str | None = None,
        variables: dict[str, Any] | None = None,
        headers: dict[str, str] | None = None,
        identity_id: str | None = None,
        profile_id: int | None = None,
        profile_type: str | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        """Make an async request to graphql."""
        headers = headers or {}

        kwargs["json"] = self.pass_graphql_body(
            query,
            operation_name=operation_name,
            variables=variables,
        )

        return await self.post(
            service_name,
            path="/graphql",
            headers=self.pass_graphql_headers(
                headers,
                identity_id=identity_id,
                profile_id=profile_id,
                profile_type=profile_type,
            ),
            correlation_id=correlation_id,
            **kwargs,
        )