# mountable_http_hook.py
from typing import Any
from urllib.parse import urlparse

from airflow.providers.http.hooks.http import HttpHook
from requests import Session
from requests.adapters import BaseAdapter
from requests_toolbelt.adapters.socket_options import (  # type: ignore[import-untyped]
    TCPKeepAliveAdapter,
)

class MountableHttpHook(HttpHook):
    """Version of AirFlow's HttpHook that allows mounting custom `requests` adapters."""

    _adapter: BaseAdapter | None

    def __init__(
        self,
        *args,
        adapter: BaseAdapter | None = None,
        tcp_keep_alive: bool = True,
        tcp_keep_alive_idle: int = 120,
        tcp_keep_alive_count: int = 20,
        tcp_keep_alive_interval: int = 30,
        **kwargs,
    ) -> None:
        super().__init__(*args, **kwargs)

        # ensure HttpHook won't override our mounts.
        # set manually instead of by constructor in case overrides
        # pass all params on to a class that doesn't know this parameter,
        # such as is the case with Oauth2HttpHook for requests.Session.
        self.tcp_keep_alive = False

        if adapter is None and tcp_keep_alive:
            # default to the HttpHook's adapter
            adapter = TCPKeepAliveAdapter(
                idle=tcp_keep_alive_idle,
                count=tcp_keep_alive_count,
                interval=tcp_keep_alive_interval,
            )
        self._adapter = adapter

    def get_conn(self, headers: dict[Any, Any] | None = None) -> Session:
        """Add our adapter to the `requests.Session`."""
        session = super().get_conn(headers=headers)
        if self._adapter:
            scheme = urlparse(self.base_url).scheme if self.base_url else "https"
            session.adapters = {scheme: self._adapter}  # type: ignore
        return session