"""Client implementation for workload identity authentication."""

import os
import time
import threading
import logging
from typing import Optional, Dict, Any
from datetime import datetime, timedelta

import requests

from .exceptions import ConfigurationError, TokenRetrievalError, TokenExchangeError

logger = logging.getLogger(__name__)


class TokenCache:
    """Thread-safe token cache with expiration support."""

    def __init__(self):
        self._cache: Dict[str, Any] = {}
        self._lock = threading.RLock()

    def get(self, key: str) -> Optional[str]:
        """Get token from cache if it exists and is not expired."""
        with self._lock:
            if key not in self._cache:
                return None

            token_info = self._cache[key]
            expires_at = token_info.get('expires_at')

            if expires_at and time.time() < expires_at:
                return token_info.get('token')
            else:
                # Token expired, remove from cache
                del self._cache[key]
                return None

    def set(self, key: str, token: str, expires_in: int):
        """Set token in cache with expiration time."""
        with self._lock:
            expires_at = time.time() + expires_in - 60  # 1 minute buffer
            self._cache[key] = {
                'token': token,
                'expires_at': expires_at
            }


class Client:
    """Workload Identity SDK Client."""

    def __init__(self):
        """Initialize the client with configuration from environment variables."""
        self._load_configuration()
        self._token_cache = TokenCache()
        self._session = requests.Session()

        # Configure session with reasonable timeouts
        self._session.timeout = 30

        logger.debug("Workload Identity Client initialized")

    def _load_configuration(self):
        """Load required configuration from environment variables."""
        required_vars = {
            'COZE_WORKLOAD_IDENTITY_CLIENT_ID': 'client_id',
            'COZE_WORKLOAD_IDENTITY_CLIENT_SECRET': 'client_secret',
            'COZE_WORKLOAD_IDENTITY_TOKEN_ENDPOINT': 'token_endpoint',
            'COZE_WORKLOAD_ACCESS_TOKEN_ENDPOINT': 'access_token_endpoint'
        }

        missing_vars = []
        for env_var, attr_name in required_vars.items():
            value = os.environ.get(env_var)
            if not value:
                missing_vars.append(env_var)
            else:
                setattr(self, f"_{attr_name}", value)

        if missing_vars:
            raise ConfigurationError(
                f"Missing required environment variables: {', '.join(missing_vars)}"
            )

        # Load lane configuration
        self._lane_env = os.environ.get('COZE_SERVER_ENV', 'NONE')
        logger.debug(f"Lane environment: {self._lane_env}")

    def _make_token_request(self, endpoint: str, data: Dict[str, str]) -> Dict[str, Any]:
        """Make a token request to the specified endpoint."""
        headers = {
            'Content-Type': 'application/x-www-form-urlencoded',
            'User-Agent': 'coze-workload-identity/0.1.0'
        }

        # Add lane headers if COZE_SERVER_ENV is set and not NONE
        if self._lane_env != 'NONE':
            headers['x-tt-env'] = self._lane_env
            if self._lane_env.startswith('ppe_'):
                headers['x-use-ppe'] = '1'
            elif self._lane_env.startswith('boe_'):
                # boe_ lanes don't need x-use-ppe header
                pass

        try:
            response = self._session.post(endpoint, data=data, headers=headers)
            response.raise_for_status()

            result = response.json()
            if 'error' in result:
                raise TokenRetrievalError(f"Token request failed: {result.get('error_description', result['error'])}")

            return result

        except requests.exceptions.RequestException as e:
            logger.error(f"Token request failed: {e}")
            raise TokenRetrievalError(f"Failed to retrieve token: {e}")
        except ValueError as e:
            logger.error(f"Failed to parse token response: {e}")
            raise TokenRetrievalError(f"Invalid token response format: {e}")

    def _get_id_token(self) -> str:
        """Get ID token using client credentials (no caching)."""
        logger.debug("Requesting new ID token")

        data = {
            'client_id': self._client_id,
            'client_secret': self._client_secret,
            'grant_type': 'client_credentials'
        }

        result = self._make_token_request(self._token_endpoint, data)

        id_token = result.get('access_token')
        if not id_token:
            raise TokenRetrievalError("No access_token found in ID token response")

        expires_in = result.get('expires_in', 3600)  # Default to 1 hour
        logger.debug(f"ID token retrieved, expires in {expires_in} seconds")
        return id_token

    def get_access_token(self) -> str:
        """
        Get access token using OAuth2.0 token exchange.

        Returns:
            str: The access token

        Raises:
            ConfigurationError: If required configuration is missing
            TokenRetrievalError: If token retrieval fails
            TokenExchangeError: If token exchange fails
        """
        cache_key = "access_token"
        cached_token = self._token_cache.get(cache_key)

        if cached_token:
            logger.debug("Using cached access token")
            return cached_token

        logger.debug("Requesting new access token")

        # Step 1: Get ID token
        id_token = self._get_id_token()

        # Step 2: Exchange ID token for access token
        data = {
            'client_id': self._client_id,
            'subject_token': id_token,
            'subject_token_type': 'urn:ietf:params:oauth:token-type:id_token',
            'grant_type': 'urn:ietf:params:oauth:grant-type:token-exchange'
        }

        result = self._make_token_request(self._access_token_endpoint, data)

        access_token = result.get('access_token')
        if not access_token:
            raise TokenExchangeError("No access_token found in access token response")

        expires_in = result.get('expires_in', 3600)  # Default to 1 hour
        self._token_cache.set(cache_key, access_token, expires_in)

        logger.debug(f"Access token retrieved, expires in {expires_in} seconds")
        return access_token

    def __enter__(self):
        """Context manager entry."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        self.close()

    def close(self):
        """Close the client and cleanup resources."""
        if hasattr(self, '_session'):
            self._session.close()
        logger.debug("Workload Identity Client closed")