"""
TOTP (Time-based One-Time Password) service for admin sudo authentication.

This module provides secure OTP generation, verification, and rate limiting
for the admin sudo authentication system.
"""

from __future__ import annotations

import hmac
import json
import logging
import secrets
import string
from typing import TYPE_CHECKING

import pyotp
from django.conf import settings
from django.core.cache import cache
from django.utils import timezone

if TYPE_CHECKING:
    from labs.models import AdminOTPDevice, User

logger = logging.getLogger(__name__)

# Cache key prefixes for rate limiting
RATE_LIMIT_KEY_PREFIX = "admin_otp_rate_limit:"
RATE_LIMIT_COUNT_KEY_PREFIX = "admin_otp_rate_limit_count:"

# Recovery code configuration
RECOVERY_CODE_LENGTH = 10
RECOVERY_CODE_COUNT = 8


class OTPRateLimitError(Exception):
    """Raised when too many failed OTP verification attempts."""

    def __init__(self, remaining_seconds: int):
        self.remaining_seconds = remaining_seconds
        super().__init__(
            f"Too many failed attempts. Please try again in {remaining_seconds} seconds."
        )


class OTPVerificationError(Exception):
    """Raised when OTP verification fails."""

    pass


class OTPReplayError(OTPVerificationError):
    """Raised when a TOTP code is reused within the same time window."""

    pass


class TOTPService:
    """
    Service class for TOTP operations.

    Handles secret generation, code verification, recovery codes,
    and rate limiting for the admin OTP sudo system.
    """

    def __init__(self) -> None:
        self.issuer_name = getattr(settings, "ADMIN_OTP_ISSUER_NAME", "TiDB Labs Admin")
        self.max_attempts = getattr(settings, "ADMIN_OTP_MAX_ATTEMPTS", 5)
        self.lockout_duration = getattr(
            settings, "ADMIN_OTP_LOCKOUT_DURATION", 15 * 60
        )  # 15 minutes
        # TOTP time step in seconds (RFC 6238 default is 30 seconds)
        self.totp_period = getattr(settings, "ADMIN_OTP_TOTP_PERIOD", 30)

    @staticmethod
    def generate_secret() -> str:
        """
        Generate a new random TOTP secret.

        Returns:
            A base32-encoded secret string suitable for TOTP.
        """
        return pyotp.random_base32()

    @staticmethod
    def generate_recovery_codes() -> list[str]:
        """
        Generate a set of one-time recovery codes.

        Returns:
            A list of uppercase alphanumeric recovery codes.
        """
        alphabet = string.ascii_uppercase + string.digits
        # Remove confusing characters (0, O, 1, I, L)
        alphabet = alphabet.replace("0", "").replace("O", "")
        alphabet = alphabet.replace("1", "").replace("I", "").replace("L", "")

        codes = []
        for _ in range(RECOVERY_CODE_COUNT):
            code = "".join(
                secrets.choice(alphabet) for _ in range(RECOVERY_CODE_LENGTH)
            )
            codes.append(code)
        return codes

    def get_totp(self, secret: str) -> pyotp.TOTP:
        """Create a TOTP instance from a secret."""
        return pyotp.TOTP(secret, interval=self.totp_period)

    def get_provisioning_uri(self, user: User, secret: str) -> str:
        """
        Generate a provisioning URI for authenticator apps.

        This URI can be encoded as a QR code for easy setup.

        Args:
            user: The user to generate the URI for.
            secret: The TOTP secret.

        Returns:
            A otpauth:// URI string.
        """
        totp = self.get_totp(secret)
        return totp.provisioning_uri(name=user.email, issuer_name=self.issuer_name)

    def _get_rate_limit_key(self, user_id: int) -> str:
        """Get the cache key for rate limit lock."""
        return f"{RATE_LIMIT_KEY_PREFIX}{user_id}"

    def _get_rate_limit_count_key(self, user_id: int) -> str:
        """Get the cache key for failed attempt count."""
        return f"{RATE_LIMIT_COUNT_KEY_PREFIX}{user_id}"

    def is_rate_limited(self, user: User) -> tuple[bool, int]:
        """
        Check if the user is currently rate-limited.

        Args:
            user: The user to check.

        Returns:
            A tuple of (is_limited, remaining_seconds).
        """
        # Django's cache API is intentionally minimal and does not expose TTL.
        # To stay backend-agnostic, we store an absolute lock expiry timestamp
        # (epoch seconds) in the cache value and compute remaining time from it.
        lock_key = self._get_rate_limit_key(user.id)
        locked_until_ts = cache.get(lock_key)
        if not locked_until_ts:
            return False, 0
        try:
            locked_until_ts_int = int(locked_until_ts)
        except (TypeError, ValueError):
            cache.delete(lock_key)
            return False, 0

        now_ts = int(timezone.now().timestamp())
        remaining = locked_until_ts_int - now_ts
        if remaining <= 0:
            cache.delete(lock_key)
            return False, 0
        return True, remaining

    def _record_failed_attempt(self, user: User) -> None:
        """
        Record a failed verification attempt and apply rate limiting if needed.
        """
        count_key = self._get_rate_limit_count_key(user.id)
        lock_key = self._get_rate_limit_key(user.id)

        # Avoid relying on backend-specific cache methods like ttl()/expire()/incr().
        # Use simple get/set so it works across Django cache backends.
        current = cache.get(count_key) or 0
        try:
            count = int(current) + 1
        except (TypeError, ValueError):
            count = 1

        if count >= self.max_attempts:
            # Lock the user out
            locked_until_ts = int(timezone.now().timestamp()) + int(
                self.lockout_duration
            )
            cache.set(lock_key, locked_until_ts, timeout=self.lockout_duration)
            # Reset the count
            cache.delete(count_key)
            logger.warning(
                "OTP rate limit triggered for user %s (id=%d)",
                user.email,
                user.id,
            )
            return

        # Persist updated count with a reasonable expiry window.
        cache.set(count_key, count, timeout=self.lockout_duration)

    def _clear_failed_attempts(self, user: User) -> None:
        """Clear failed attempt count after successful verification."""
        count_key = self._get_rate_limit_count_key(user.id)
        cache.delete(count_key)

    def _check_replay_attack(self, device: AdminOTPDevice) -> None:
        """
        Check if a TOTP code verification would be a replay attack.

        A replay attack occurs when the same TOTP code is used multiple times
        within its validity window. To prevent this, we ensure that at least
        one TOTP period has passed since the last successful verification.

        Args:
            device: The OTP device to check.

        Raises:
            OTPReplayError: If the code would be a replay.
        """
        if device.last_totp_verified_at is None:
            return

        time_since_last = (
            timezone.now() - device.last_totp_verified_at
        ).total_seconds()
        if time_since_last < self.totp_period:
            logger.warning(
                "TOTP replay attack detected for user %s (id=%d): "
                "%.1f seconds since last verification",
                device.user.email,
                device.user.id,
                time_since_last,
            )
            raise OTPReplayError(
                "This code has already been used. Please wait for a new code."
            )

    def _update_totp_verified_timestamp(self, device: AdminOTPDevice) -> None:
        """
        Update the last TOTP verification timestamp on the device.

        Args:
            device: The OTP device to update.
        """
        device.last_totp_verified_at = timezone.now()
        device.save(update_fields=["last_totp_verified_at", "updated_at"])

    def verify_code(
        self, device: AdminOTPDevice, code: str, check_rate_limit: bool = True
    ) -> bool:
        """
        Verify a TOTP code against the device's secret.

        This method includes protection against replay attacks by tracking
        the timestamp of the last successful verification.

        Args:
            device: The OTP device to verify against.
            code: The 6-digit code to verify.
            check_rate_limit: Whether to check and apply rate limiting.

        Returns:
            True if the code is valid.

        Raises:
            OTPRateLimitError: If the user is rate-limited.
            OTPReplayError: If the code has already been used.
            OTPVerificationError: If the code is invalid.
        """
        user = device.user

        if check_rate_limit:
            is_limited, remaining = self.is_rate_limited(user)
            if is_limited:
                raise OTPRateLimitError(remaining)

        # Check for replay attack before verifying the code
        self._check_replay_attack(device)

        # Get the decrypted secret
        secret = device.secret_encrypted
        totp = self.get_totp(secret)

        # Verify with a window of 1 (allows for slight clock drift)
        # This accepts codes from 1 period before and after the current time
        is_valid = totp.verify(code, valid_window=1)

        if is_valid:
            self._clear_failed_attempts(user)
            # Update the TOTP verification timestamp for replay protection
            self._update_totp_verified_timestamp(device)
            device.mark_used()
            logger.info(
                "OTP verification successful for user %s (id=%d)",
                user.email,
                user.id,
            )
            return True

        if check_rate_limit:
            self._record_failed_attempt(user)

        logger.warning(
            "OTP verification failed for user %s (id=%d)",
            user.email,
            user.id,
        )
        raise OTPVerificationError("Invalid verification code")

    def verify_recovery_code(self, device: AdminOTPDevice, code: str) -> bool:
        """
        Verify and consume a recovery code.

        Args:
            device: The OTP device.
            code: The recovery code to verify.

        Returns:
            True if the code is valid and was consumed.

        Raises:
            OTPRateLimitError: If the user is rate-limited.
            OTPVerificationError: If the code is invalid.
        """
        user = device.user

        is_limited, remaining = self.is_rate_limited(user)
        if is_limited:
            raise OTPRateLimitError(remaining)

        if not device.recovery_codes_encrypted:
            self._record_failed_attempt(user)
            raise OTPVerificationError("No recovery codes available")

        # Parse recovery codes
        try:
            codes = json.loads(device.recovery_codes_encrypted)
        except json.JSONDecodeError:
            self._record_failed_attempt(user)
            raise OTPVerificationError("Invalid recovery code data")

        # Normalize the input code (uppercase, no spaces/dashes)
        normalized_code = code.upper().replace(" ", "").replace("-", "")

        # Use constant-time comparison to prevent timing attacks.
        # We iterate through all codes to avoid early exit that could
        # leak information about which codes are valid.
        matched_code: str | None = None
        for stored_code in codes:
            # hmac.compare_digest provides constant-time comparison
            if hmac.compare_digest(normalized_code, stored_code):
                matched_code = stored_code
                # Continue checking all codes to maintain constant time

        if matched_code is None:
            self._record_failed_attempt(user)
            logger.warning(
                "Recovery code verification failed for user %s (id=%d)",
                user.email,
                user.id,
            )
            raise OTPVerificationError("Invalid recovery code")

        # Remove the used code
        codes.remove(matched_code)
        device.recovery_codes_encrypted = json.dumps(codes)
        device.last_used_at = timezone.now()
        device.save(
            update_fields=["recovery_codes_encrypted", "last_used_at", "updated_at"]
        )

        self._clear_failed_attempts(user)
        logger.info(
            "Recovery code used for user %s (id=%d). Remaining codes: %d",
            user.email,
            user.id,
            len(codes),
        )
        return True

    def get_remaining_recovery_codes_count(self, device: AdminOTPDevice) -> int:
        """Get the number of remaining recovery codes."""
        if not device.recovery_codes_encrypted:
            return 0
        try:
            codes = json.loads(device.recovery_codes_encrypted)
            return len(codes)
        except json.JSONDecodeError:
            return 0


# Global service instance
totp_service = TOTPService()