"""
Custom client entrypoints that extend the generated BaseClient/AsyncBaseClient.

Adds support for:
- `access_token` as an alternative to `api_key` with the following rules:
  - If `access_token` is provided, it takes precedence and sets `Authorization: bearer <token>`
  - When `access_token` is used, `api_key` is forced to "token" to satisfy the generator,
    but the Authorization header is overridden for all HTTP and WebSocket requests.
- `session_id` as a header sent with every request and websocket connection:
  - If `session_id` is provided, it will be used; otherwise, a UUID is auto-generated
  - The session_id is sent as the `x-deepgram-session-id` header
- `transport_factory` to replace the default `websockets` transport with a custom one:
  - A callable: ``factory(url, headers) -> transport`` returning an object with
    ``send()``, ``recv()``, iteration, and ``close()`` support.
- `reconnect` flag (default `True`) declaring whether the SDK is expected to
  manage WebSocket reconnects for this client. When a custom
  ``transport_factory`` is set, ``reconnect`` is auto-disabled because the
  custom transport owns its own retry/reconnect lifecycle; double-stacking
  retries on top would cause storm-on-storm under burst load. The Python
  SDK has no wrapper reconnect layer today, so this flag is declarative
  only -- it documents intent and is reserved for any future SDK-side
  reconnect logic.
"""

import os
import types
import uuid
from typing import Any, Callable, Dict, Optional

from ._secure_logging import install_websocket_log_redaction
from .base_client import AsyncBaseClient, BaseClient
from .transport import install_transport

from deepgram.core.client_wrapper import BaseClientWrapper


def _apply_bearer_authorization_override(client_wrapper: BaseClientWrapper, bearer_token: str) -> None:
    """Override header providers to always use a Bearer authorization token.

    This updates both:
    - client_wrapper.get_headers() used by WebSocket clients
    - client_wrapper.httpx_client.base_headers used by HTTP clients
    """
    original_get_headers = client_wrapper.get_headers

    def _get_headers_with_bearer(_self: Any) -> Dict[str, str]:
        headers = original_get_headers()
        headers["Authorization"] = f"bearer {bearer_token}"
        return headers

    # Override on wrapper for WebSockets
    client_wrapper.get_headers = types.MethodType(_get_headers_with_bearer, client_wrapper)  # type: ignore[method-assign]

    # Override on HTTP client for REST requests
    if hasattr(client_wrapper, "httpx_client") and hasattr(client_wrapper.httpx_client, "base_headers"):
        client_wrapper.httpx_client.base_headers = client_wrapper.get_headers


class DeepgramClient(BaseClient):
    """
    Custom Deepgram client that extends the generated BaseClient.

    Supports:
    - `session_id`: Optional session identifier. If not provided, a UUID is auto-generated.
                     Sent as `x-deepgram-session-id` header in all requests and websocket connections.
    - `access_token`: Alternative to `api_key`. If provided, uses Bearer token authentication.
    - `transport_factory`: Custom sync WebSocket transport factory. A callable
                           ``factory(url, headers) -> transport`` whose return value must support
                           ``send()``, ``recv()``, iteration, and ``close()``.
    - `reconnect`: Declarative flag (default ``True``) signalling whether the SDK is
                    expected to manage WebSocket reconnects for this client. Auto-disabled
                    when ``transport_factory`` is set, since the custom transport owns its
                    retry/reconnect lifecycle. The Python SDK has no wrapper reconnect layer
                    today, so this flag is declarative only -- it documents intent and is
                    reserved for any future SDK-side reconnect logic.
    - `redact_credentials_in_logs`: Mask the `Authorization` header (API key / access token)
                    in the `websockets` library's DEBUG handshake logs. Defaults to `True`; set
                    to `False` to opt out and manage credential redaction yourself.
    - `telemetry_opt_out`: Telemetry opt-out flag (maintained for backwards compatibility, no-op).
    - `telemetry_handler`: Telemetry handler (maintained for backwards compatibility, no-op).
    """

    def __init__(self, *args, **kwargs) -> None:
        access_token: Optional[str] = kwargs.pop("access_token", None)
        session_id: Optional[str] = kwargs.pop("session_id", None)
        transport_factory: Optional[Callable] = kwargs.pop("transport_factory", None)
        reconnect: bool = bool(kwargs.pop("reconnect", True))
        redact_credentials_in_logs: bool = bool(kwargs.pop("redact_credentials_in_logs", True))
        telemetry_opt_out: bool = bool(kwargs.pop("telemetry_opt_out", True))
        telemetry_handler: Optional[Any] = kwargs.pop("telemetry_handler", None)

        # Mask the Authorization header (API key / access token) in the
        # third-party `websockets` library's DEBUG handshake logs so the
        # credential is never written to application logs. On by default;
        # pass `redact_credentials_in_logs=False` to manage redaction yourself.
        if redact_credentials_in_logs:
            install_websocket_log_redaction()

        # Use provided session_id or generate one
        final_session_id = session_id if session_id is not None else str(uuid.uuid4())

        # Ensure headers object exists for pass-through custom headers
        headers: Optional[Dict[str, str]] = kwargs.get("headers")
        if headers is None:
            headers = {}
            kwargs["headers"] = headers

        # Ensure every request has a session identifier header
        headers["x-deepgram-session-id"] = final_session_id

        # Handle access_token: if provided, it takes precedence over api_key
        # The base client requires api_key, so we set a placeholder if needed
        # The Authorization header will be overridden to use Bearer token
        if access_token is not None:
            # Set a placeholder api_key if none provided (base client requires it)
            if kwargs.get("api_key") is None:
                kwargs["api_key"] = "token"
        elif "api_key" not in kwargs:
            # The generated base client takes os.getenv("DEEPGRAM_API_KEY") as a default
            # argument, so it is read once at import. Re-read it here so a key set after
            # import (load_dotenv below the imports) is still picked up.
            kwargs["api_key"] = os.getenv("DEEPGRAM_API_KEY")

        super().__init__(*args, **kwargs)
        self.session_id = final_session_id

        # Override Authorization header to use Bearer token if access_token was provided
        if access_token is not None:
            _apply_bearer_authorization_override(self._client_wrapper, access_token)

        # Install custom WebSocket transport if provided. Auto-disable
        # `reconnect`: a custom transport owns its retry lifecycle, so flip
        # the flag off even if the caller left it at the default.
        if transport_factory is not None:
            install_transport(sync_factory=transport_factory)
            reconnect = False
        self.reconnect = reconnect

        # Store telemetry handler for backwards compatibility (no-op, telemetry not implemented)
        self._telemetry_handler = None


class AsyncDeepgramClient(AsyncBaseClient):
    """
    Custom async Deepgram client that extends the generated AsyncBaseClient.

    Supports:
    - `session_id`: Optional session identifier. If not provided, a UUID is auto-generated.
                     Sent as `x-deepgram-session-id` header in all requests and websocket connections.
    - `access_token`: Alternative to `api_key`. If provided, uses Bearer token authentication.
    - `transport_factory`: Custom async WebSocket transport factory. A callable
                           ``factory(url, headers) -> transport`` whose return value must support
                           ``send()``, ``recv()``, async iteration, and ``close()``.
    - `reconnect`: Declarative flag (default ``True``) signalling whether the SDK is
                    expected to manage WebSocket reconnects for this client. Auto-disabled
                    when ``transport_factory`` is set, since the custom transport owns its
                    retry/reconnect lifecycle. The Python SDK has no wrapper reconnect layer
                    today, so this flag is declarative only -- it documents intent and is
                    reserved for any future SDK-side reconnect logic.
    - `redact_credentials_in_logs`: Mask the `Authorization` header (API key / access token)
                    in the `websockets` library's DEBUG handshake logs. Defaults to `True`; set
                    to `False` to opt out and manage credential redaction yourself.
    - `telemetry_opt_out`: Telemetry opt-out flag (maintained for backwards compatibility, no-op).
    - `telemetry_handler`: Telemetry handler (maintained for backwards compatibility, no-op).
    """

    def __init__(self, *args, **kwargs) -> None:
        access_token: Optional[str] = kwargs.pop("access_token", None)
        session_id: Optional[str] = kwargs.pop("session_id", None)
        transport_factory: Optional[Callable] = kwargs.pop("transport_factory", None)
        reconnect: bool = bool(kwargs.pop("reconnect", True))
        redact_credentials_in_logs: bool = bool(kwargs.pop("redact_credentials_in_logs", True))
        telemetry_opt_out: bool = bool(kwargs.pop("telemetry_opt_out", True))
        telemetry_handler: Optional[Any] = kwargs.pop("telemetry_handler", None)

        # Mask the Authorization header (API key / access token) in the
        # third-party `websockets` library's DEBUG handshake logs so the
        # credential is never written to application logs. On by default;
        # pass `redact_credentials_in_logs=False` to manage redaction yourself.
        if redact_credentials_in_logs:
            install_websocket_log_redaction()

        # Use provided session_id or generate one
        final_session_id = session_id if session_id is not None else str(uuid.uuid4())

        # Ensure headers object exists for pass-through custom headers
        headers: Optional[Dict[str, str]] = kwargs.get("headers")
        if headers is None:
            headers = {}
            kwargs["headers"] = headers

        # Ensure every request has a session identifier header
        headers["x-deepgram-session-id"] = final_session_id

        # Handle access_token: if provided, it takes precedence over api_key
        # The base client requires api_key, so we set a placeholder if needed
        # The Authorization header will be overridden to use Bearer token
        if access_token is not None:
            # Set a placeholder api_key if none provided (base client requires it)
            if kwargs.get("api_key") is None:
                kwargs["api_key"] = "token"
        elif "api_key" not in kwargs:
            # The generated base client takes os.getenv("DEEPGRAM_API_KEY") as a default
            # argument, so it is read once at import. Re-read it here so a key set after
            # import (load_dotenv below the imports) is still picked up.
            kwargs["api_key"] = os.getenv("DEEPGRAM_API_KEY")

        super().__init__(*args, **kwargs)
        self.session_id = final_session_id

        # Override Authorization header to use Bearer token if access_token was provided
        if access_token is not None:
            _apply_bearer_authorization_override(self._client_wrapper, access_token)

        # Install custom WebSocket transport if provided. Auto-disable
        # `reconnect`: a custom transport owns its retry lifecycle, so flip
        # the flag off even if the caller left it at the default.
        if transport_factory is not None:
            install_transport(async_factory=transport_factory)
            reconnect = False
        self.reconnect = reconnect

        # Store telemetry handler for backwards compatibility (no-op, telemetry not implemented)
        self._telemetry_handler = None