"""Tests for the authorization header extractor."""

import os
import sys
from fastapi import HTTPException
import random
import string
import pytest


# Add root folder (parent) to sys.path
sys.path.append(
    os.path.dirname(
        os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    )
)

from gateway.common.authorization import extract_authorization_from_headers, INVARIANT_AUTHORIZATION_HEADER, API_KEYS_SEPARATOR

@pytest.mark.parametrize("push_to_explorer", [True, False])
@pytest.mark.parametrize("invariant_authorization", [True, False])
@pytest.mark.parametrize("invariant_authorization_appended_to_llm_provider_api_key", [True, False])
@pytest.mark.parametrize("use_fallback_header", [True, False])
def test_extract_authorization_from_headers(
    push_to_explorer: bool,
    invariant_authorization: bool,
    invariant_authorization_appended_to_llm_provider_api_key: bool,
    use_fallback_header: bool,
):
    """Test the extract_authorization_from_headers function."""

    llm_apikey = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
    inv_apikey = ''.join(random.choices(string.ascii_letters + string.digits, k=10))
    dataset_name = "test-dataset" if push_to_explorer else None

    headers: dict[str, str] = {} 

    llm_provider_api_key = "fallback-header" if use_fallback_header else "llm-provider-api-key"
    headers[llm_provider_api_key] = llm_apikey

    if invariant_authorization:
        print("invariant_authorization - TRUE")
        if invariant_authorization_appended_to_llm_provider_api_key:
            headers[llm_provider_api_key] = f"{headers.get(llm_provider_api_key, '')}{API_KEYS_SEPARATOR}{inv_apikey}"
        else:
            headers[INVARIANT_AUTHORIZATION_HEADER] = f"Bearer {inv_apikey}"

    # Mock request headers
    class MockRequest:
        def __init__(self, headers):
            self.headers = headers

    request = MockRequest(headers) 

    # Call the function
    try:
        invariant_auth, llm_provider_api_key = extract_authorization_from_headers(
            request,
            dataset_name=dataset_name,
            llm_provider_api_key_header="llm-provider-api-key",
            llm_provider_fallback_api_key_headers=["fallback-header"],
        )
    except HTTPException as e:
        # If an exception is raised, check if it is the expected one
        if not invariant_authorization:
            assert e.status_code == 400
            assert e.detail == "Missing invariant api key"
            return
        else:
            raise e
    # Verify the results
    if invariant_authorization:
        if not push_to_explorer and invariant_authorization_appended_to_llm_provider_api_key:
            assert llm_provider_api_key.split(API_KEYS_SEPARATOR)[0] == llm_apikey
            assert llm_provider_api_key.split(API_KEYS_SEPARATOR)[1] == inv_apikey
        else:
            assert invariant_auth == ("Bearer " + inv_apikey)
    else:
        assert invariant_auth is None
    if not(not push_to_explorer and invariant_authorization_appended_to_llm_provider_api_key):
        assert llm_provider_api_key == llm_apikey