# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import base64
import inspect
import json
import logging
from collections.abc import Awaitable, Callable
from typing import Any, Literal, TypeVar, Union, overload
from uuid import uuid4

import httpx
from agent_framework._telemetry import get_user_agent
from agent_framework.observability import get_tracer
from azure.core.credentials import TokenCredential
from azure.core.credentials_async import AsyncTokenCredential
from opentelemetry import trace

from ._exceptions import (
    PurviewAuthenticationError,
    PurviewPaymentRequiredError,
    PurviewRateLimitError,
    PurviewRequestError,
    PurviewServiceError,
)
from ._models import (
    ContentActivitiesRequest,
    ContentActivitiesResponse,
    ProcessContentRequest,
    ProcessContentResponse,
    ProtectionScopesRequest,
    ProtectionScopesResponse,
)
from ._settings import PurviewSettings, get_purview_scopes

AzureCredentialTypes = Union[TokenCredential, AsyncTokenCredential]
AzureTokenProvider = Callable[[], Union[str, Awaitable[str]]]

logger = logging.getLogger("agent_framework.purview")

ResponseT = TypeVar("ResponseT")


class PurviewClient:
    """Async client for calling Graph Purview endpoints.

    Supports synchronous TokenCredential, asynchronous AsyncTokenCredential,
    or callable token providers. A sync credential will be invoked in a thread
    to avoid blocking the event loop.
    """

    def __init__(
        self,
        credential: AzureCredentialTypes | AzureTokenProvider,
        settings: PurviewSettings,
        *,
        timeout: float | None = 10.0,
    ):
        self._credential: AzureCredentialTypes | AzureTokenProvider = credential
        self._settings = settings
        self._graph_uri = (settings.get("graph_base_uri") or "https://graph.microsoft.com/v1.0/").rstrip("/")
        self._timeout = timeout
        self._client = httpx.AsyncClient(timeout=timeout)

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

    async def _get_token(self, *, tenant_id: str | None = None) -> str:
        """Acquire an access token using either async or sync credential, or callable token provider."""
        cred = self._credential
        # Callable token provider — returns a token string directly
        if callable(cred) and not isinstance(cred, (TokenCredential, AsyncTokenCredential)):
            result = cred()
            return await result if inspect.isawaitable(result) else result
        scopes = get_purview_scopes(self._settings)
        token = cred.get_token(*scopes, tenant_id=tenant_id)
        token = await token if inspect.isawaitable(token) else token
        return token.token

    @staticmethod
    def _extract_token_info(token: str) -> dict[str, Any]:
        parts = token.split(".")
        if len(parts) < 2:
            raise ValueError("Invalid JWT token format")
        payload = parts[1]
        rem = len(payload) % 4
        if rem:
            payload += "=" * (4 - rem)
        decoded = base64.urlsafe_b64decode(payload)
        data = json.loads(decoded.decode("utf-8"))
        return {
            "user_id": data.get("oid") if data.get("idtyp") == "user" else None,
            "tenant_id": data.get("tid"),
            "client_id": data.get("appid"),
        }

    async def get_user_info_from_token(self, *, tenant_id: str | None = None) -> dict[str, Any]:
        token = await self._get_token(tenant_id=tenant_id)
        return self._extract_token_info(token)

    async def process_content(self, request: ProcessContentRequest) -> ProcessContentResponse:
        with get_tracer().start_as_current_span("purview.process_content"):
            token = await self._get_token(tenant_id=request.tenant_id)
            url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/processContent"
            headers: dict[str, str] = {}
            # Add If-None-Match header if scope_identifier is present
            if hasattr(request, "scope_identifier") and request.scope_identifier:
                headers["If-None-Match"] = request.scope_identifier
            # Add Prefer: evaluateInline header if process_inline is True
            if hasattr(request, "process_inline") and request.process_inline:
                headers["Prefer"] = "evaluateInline"

            response: ProcessContentResponse | tuple[ProcessContentResponse, httpx.Headers] = await self._post(
                url, request, ProcessContentResponse, token, headers=headers, return_response=True
            )

            if isinstance(response, tuple) and len(response) == 2:
                response_obj, _ = response
                return response_obj

            return response

    async def get_protection_scopes(self, request: ProtectionScopesRequest) -> ProtectionScopesResponse:
        with get_tracer().start_as_current_span("purview.get_protection_scopes"):
            token = await self._get_token()
            url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/protectionScopes/compute"
            response: ProtectionScopesResponse | tuple[ProtectionScopesResponse, httpx.Headers] = await self._post(
                url, request, ProtectionScopesResponse, token, return_response=True
            )

            # Extract etag from response headers
            if isinstance(response, tuple) and len(response) == 2:
                response_obj, headers = response
                if "etag" in headers:
                    etag_value = headers["etag"].strip('"')
                    response_obj.scope_identifier = etag_value
                return response_obj

            return response

    async def send_content_activities(self, request: ContentActivitiesRequest) -> ContentActivitiesResponse:
        with get_tracer().start_as_current_span("purview.send_content_activities"):
            token = await self._get_token()
            url = f"{self._graph_uri}/users/{request.user_id}/dataSecurityAndGovernance/activities/contentActivities"
            return await self._post(url, request, ContentActivitiesResponse, token)

    @overload
    async def _post(
        self,
        url: str,
        model: Any,
        response_type: type[ResponseT],
        token: str,
        headers: dict[str, str] | None = None,
        return_response: Literal[False] = False,
    ) -> ResponseT: ...

    @overload
    async def _post(
        self,
        url: str,
        model: Any,
        response_type: type[ResponseT],
        token: str,
        headers: dict[str, str] | None = None,
        return_response: Literal[True] = True,
    ) -> tuple[ResponseT, httpx.Headers]: ...

    async def _post(
        self,
        url: str,
        model: Any,
        response_type: type[ResponseT],
        token: str,
        headers: dict[str, str] | None = None,
        return_response: bool = False,
    ) -> ResponseT | tuple[ResponseT, httpx.Headers]:
        if hasattr(model, "correlation_id") and not model.correlation_id:
            model.correlation_id = str(uuid4())

        correlation_id = getattr(model, "correlation_id", None)
        if correlation_id:
            span = trace.get_current_span()
            if span and span.is_recording():
                span.set_attribute("correlation_id", correlation_id)
            logger.info(f"Purview request to {url} with correlation_id: {correlation_id}")

        payload = model.model_dump(by_alias=True, exclude_none=True, mode="json")
        request_headers = {
            "Authorization": f"Bearer {token}",
            "User-Agent": get_user_agent(),
            "Content-Type": "application/json",
        }
        if correlation_id:
            request_headers["client-request-id"] = correlation_id

        if headers:
            request_headers.update(headers)
        resp = await self._client.post(url, json=payload, headers=request_headers)

        if resp.status_code in (401, 403):
            raise PurviewAuthenticationError(f"Auth failure {resp.status_code}: {resp.text}")
        if resp.status_code == 402:
            if self._settings.get("ignore_payment_required", False):
                return response_type()
            raise PurviewPaymentRequiredError(f"Payment required {resp.status_code}: {resp.text}")
        if resp.status_code == 429:
            raise PurviewRateLimitError(f"Rate limited {resp.status_code}: {resp.text}")
        if resp.status_code not in (200, 201, 202):
            raise PurviewRequestError(f"Purview request failed {resp.status_code}: {resp.text}")
        try:
            data = resp.json()
        except ValueError:
            data = {}

        try:
            # Prefer pydantic-style model_validate if present, else fall back to constructor.
            model_validate = getattr(response_type, "model_validate", None)
            response_obj = model_validate(data) if callable(model_validate) else response_type(**data)

            # Extract correlation_id from response headers if response object supports it
            if "client-request-id" in resp.headers and hasattr(response_obj, "correlation_id"):
                response_correlation_id = resp.headers["client-request-id"]
                response_obj.correlation_id = response_correlation_id  # pyright: ignore[reportAttributeAccessIssue]
                logger.info(f"Purview response from {url} with correlation_id: {response_correlation_id}")

            typed_response_obj = response_obj if isinstance(response_obj, response_type) else response_type(**data)
            if return_response:
                return (typed_response_obj, resp.headers)
            return typed_response_obj
        except Exception as ex:
            raise PurviewServiceError(f"Failed to deserialize Purview response: {ex}") from ex
