"""Legacy API key rotation and validation with PBKDF2 hashing.

Legacy keys use different format and hashing than managed keys:
- Format: lls_{token_urlsafe(32)} (vs lls_mk_... for managed)
- Hashing: PBKDF2-SHA256 with salt (vs plain SHA-256 for managed)
- Key ID: key_{token_hex(8)} (vs mk_... for managed)
- Lookup: SHA-256 lookup_hash GSI for O(1) validation
"""

import hashlib
import logging
import os
import secrets
from datetime import datetime
from typing import TYPE_CHECKING, Any

from ai_lls_lib.key_management import compute_key_hash

try:
    import boto3
    from botocore.exceptions import ClientError

    HAS_BOTO3 = True
except ImportError:
    boto3 = None  # type: ignore[assignment]
    HAS_BOTO3 = False

if TYPE_CHECKING:
    from mypy_boto3_dynamodb.service_resource import Table

logger = logging.getLogger(__name__)

LEGACY_KEY_PREFIX = "lls_"
LEGACY_KEY_ID_PREFIX = "key_"
PBKDF2_ITERATIONS = 100_000


def compute_legacy_key_hash(key: str, salt: bytes | None = None) -> tuple[str, str]:
    """Compute PBKDF2-SHA256 hash of a legacy API key.

    Args:
        key: Plaintext API key.
        salt: Optional salt bytes (generated if not provided).

    Returns:
        Tuple of (hash_hex, salt_hex).
    """
    if salt is None:
        salt = secrets.token_bytes(32)

    hash_bytes = hashlib.pbkdf2_hmac("sha256", key.encode(), salt, PBKDF2_ITERATIONS)
    return hash_bytes.hex(), salt.hex()


class LegacyApiKeyService:
    """Manages legacy single-key-per-user API keys with PBKDF2 hashing.

    DynamoDB table schema:
        - Hash key: api_key_id (S)
        - GSI: gsi_user on user_id (HASH)
    """

    table: "Table | None"

    def __init__(self, table_name: str | None = None):
        """Initialize with DynamoDB table."""
        if not HAS_BOTO3 or not boto3:
            raise RuntimeError("boto3 is required for LegacyApiKeyService")

        self.dynamodb = boto3.resource("dynamodb")
        self.table_name = table_name if table_name else os.environ["API_KEYS_TABLE"]

        try:
            self.table = self.dynamodb.Table(self.table_name)
        except Exception as e:
            logger.error(f"Failed to connect to DynamoDB table {self.table_name}: {e}")
            self.table = None

    def rotate_key(self, user_id: str, email: str | None = None) -> dict[str, Any]:
        """Rotate (or create) a user's legacy API key.

        Deactivates all existing active keys for the user, then creates
        a new key with PBKDF2 hashing.

        Args:
            user_id: User identifier.
            email: Optional user email stored for reference.

        Returns:
            Dict with api_key_id, api_key (plaintext), created_at, last4, message.
        """
        if not self.table:
            raise RuntimeError(f"DynamoDB table {self.table_name} not accessible")

        # Find and deactivate existing keys
        self._deactivate_user_keys(user_id)

        # Generate new key
        api_key = f"{LEGACY_KEY_PREFIX}{secrets.token_urlsafe(32)}"
        api_key_id = f"{LEGACY_KEY_ID_PREFIX}{secrets.token_hex(8)}"
        key_hash, salt_hex = compute_legacy_key_hash(api_key)
        lookup_hash = compute_key_hash(api_key)
        now = datetime.utcnow().isoformat()

        item: dict[str, Any] = {
            "api_key_id": api_key_id,
            "user_id": user_id,
            "key_hash": key_hash,
            "salt": salt_hex,
            "lookup_hash": lookup_hash,
            "created_at": now,
            "active": True,
            "last_used": None,
            "usage_count": 0,
        }
        if email:
            item["email"] = email

        self.table.put_item(Item=item)

        logger.info(f"Rotated legacy key for user {user_id}, new key_id: {api_key_id}")
        return {
            "api_key_id": api_key_id,
            "api_key": api_key,
            "created_at": now,
            "last4": api_key[-4:],
            "message": "API key rotated successfully",
        }

    def _deactivate_user_keys(self, user_id: str) -> None:
        """Deactivate all active keys for a user via the gsi_user GSI."""
        if not self.table:
            return

        try:
            response = self.table.query(
                IndexName="gsi_user",
                KeyConditionExpression="user_id = :uid",
                ExpressionAttributeValues={":uid": user_id},
            )
            now = datetime.utcnow().isoformat()

            for item in response.get("Items", []):
                if item.get("active"):
                    self.table.update_item(
                        Key={"api_key_id": item["api_key_id"]},
                        UpdateExpression=("SET active = :false_val, deactivated_at = :now"),
                        ExpressionAttributeValues={
                            ":false_val": False,
                            ":now": now,
                        },
                    )
        except ClientError as e:
            logger.error(f"Error deactivating keys for {user_id}: {e}")
            raise

    def validate_by_lookup_hash(self, api_key: str) -> dict[str, Any] | None:
        """Validate a legacy API key using the lookup_hash GSI (O(1) path).

        Computes the SHA-256 lookup hash, queries the GSI, then verifies
        with PBKDF2 against the stored hash and salt.

        Args:
            api_key: Plaintext API key to validate.

        Returns:
            The DynamoDB item dict if valid, or None.
        """
        if not self.table:
            raise RuntimeError(f"DynamoDB table {self.table_name} not accessible")

        lookup_hash = compute_key_hash(api_key)

        response = self.table.query(
            IndexName="gsi_lookup_hash",
            KeyConditionExpression="lookup_hash = :lh",
            ExpressionAttributeValues={":lh": lookup_hash},
        )

        for item in response.get("Items", []):
            if not item.get("active"):
                continue
            stored_hash = str(item["key_hash"])
            salt_bytes = bytes.fromhex(str(item["salt"]))
            computed_hash, _ = compute_legacy_key_hash(api_key, salt_bytes)
            if computed_hash == stored_hash:
                return dict(item)

        return None

    def backfill_lookup_hash(self, api_key_id: str, lookup_hash: str) -> None:
        """Write lookup_hash to an existing item for GSI indexing.

        Args:
            api_key_id: The key's primary key.
            lookup_hash: SHA-256 hex digest of the plaintext key.
        """
        if not self.table:
            raise RuntimeError(f"DynamoDB table {self.table_name} not accessible")

        self.table.update_item(
            Key={"api_key_id": api_key_id},
            UpdateExpression="SET lookup_hash = :lh",
            ExpressionAttributeValues={":lh": lookup_hash},
        )
        logger.info(f"Backfilled lookup_hash for {api_key_id}")

    def validate_key(self, api_key: str) -> dict[str, Any] | None:
        """Validate a legacy API key using GSI-first with scan fallback.

        Tries the O(1) lookup_hash GSI path first. If no match (key
        predates the GSI), falls back to a full table scan with PBKDF2
        verification and lazily backfills the lookup_hash.

        Args:
            api_key: Plaintext API key to validate.

        Returns:
            The DynamoDB item dict if valid, or None.
        """
        if not self.table:
            raise RuntimeError(f"DynamoDB table {self.table_name} not accessible")

        # O(1) GSI path
        result = self.validate_by_lookup_hash(api_key)
        if result:
            return result

        # O(n) scan fallback for items without lookup_hash
        lookup_hash = compute_key_hash(api_key)
        response = self.table.scan()

        for item in response.get("Items", []):
            if not item.get("active"):
                continue
            salt_hex = item.get("salt")
            if not salt_hex:
                continue
            salt_bytes = bytes.fromhex(str(salt_hex))
            computed_hash, _ = compute_legacy_key_hash(api_key, salt_bytes)
            if computed_hash == str(item["key_hash"]):
                # Lazy backfill
                self.backfill_lookup_hash(str(item["api_key_id"]), lookup_hash)
                return dict(item)

        return None
