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

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}
API_KEY_CACHE = {}

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


def fetch_valid_public_keys(url):
    """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
    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 Exception as e:
        raise RuntimeError(f"Failed to fetch valid public keys: {e}")


def compute_public_key(secret_key):
    """Compute the public key from the secret key using argon2id."""
    try:
        # We don't need a fixed salt - when verifying, we extract the salt from the public key
        # The public key already contains all the information needed to verify
        return PH.hash(secret_key)
    except HashingError as e:
        raise RuntimeError(f"Failed to compute public key: {e}")


def apikey_login(func=None, *, public_keys_url=None, service=None):
    """
    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)
    """
    # 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")
    # Default URL using JSONBin.io
    DEFAULT_URL = "https://api.jsonbin.io/v3/b/691ec6a543b1c97be9b8ea6d"

    # 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)

    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=True,
            help='Your API secret key'
        )

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

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

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

        # Fast fail validations
        # 1. Check if secret key is empty
        if not secret_key:
            raise PermissionError("API key validation failed: API key cannot be empty")

        # 2. 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 API_KEY_CACHE:
            expiration_time = API_KEY_CACHE[cache_key]
            if current_time < expiration_time:
                # Cache hit: key-service pair is still valid
                return func(*args, **kwargs)

        # Cache miss or expired: proceed with full validation
        # Fetch valid public keys
        url = public_keys_url or DEFAULT_URL
        all_valid_keys = fetch_valid_public_keys(url)

        # 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 = 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 PH.verify(public_key, secret_key):
                    valid = True
                    break
            except Exception:
                # If verification fails for any reason (wrong key, corrupted hash, etc.), 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
        API_KEY_CACHE[cache_key] = current_time + CACHE_TTL_SECONDS

        try:
            # 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()
