from cryptography.fernet import Fernet, InvalidToken
from sqlalchemy.orm import Session
from db_models import ApiKey
import os

# Initialize the encryption key from environment or generate a new one
ENCRYPTION_KEY = os.getenv("ENCRYPTION_KEY", Fernet.generate_key().decode())
cipher_suite = Fernet(ENCRYPTION_KEY.encode())


class ApiKeyService:
    """
    Service for managing API keys in the database
    """

    @staticmethod
    def encrypt_api_key(api_key: str) -> str:
        """
        Encrypt an API key
        """
        encrypted_key = cipher_suite.encrypt(api_key.encode())
        return encrypted_key.decode()

    @staticmethod
    def decrypt_api_key(encrypted_api_key: str) -> str:
        """
        Decrypt an API key
        """
        decrypted_key = cipher_suite.decrypt(encrypted_api_key.encode())
        return decrypted_key.decode()

    @staticmethod
    def get_api_key(db: Session, provider: str) -> str:
        """
        Retrieve and decrypt an API key for a provider
        """
        api_key_record = db.query(ApiKey).filter(ApiKey.provider == provider).first()
        if not api_key_record or not api_key_record.api_key_encrypted:
            return None
        try:
            return ApiKeyService.decrypt_api_key(api_key_record.api_key_encrypted)
        except InvalidToken:
            return None

    @staticmethod
    def set_api_key(db: Session, provider: str, api_key: str) -> bool:
        """
        Encrypt and store an API key for a provider
        """
        encrypted_key = ApiKeyService.encrypt_api_key(api_key)

        # Check if record exists
        api_key_record = db.query(ApiKey).filter(ApiKey.provider == provider).first()

        if api_key_record:
            # Update existing record
            api_key_record.api_key_encrypted = encrypted_key
        else:
            # Create new record
            api_key_record = ApiKey(
                provider=provider,
                api_key_encrypted=encrypted_key
            )
            db.add(api_key_record)

        try:
            db.commit()
            return True
        except Exception:
            db.rollback()
            return False