# 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 base64
from typing import Optional, Union

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

from coreason_vault.auth_async import VaultAuthenticationAsync
from coreason_vault.exceptions import EncryptionError
from coreason_vault.utils.logger import logger

RETRYABLE_EXCEPTIONS = (
    httpx.RequestError,
    httpx.HTTPStatusError,
)


class TransitCipherAsync:
    """Async version of TransitCipher."""

    def __init__(self, auth: VaultAuthenticationAsync):
        self._auth = auth

    async def encrypt(self, plaintext: Union[str, bytes], key_name: str, context: Optional[str] = None) -> str:
        try:
            return await self._encrypt_impl(plaintext, key_name, context)  # type: ignore[no-any-return]
        except Exception as e:
            if isinstance(e, TypeError):
                raise
            logger.error(f"Encryption failed after retries: {e}")
            raise EncryptionError(f"Encryption failed: {e}") from e

    @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 _encrypt_impl(self, plaintext: Union[str, bytes], key_name: str, context: Optional[str] = None) -> str:
        client = await self._auth.get_client()

        encoded_plaintext = self._encode_base64(plaintext)
        encoded_context = self._encode_base64(context) if context else None

        # Payload for transit encrypt
        payload = {"plaintext": encoded_plaintext, "context": encoded_context}

        try:
            response = await client.post(f"v1/transit/encrypt/{key_name}", json=payload)
            response.raise_for_status()
            return response.json()["data"]["ciphertext"]  # type: ignore[no-any-return]
        except httpx.HTTPStatusError as e:
            logger.error(f"Encryption failed: {e}")
            raise
        except Exception:
            raise

    async def decrypt(self, ciphertext: str, key_name: str, context: Optional[str] = None) -> Union[str, bytes]:
        try:
            return await self._decrypt_impl(ciphertext, key_name, context)  # type: ignore[no-any-return]
        except Exception as e:
            logger.error(f"Decryption failed after retries: {e}")
            raise EncryptionError(f"Decryption failed: {e}") from e

    @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 _decrypt_impl(self, ciphertext: str, key_name: str, context: Optional[str] = None) -> Union[str, bytes]:
        client = await self._auth.get_client()
        encoded_context = self._encode_base64(context) if context else None

        payload = {"ciphertext": ciphertext, "context": encoded_context}

        try:
            response = await client.post(f"v1/transit/decrypt/{key_name}", json=payload)
            response.raise_for_status()

            encoded_plaintext = response.json()["data"]["plaintext"]
            plaintext_bytes = base64.b64decode(encoded_plaintext, validate=True)

            try:
                return plaintext_bytes.decode("utf-8")
            except UnicodeDecodeError:
                return plaintext_bytes
        except httpx.HTTPStatusError as e:
            logger.error(f"Decryption failed: {e}")
            raise
        except Exception:
            raise

    def _encode_base64(self, data: Union[str, bytes]) -> str:
        if isinstance(data, str):
            data_bytes = data.encode("utf-8")
        elif isinstance(data, bytes):
            data_bytes = data
        else:
            raise TypeError(f"Expected str or bytes, got {type(data)}")
        return base64.b64encode(data_bytes).decode("utf-8")
