import httpx
from httpx import Response, AsyncClient
from typing import Optional, Dict, Any
import base64
import asyncio
from datetime import datetime, timedelta
import json

class TokenGenerationError(Exception):
    """Custom exception for token generation errors."""
    pass

class GraphQLServiceManager:
    def __init__(self, config, tenant_id: str):
        graphql_config = json.loads(config.graphql_config)

        self.endpoint = graphql_config.get("endpoint")
        self.headers = graphql_config.get("headers", {"Content-Type": "application/json"})
        self.tenant_id = tenant_id
        self._client: AsyncClient = None
        
        # OAuth 2.0 configuration 
        self.client_id = graphql_config.get("client_id")
        self.client_secret = graphql_config.get("client_secret")
        self.token_url = graphql_config.get("token_url")
        self.scope = graphql_config.get("scope")
        
        # Token management
        self._access_token: Optional[str] = None
        self._token_expires_at: Optional[datetime] = None
        self._token_lock = asyncio.Lock()

    async def _generate_oauth_token(self) -> str:
        """
        Generates an OAuth 2.0 access token using client credentials flow.
        
        :return: The access token string.
        :raises ValueError: If OAuth configuration is missing.
        :raises httpx.HTTPStatusError: If token request fails.
        """
        if not self.is_oauth_configured():
            raise ValueError("OAuth configuration incomplete. client_id, client_secret, and token_url are required.")
        
        # Create basic auth header
        credentials = f"{self.client_id}:{self.client_secret}"
        encoded_credentials = base64.b64encode(credentials.encode()).decode()
        
        headers = {
            "Authorization": f"Basic {encoded_credentials}",
            "Content-Type": "application/x-www-form-urlencoded"
        }
        
        # Prepare the token request data
        data = {
            "grant_type": "client_credentials"
        }
        
        if self.scope:
            data["scope"] = self.scope
        
        self._ensure_client()
        
        try:
            response = await self._client.post(
                self.token_url,
                headers=headers,
                data=data
            )
            response.raise_for_status()
            
            token_data = response.json()
            access_token = token_data.get("access_token")
            
            if not access_token:
                raise ValueError("No access_token found in token response")
            
            # Calculate expiration time (default to 1 hour if not provided)
            expires_in = token_data.get("expires_in", 3600)  # Default to 1 hour
            self._token_expires_at = datetime.now() + timedelta(seconds=expires_in - 60)  # Refresh 1 minute early
            
            return access_token
            
        except httpx.HTTPStatusError as e:
            raise httpx.HTTPStatusError(
                f"Token request failed: {e.response.status_code} - {e.response.text}",
                request=e.request,
                response=e.response
            )
        except Exception as e:
            raise TokenGenerationError("Failed to generate OAuth token") from e

    async def _get_valid_token(self) -> Optional[str]:
        """
        Returns a valid access token, generating a new one if necessary.
        
        :return: Valid access token or None if OAuth is not configured.
        """
        if not self.is_oauth_configured():
            raise ValueError("OAuth configuration incomplete. client_id, client_secret, and token_url are required.")
        
        async with self._token_lock:
            # Check if we need a new token
            if (not self._access_token or 
                not self._token_expires_at or 
                datetime.now() >= self._token_expires_at):
                
                self._access_token = await self._generate_oauth_token()
            
            return self._access_token

    async def refresh_token(self) -> str:
        """
        Manually refreshes the OAuth token.
        
        :return: New access token.
        :raises ValueError: If OAuth configuration is missing.
        """
        async with self._token_lock:
            self._access_token = None
            self._token_expires_at = None
            return await self._generate_oauth_token()

    def is_oauth_configured(self) -> bool:
        """
        Checks if OAuth 2.0 is properly configured.
        
        :return: True if OAuth is configured, False otherwise.
        """
        return all([self.client_id, self.client_secret, self.token_url])

    async def execute(
        self,
        query: str,
        variables: Optional[Dict[str, Any]] = None,
        operation_name: Optional[str] = None,
        method: str = "POST",
        tenant_id: Optional[str] = None
    ) -> Response:
        """
        Executes a GraphQL query against the specified endpoint.

        :param query: The GraphQL query string.
        :param variables: Optional dictionary of variables to be used in the query.
        :param operation_name: Optional name of the operation to execute.
        :param method: The HTTP method to use for the request, either 'POST' or 'GET'. Defaults to 'POST'.
        :param tenant_id: Optional tenant ID to override the default tenant ID.
        :return: A dictionary containing the response data from the GraphQL server.
        :raises ValueError: If an unsupported HTTP method is provided.
        """
        self._ensure_client()
        
        # Get OAuth token if configured and add to headers
        token = await self._get_valid_token()
        request_headers = self.headers.copy()
        
        if token:
            request_headers["Authorization"] = f"Bearer {token}"

        tenant_id = tenant_id or self.tenant_id
        url = f"{self.endpoint}/{tenant_id}/graphql-v2"
        payload = {"query": query}
        
        
        if variables is not None:
            payload["variables"] = variables
        if operation_name:
            payload["operationName"] = operation_name

        if method.upper() == "POST":
            response = await self._client.post(url, json=payload, headers=request_headers)
        elif method.upper() == "GET":
            response = await self._client.get(url, params=payload, headers=request_headers)
        else:
            raise ValueError(f"Unsupported HTTP method for GraphQL: {method}")

        return response

    def _ensure_client(self):
        if not self._client or self._client.is_closed:
            self._client = httpx.AsyncClient()

    async def aclose(self):
        if self._client:
            await self._client.aclose()
