import argparse
import urllib.request
import urllib.error
import json
import os
import time
import sys
from functools import wraps
from argon2 import PasswordHasher
from argon2.exceptions import HashingError
import urllib.error

# Try to import FastAPI types for middleware (optional dependency)
try:
    from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
    from starlette.requests import Request
    from starlette.responses import Response, JSONResponse
    from fastapi import FastAPI
    FASTAPI_AVAILABLE = True
except ImportError:
    # If FastAPI is not installed, middleware-related functionality will be disabled
    FASTAPI_AVAILABLE = False
    # Create mock classes for type hinting
    class Request:
        pass

    class Response:
        pass

    class JSONResponse:
        pass

    class BaseHTTPMiddleware:
        pass

    class RequestResponseEndpoint:
        pass

    class FastAPI:
        pass

# Singleton API Key Gateway instance
class APIKeyGateway:
    """Core API key gateway logic for validation and caching"""

    def __init__(self):
        self.PH = PasswordHasher(
            memory_cost=65536,
            time_cost=3,
            parallelism=4,
            hash_len=32,
            salt_len=16
            # argon2id is the default type in newer versions of argon2-cffi
        )

        # Cache to store valid API key-service pairs
        # Format: {(secret_key, service): expiration_time}
        self.API_KEY_CACHE = {}

        # Cache TTL: 1 day (86400 seconds)
        self.CACHE_TTL_SECONDS = 86400

        # Default URL using JSONBin.io
        self.DEFAULT_URL = "https://api.jsonbin.io/v3/b/691ec6a543b1c97be9b8ea6d"

    @staticmethod
    def retry_on_exception(exceptions, max_retries=3, delay=0.5, backoff=1):
        """
        Decorator that retries a function on specific exceptions.
        """
        def decorator(func):
            @wraps(func)
            def wrapper(*args, **kwargs):
                verbose = kwargs.get('verbose', True)
                func_kwargs = {k: v for k, v in kwargs.items() if k != 'verbose'}
                retries = 0
                while retries <= max_retries:
                    try:
                        return func(*args, **func_kwargs)
                    except exceptions as e:
                        retries += 1
                        if retries > max_retries:
                            if verbose:
                                print(f"   All {max_retries} retry attempts failed. Exception: {e}")
                            raise
                        wait_time = delay * (backoff ** (retries - 1))
                        if verbose:
                            print(f"   Retry {retries}/{max_retries} in {wait_time:.2f} seconds due to: {type(e).__name__}: {e}")
                        time.sleep(wait_time)
                return None  # This line should never be reached
            return wrapper
        return decorator

    @retry_on_exception(
        urllib.error.URLError,
        max_retries=3,
        delay=0.5,
        backoff=1  # constant delay, not exponential
    )
    def fetch_valid_public_keys(self, url, verbose=True):
        """Fetch the list of valid public keys from the specified URL."""
        try:
            # Set proper headers for JSONBin.io API
            headers = {
                'accept': 'application/json',
            }

            # Add JSONBin API key if available
            jsonbin_api_key = os.environ.get('JSONBIN_API_KEY')
            if jsonbin_api_key:
                headers['X-Access-Key'] = jsonbin_api_key

            # Create a request object with the headers
            req = urllib.request.Request(url, headers=headers)

            # Set timeout to avoid hanging (10 seconds)
            with urllib.request.urlopen(req, timeout=10) as response:
                data = json.loads(response.read().decode('utf-8'))

            # JSONBin API wraps the data in a 'record' field
            if isinstance(data, dict) and 'record' in data:
                data = data['record']
                all_keys = []

                # Expected structure: {service: {key_id: public_key, ...}}
                for service_name, service_data in data.items():
                    if isinstance(service_data, dict):
                        # Service has multiple API keys
                        for key_id, public_key in service_data.items():
                            all_keys.append({
                                "service": service_name,
                                "public_key": public_key
                            })
                    else:
                        # Service has a single API key
                        all_keys.append({
                            "service": service_name,
                            "public_key": service_data
                        })
                return all_keys
            raise RuntimeError(f"Unexpected response structure from JSONBin API: {data}")
        except urllib.error.HTTPError as e:
            if e.code == 401:
                # 401 Unauthorized typically means missing or invalid API key
                jsonbin_api_key = os.environ.get('JSONBIN_API_KEY')
                if not jsonbin_api_key:
                    raise RuntimeError("Failed to fetch valid public keys from jsonbin.io: JSONBIN_API_KEY environment variable is missing")
                else:
                    raise RuntimeError(f"Failed to fetch valid public keys from jsonbin.io: Invalid API key provided (HTTP 401). Check your JSONBIN_API_KEY environment variable.")
            raise RuntimeError(f"Failed to fetch valid public keys from jsonbin.io: HTTP {e.code} error - {e.reason}")
        except urllib.error.URLError:
            # Re-raise URLError so the retry decorator can handle it
            raise
        except Exception as e:
            # Wrap all other exceptions in RuntimeError
            raise RuntimeError(f"Failed to fetch valid public keys: {e}")

    def compute_public_key(self, secret_key):
        """Compute the public key from the secret key using argon2id."""
        try:
            return self.PH.hash(secret_key)
        except HashingError as e:
            raise RuntimeError(f"Failed to compute public key: {e}")

    def validate_api_key(self, secret_key, service, public_keys_url=None, verbose=True):
        """Validate an API key for a specific service"""
        # Fast fail validations
        if not secret_key:
            raise PermissionError("API key validation failed: API key cannot be empty")

        # Check if secret key has correct length
        expected_length = 43  # Based on test secret key: LVYYllZk6xLIRoGVfGK9A78Dl1ehJbZOpRo6JoF2LAM
        if len(secret_key) != expected_length:
            raise PermissionError(f"API key validation failed: API key must be {expected_length} characters long")

        # Check cache first
        cache_key = (secret_key, service)
        current_time = time.time()

        if cache_key in self.API_KEY_CACHE:
            expiration_time = self.API_KEY_CACHE[cache_key]
            if current_time < expiration_time:
                # Cache hit: key-service pair is still valid
                if verbose:
                    print(f"✓ API key validation passed (cached until {time.ctime(expiration_time)})")
                return True

        # Cache miss or expired: proceed with full validation
        if verbose:
            print(f"⟳ Checking API key validity for service '{service}'...")

        # Fetch valid public keys
        url = public_keys_url or self.DEFAULT_URL
        if verbose:
            print(f"   Fetching valid public keys from {url}...")

        all_valid_keys = self.fetch_valid_public_keys(url, verbose=verbose)

        # Filter keys by the specified service
        valid_public_keys = [key['public_key'] for key in all_valid_keys if key['service'] == service]

        # Check if service exists
        if not valid_public_keys:
            # Check if the service name exists in the fetched keys
            available_services = set(key['service'] for key in all_valid_keys)
            if service not in available_services:
                raise PermissionError(
                    f"API key validation failed: Service '{service}' does not exist. Available services: {', '.join(available_services) if available_services else 'None'}"
                )

        # Compute public key from secret key
        try:
            computed_public_key = self.compute_public_key(secret_key)
        except RuntimeError as e:
            raise RuntimeError(f"API key validation failed: {e}")

        # Check if any of the valid public keys match the secret key
        valid = False
        for public_key in valid_public_keys:
            try:
                if self.PH.verify(public_key, secret_key):
                    valid = True
                    break
            except Exception:
                # If verification fails for any reason, try next one
                continue

        if not valid:
            raise PermissionError(
                "API key validation failed: API key is not correct for the specified service. Please apply for a valid API key from yanru@cyanru.com"
            )

        # Update cache with valid key-service pair
        self.API_KEY_CACHE[cache_key] = current_time + self.CACHE_TTL_SECONDS
        expiration_time = self.API_KEY_CACHE[cache_key]

        if verbose:
            print(f"✓ API key validation passed for service '{service}'")
            print(f"   Caching result until {time.ctime(expiration_time)}")

        return True

# Create singleton instance
gateway = APIKeyGateway()


class APIKeyGatewayMiddleware(BaseHTTPMiddleware):
    """
    FastAPI middleware for API key gateway validation.

    Best practice:
    - Validates API keys for all incoming requests to protected routes
    - Extracts API key from configurable headers (default: X-API-Key, Authorization Bearer)
    - Supports service-specific validation
    - Handles proper error responses with appropriate status codes
    """

    def __init__(
        self,
        app: FastAPI,
        service: str,
        public_keys_url: str | None = None,
        api_key_headers: list[str] | None = None,
        verbose: bool = True
    ):
        """
        Initialize the API key gateway middleware.

        :param app: FastAPI app instance
        :param service: Service name to validate against (required)
        :param public_keys_url: Custom URL for public keys (defaults to gateway's DEFAULT_URL)
        :param api_key_headers: List of headers to check for API key (default: X-API-Key, Authorization Bearer)
        :param verbose: Verbose logging
        """
        super().__init__(app)
        self.service = service
        self.public_keys_url = public_keys_url
        self.api_key_headers = api_key_headers or [
            "x-api-key",
            "authorization"  # for Bearer token format
        ]
        self.verbose = verbose
        self.gateway_instance = gateway

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        """
        Middleware dispatch method that handles API key validation.
        """
        secret_key = None

        # Extract API key from configured headers
        headers = request.headers
        for header_name in self.api_key_headers:
            header_value = headers.get(header_name)
            if header_value:
                if header_name.lower() == "authorization":
                    # Handle Bearer token format
                    if header_value.lower().startswith("bearer "):
                        secret_key = header_value[7:].strip()
                        break
                else:
                    # Handle direct header value
                    secret_key = header_value.strip()
                    break

        try:
            # Validate the API key using the gateway instance
            if secret_key:
                self.gateway_instance.validate_api_key(
                    secret_key=secret_key,
                    service=self.service,
                    public_keys_url=self.public_keys_url,
                    verbose=self.verbose
                )
            else:
                # No API key found in headers
                raise PermissionError("API key validation failed: No API key found in request headers")

            # If validation passes, call the next endpoint
            response = await call_next(request)
            return response

        except (PermissionError, RuntimeError) as e:
            # Return appropriate error response
            return JSONResponse(
                status_code=401 if isinstance(e, PermissionError) else 500,
                content={"error": str(e)}
            )


# Function to add middleware to FastAPI app (convenience wrapper)
def add_apikey_gateway_middleware(
    app: FastAPI,
    service: str,
    public_keys_url: str | None = None,
    api_key_headers: list[str] | None = None,
    verbose: bool = True
):
    """
    Convenience function to add APIKeyGatewayMiddleware to a FastAPI app.

    :param app: FastAPI app instance
    :param service: Service name to validate against
    :param public_keys_url: Custom URL for public keys
    :param api_key_headers: List of headers to check for API key
    :param verbose: Verbose logging
    """
    if not FASTAPI_AVAILABLE:
        raise RuntimeError("FastAPI is not installed. Please install it with 'pip install fastapi'.")

    # Add middleware directly to the app
    app.add_middleware(
        APIKeyGatewayMiddleware,
        service=service,
        public_keys_url=public_keys_url,
        api_key_headers=api_key_headers,
        verbose=verbose
    )

# Expose some constants and helper functions for backward compatibility
CACHE_TTL_SECONDS = gateway.CACHE_TTL_SECONDS
API_KEY_CACHE = gateway.API_KEY_CACHE
PH = gateway.PH


# Expose the retry_on_exception decorator from the gateway for backward compatibility
def retry_on_exception(exceptions, max_retries=3, delay=0.5, backoff=1):
    """
    Decorator that retries a function on specific exceptions.

    :param exceptions: Tuple of exceptions to retry on
    :param max_retries: Maximum number of retries (default: 3)
    :param delay: Initial delay between retries in seconds (default: 0.5)
    :param backoff: Factor by which delay increases each retry (default: 1 for constant delay)
    """
    return APIKeyGateway.retry_on_exception(exceptions, max_retries, delay, backoff)

# Expose fetch_valid_public_keys from gateway for backward compatibility
def fetch_valid_public_keys(url, verbose=True):
    """Fetch the list of valid public keys from the specified URL."""
    return gateway.fetch_valid_public_keys(url, verbose=verbose)

# Expose compute_public_key from gateway for backward compatibility
def compute_public_key(secret_key):
    """Compute the public key from the secret key using argon2id."""
    return gateway.compute_public_key(secret_key)

def apikey_login(func=None, *, public_keys_url=None, service=None, verbose=True):
    """
    Decorator that adds API key validation to a function.

    :param func: The function to decorate (required if using as decorator without params)
    :param public_keys_url: Optional URL to fetch valid public keys from (defaults to the official URL)
    :param service: The service name to validate the API key against (required)
    :param verbose: Whether to print verbose messages (default: True)
    """
    # Check if service parameter is provided
    if func is None:
        # If we're in the parameterized decorator mode, we'll check service later
        pass
    elif service is None:
        raise ValueError("service parameter is required for apikey_login decorator")

    # Handle both decorator with and without parameters
    if func is None:
        # Return a new decorator with the provided parameters
        return lambda f: apikey_login(f, public_keys_url=public_keys_url, service=service, verbose=verbose)

    def wrapper(*args, **kwargs):
        # Create argument parser
        parser = argparse.ArgumentParser(add_help=False)  # Don't add help to avoid conflict
        parser.add_argument(
            '-a', '--apikey',
            required=False,
            help='Your API secret key (fallbacks to AKGATEWAY_API_KEY env var)'
        )

        # Save original argv to restore later
        original_argv = sys.argv.copy()

        # Parse only API key arguments
        args_parsed, remaining_args = parser.parse_known_args()

        # Get API key from command line or environment variable
        secret_key = args_parsed.apikey or os.environ.get('AKGATEWAY_API_KEY')

        # If no API key provided, raise an error
        if not secret_key:
            parser.error('API key is required. Use -a/--apikey or set AKGATEWAY_API_KEY environment variable')

        # Replace sys.argv with remaining args (excluding API key) for the wrapped function
        sys.argv = [sys.argv[0]] + remaining_args

        try:
            # Use the gateway singleton to validate the API key
            gateway.validate_api_key(secret_key, service, public_keys_url=public_keys_url, verbose=verbose)

            # If validation passes, call the decorated function
            return func(*args, **kwargs)
        finally:
            # Always restore original argv
            sys.argv = original_argv

    return wrapper


# Example usage
if __name__ == "__main__":
    @apikey_login(service="media-match")
    def app():
        print("API key validated successfully! Welcome to the app.")

    app()
