"""Authentication flow."""

import contextlib
import socket
import string
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timedelta
from secrets import SystemRandom

import msgspec
from pytz import UTC

from .config import settings
from .server import wait_for_authorization_response
from .token import OAuthToken

DEVICE_FLOW_POLLING_FREQUENCY = 2


_OAUTH2_CLIENT = None


def _get_oauth2_client():
    # lazily loaded, because otherwise, it slows import time
    global _OAUTH2_CLIENT  # noqa: PLW0603 one time init using global state
    if _OAUTH2_CLIENT is None:
        from slices_clientlib_core.util.disable_authlib_warning import disable_authlib_warning

        disable_authlib_warning()
        from authlib.integrations.httpx_client import OAuth2Client

        _OAUTH2_CLIENT = OAuth2Client(
            client_id=settings.client_id,
            client_secret=settings.client_secret,
            scope=settings.scope,
            token_endpoint_auth_method="client_secret_post",  # noqa: S106
        )
    return _OAUTH2_CLIENT


_CODE_VERIFIER_CHARS = None
_CODE_VERIFIER_RANDOM = None


def _generate_code_verifier():
    # lazily loaded, because otherwise, it slows import time
    global _CODE_VERIFIER_CHARS, _CODE_VERIFIER_RANDOM  # noqa: PLW0603 one time init using global state
    if _CODE_VERIFIER_CHARS is None or _CODE_VERIFIER_RANDOM is None:
        _CODE_VERIFIER_CHARS = string.ascii_letters + string.digits + "-._~"
        _CODE_VERIFIER_RANDOM = SystemRandom()
    return "".join(_CODE_VERIFIER_RANDOM.choice(_CODE_VERIFIER_CHARS) for _ in range(128))


class AuthenticationFlow:
    """Authentication Flow."""

    authorization_url: str
    """URL that the user must visit to authenticate."""

    def __init__(self) -> None:
        """Initialize the AuthenticationFlow."""

        callback_port = self._find_open_port()

        self._callback_server_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="AuthFlowServer")
        self._callback_server_future = self._callback_server_executor.submit(
            wait_for_authorization_response, "localhost", callback_port
        )

        code_verifier = _generate_code_verifier()

        self._redirect_uri = f"http://localhost:{callback_port}/callback"
        self.authorization_url, self._state = _get_oauth2_client().create_authorization_url(
            settings.authorization_url,
            redirect_uri=self._redirect_uri,
            code_verifier=code_verifier,
        )

    def wait_for_authorization_token(self) -> OAuthToken:
        """Blocking method which waits for the user to authenticate and returns the token."""
        authorization_response = self._callback_server_future.result()

        token = _get_oauth2_client().fetch_token(
            settings.token_endpoint_url,
            state=self._state,
            authorization_response=authorization_response,
            redirect_uri=self._redirect_uri,
        )
        if token:
            return msgspec.convert(token, OAuthToken)

        raise ValueError("Did not receive a token")

    @staticmethod
    def _find_open_port(start_port=None, stop_port=None):
        start = start_port or 8080
        stop = stop_port or start + 100
        for port in range(start, stop):
            with contextlib.closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:
                try:
                    sock.bind(("127.0.0.1", port))
                    sock.listen(1)
                except OSError:
                    is_open = False
                else:
                    is_open = True
            if is_open:
                return port
        raise ConnectionError("Could not find an open port")


class DeviceFlow:
    """Device Authentication Flow."""

    def __init__(self):
        """Perform OAuth Login flow and return authorization code."""

        code_verifier = _generate_code_verifier()

        from slices_clientlib_core.util.disable_authlib_warning import disable_authlib_warning

        disable_authlib_warning()
        from authlib.oauth2.rfc8628 import DEVICE_CODE_GRANT_TYPE

        _, self._state = _get_oauth2_client().create_authorization_url(
            settings.authorization_url,
            code_verifier=code_verifier,
            grant_type=DEVICE_CODE_GRANT_TYPE,
        )

        device_code_response = _get_oauth2_client().request(
            "POST",
            settings.device_authorization_url,
            withhold_token=True,
            data={
                "client_id": settings.client_id,
                "client_secret": settings.client_secret,
                "scope": settings.scope,
            },
        )

        self._device_code_info = device_code_response.json()

        self.authorization_url = self._device_code_info["verification_uri_complete"]
        self._expires_at = datetime.now(UTC) + timedelta(seconds=self._device_code_info["expires_in"])

    def wait_for_authorization_token(self) -> OAuthToken:
        """Blocking method which waits for the user to authenticate and returns the token."""

        from slices_clientlib_core.util.disable_authlib_warning import disable_authlib_warning

        disable_authlib_warning()
        from authlib.integrations.httpx_client import OAuthError
        from authlib.oauth2.rfc8628 import DEVICE_CODE_GRANT_TYPE

        while self._expires_at > datetime.now(UTC):
            time.sleep(DEVICE_FLOW_POLLING_FREQUENCY)
            try:
                token = _get_oauth2_client().fetch_token(
                    settings.token_endpoint_url,
                    state=self._state,
                    grant_type=DEVICE_CODE_GRANT_TYPE,
                    device_code=self._device_code_info["device_code"],
                )
                if token:
                    return msgspec.convert(token, OAuthToken)
            except OAuthError as err:
                if err.error != "authorization_pending":
                    raise ValueError("Exception while waiting for device authorization") from err

        raise TimeoutError("Device login flow timed out.")


class RefreshFlow:
    """Refresh Flow."""

    def perform_refresh(self, refresh_token: str):
        """Perform OAuth Login flow and return authorization code."""

        return msgspec.convert(
            _get_oauth2_client().refresh_token(settings.token_endpoint_url, refresh_token), OAuthToken
        )
