# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_vault

import asyncio
import time
from typing import Optional

import httpx
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential

from coreason_vault.config import CoreasonVaultConfig
from coreason_vault.exceptions import VaultConnectionError
from coreason_vault.utils.logger import logger

# Define retryable exceptions
RETRYABLE_EXCEPTIONS = (
    httpx.RequestError,
    httpx.HTTPStatusError,
)


class VaultAuthenticationAsync:
    """Handles authentication with HashiCorp Vault (Async).

    Supports AppRole and Kubernetes authentication methods.
    """

    def __init__(self, config: CoreasonVaultConfig, client: Optional[httpx.AsyncClient] = None):
        """Initializes the VaultAuthenticationAsync instance.

        Args:
            config (CoreasonVaultConfig): The configuration object containing Vault settings.
            client (Optional[httpx.AsyncClient]): An optional external client for connection pooling.
        """
        self.config = config
        self._internal_client = client is None
        self._client = client or httpx.AsyncClient(
            base_url=str(config.VAULT_ADDR), verify=config.VAULT_VERIFY_SSL, timeout=30.0
        )
        if config.VAULT_NAMESPACE:
            self._client.headers["X-Vault-Namespace"] = config.VAULT_NAMESPACE

        self._token: Optional[str] = None
        self._last_token_check: float = 0.0
        self._lock = asyncio.Lock()

    async def close(self) -> None:
        """Closes the underlying HTTP client if it was created internally."""
        if self._internal_client:
            await self._client.aclose()

    async def get_client(self) -> httpx.AsyncClient:
        """Returns an authenticated Vault client (Async).

        Checks token validity and renews/re-authenticates if necessary.

        Returns:
            httpx.AsyncClient: An authenticated Vault client instance.
        """
        # First check (optimistic)
        if self._is_token_fresh():
            self._client.headers["X-Vault-Token"] = self._token or ""  # pragma: no cover
            return self._client

        async with self._lock:
            # Double-check inside lock
            if self._is_token_fresh():
                self._client.headers["X-Vault-Token"] = self._token or ""  # pragma: no cover
                return self._client

            if self._token is None:
                await self._authenticate()
                return self._client

            # Must validate
            try:
                await self._validate_token()
                self._last_token_check = time.time()
            except Exception:
                logger.info("Vault token expired or invalid, re-authenticating...")  # pragma: no cover
                await self._authenticate()

            return self._client

    def _is_token_fresh(self) -> bool:
        return self._token is not None and not self._should_validate_token()

    def _should_validate_token(self) -> bool:
        return (time.time() - self._last_token_check) > self.config.VAULT_TOKEN_TTL

    async def _validate_token(self) -> None:
        """Validates the current token."""
        try:
            self._client.headers["X-Vault-Token"] = self._token or ""  # pragma: no cover
            response = await self._client.get("v1/auth/token/lookup-self")
            response.raise_for_status()
            data = response.json()
            ttl = data.get("data", {}).get("ttl", 0)
            if ttl < 10:
                raise VaultConnectionError("Token expiring soon")
        except httpx.HTTPStatusError as e:
            if e.response.status_code in (403, 404):
                raise VaultConnectionError("Token invalid") from e
            raise

    @retry(  # type: ignore[misc]
        retry=retry_if_exception_type(RETRYABLE_EXCEPTIONS),
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=10),
        reraise=True,
    )
    async def _authenticate(self) -> None:
        """Authenticates to Vault using the configured method (Async)."""
        try:
            token: Optional[str] = None

            if self.config.VAULT_ROLE_ID and self.config.VAULT_SECRET_ID:
                logger.info("Authenticating to Vault via AppRole (Async)")
                payload = {
                    "role_id": self.config.VAULT_ROLE_ID,
                    "secret_id": self.config.VAULT_SECRET_ID,
                }
                response = await self._client.post("v1/auth/approle/login", json=payload)
                response.raise_for_status()
                token = response.json()["auth"]["client_token"]

            elif self.config.KUBERNETES_SERVICE_ACCOUNT_TOKEN:
                logger.info("Authenticating to Vault via Kubernetes (Async)")
                role = self.config.VAULT_K8S_ROLE
                if not role:
                    raise ValueError("Missing Kubernetes role (VAULT_K8S_ROLE)")

                payload = {
                    "role": role,
                    "jwt": self.config.KUBERNETES_SERVICE_ACCOUNT_TOKEN,
                }
                response = await self._client.post("v1/auth/kubernetes/login", json=payload)
                response.raise_for_status()
                token = response.json()["auth"]["client_token"]
            else:
                raise ValueError("Missing authentication credentials (AppRole or Kubernetes)")

            if not token:
                raise VaultConnectionError("Vault authentication failed: No token returned")

            self._token = token
            self._client.headers["X-Vault-Token"] = token
            self._last_token_check = time.time()
            logger.info("Successfully authenticated to Vault (Async)")

        except (httpx.HTTPError, VaultConnectionError) as e:
            logger.error(f"Unexpected error during Vault authentication: {e}")
            raise VaultConnectionError(f"Vault authentication failed: {e}") from e
        except Exception:
            # Let ValueError and other critical exceptions propagate
            raise
