import os

import httpx
import opik.hooks


class AcceldataAuth(httpx.Auth):
    def __init__(self, access_key: str, secret_key: str) -> None:
        self.access_key = access_key
        self.secret_key = secret_key

    def auth_flow(self, request):  # type: ignore
        request.headers.pop("authorization", None)
        request.headers["accessKey"] = self.access_key
        request.headers["secretKey"] = self.secret_key
        yield request


def setup_acceldata_auth_hook(access_key: str, secret_key: str) -> None:
    def acceldata_auth_client_hook(client: httpx.Client) -> None:
        client.auth = AcceldataAuth(access_key, secret_key)

    opik.hooks.add_httpx_client_hook(
        opik.hooks.HttpxClientHook(
            client_modifier=acceldata_auth_client_hook,
            client_init_arguments=None,
        )
    )


def setup_acceldata_auth(url: str, access_key: str, secret_key: str) -> None:
    """Boot-time setup for users of @opik.track and Opik integrations.

    Call once before any tracked function runs. Subsequent Opik clients
    (including the cached one used by @track) will pick up the auth.
    """
    os.environ.setdefault("OPIK_URL_OVERRIDE", url)
    setup_acceldata_auth_hook(access_key, secret_key)
