src/model/user_settings.py

from typing import TYPE_CHECKING, Dict
import os
import json
from sqlalchemy import ForeignKey, String
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column, relationship
from cryptography.fernet import Fernet
from ..database.db import Base, TimestampMixin


ENCRYPTION_KEY = os.getenv("ENCRYPTION_KEY")
if not ENCRYPTION_KEY:
    raise ValueError("ENCRYPTION_KEY environment variable is not set")

fernet = Fernet(ENCRYPTION_KEY.encode())


class UserSettings(Base, TimestampMixin):
    __tablename__ = "user_settings"

    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
    user_id: Mapped[int] = mapped_column(
        ForeignKey("user.id"), nullable=False, index=True, unique=True
    )
    location: Mapped[str | None] = mapped_column(String(100), nullable=True)
    custom_summary_instructions: Mapped[str | None] = mapped_column(
        String, nullable=True
    )
    usability_analysis_instructions: Mapped[str | None] = mapped_column(
        String, nullable=True
    )
    summary_model: Mapped[str | None] = mapped_column(String(100), nullable=True)
    usability_model: Mapped[str | None] = mapped_column(String(100), nullable=True)
    api_keys_encrypted: Mapped[list[dict] | None] = mapped_column(JSONB, nullable=True)

    if TYPE_CHECKING:
        from .user import User  # pragma: no cover

    user: Mapped["User"] = relationship("User", back_populates="user_settings")

    @property
    def api_keys(self) -> list[dict] | None:
        """
        Decrypt and return API keys as a list of dictionaries.
        Each dict contains: {id, slug, name, api_key}

        SECURITY WARNING: This property decrypts all API keys and should NEVER
        be used when returning data to the frontend. Use api_keys_encrypted instead.
        Only use this property internally when you need to access decrypted keys
        for making actual API requests to external services.
        """
        if self.api_keys_encrypted:
            try:
                # If it's already a list (stored in JSONB), decrypt each api_key
                if isinstance(self.api_keys_encrypted, list):
                    decrypted_list = []
                    for item in self.api_keys_encrypted:
                        decrypted_item = item.copy()
                        if "api_key" in decrypted_item and decrypted_item["api_key"]:
                            decrypted_item["api_key"] = fernet.decrypt(
                                decrypted_item["api_key"].encode()
                            ).decode()
                        decrypted_list.append(decrypted_item)
                    return decrypted_list
                # Legacy support: if it's a string, try to decrypt and parse
                elif isinstance(self.api_keys_encrypted, str):
                    decrypted = fernet.decrypt(
                        self.api_keys_encrypted.encode()
                    ).decode()
                    return json.loads(decrypted)
            except Exception as e:
                # Handle decryption errors, perhaps log and return None
                return None
        return None

    @api_keys.setter
    def api_keys(self, value: list[dict] | None):
        """
        Encrypt and set API keys from a list of dictionaries.
        Each dict should contain: {id, slug, name, api_key}
        """
        if value:
            try:
                # Encrypt only the api_key field in each dictionary
                encrypted_list = []
                for item in value:
                    encrypted_item = item.copy()
                    if "api_key" in encrypted_item and encrypted_item["api_key"]:
                        encrypted_item["api_key"] = fernet.encrypt(
                            encrypted_item["api_key"].encode()
                        ).decode()
                    encrypted_list.append(encrypted_item)
                self.api_keys_encrypted = encrypted_list
            except Exception as e:
                # Handle encryption errors
                self.api_keys_encrypted = None
        else:
            self.api_keys_encrypted = None

    @property
    def available_api_keys(self) -> list[str] | None:
        """Return the list of available API key slugs."""
        keys = self.api_keys
        if keys:
            return [item.get("slug") for item in keys if "slug" in item]
        return None

    def get_api_key(self, service_slug: str) -> str | None:
        """
        Get API key for a specific service by its slug.

        SECURITY NOTE: This method decrypts the API key and should ONLY be used
        when actually making API requests to external services. Never use this
        method when returning data to the frontend - use api_keys_encrypted instead.

        Args:
            service_slug: The slug of the service from ServiceCatalog (e.g., 'tavily', 'groq')

        Returns:
            The decrypted API key string if found, None otherwise
        """
        keys = self.api_keys
        if keys:
            for item in keys:
                if item.get("slug") == service_slug:
                    return item.get("api_key")
        return None

    def set_api_key(
        self, service_id: int, service_slug: str, service_name: str, api_key: str
    ) -> None:
        """
        Set or update API key for a specific service.

        Args:
            service_id: The ID of the service from ServiceCatalog
            service_slug: The slug of the service from ServiceCatalog
            service_name: The name of the service
            api_key: The API key value to store (will be encrypted)
        """
        keys = self.api_keys or []

        # Check if service already exists, update it
        found = False
        for item in keys:
            if item.get("slug") == service_slug:
                item["api_key"] = api_key
                item["id"] = service_id
                item["name"] = service_name
                found = True
                break

        # If not found, add new entry
        if not found:
            keys.append(
                {
                    "id": service_id,
                    "slug": service_slug,
                    "name": service_name,
                    "api_key": api_key,
                }
            )

        self.api_keys = keys  # This triggers the setter which encrypts

    def remove_api_key(self, service_slug: str) -> bool:
        """
        Remove API key for a specific service.

        Args:
            service_slug: The slug of the service from ServiceCatalog

        Returns:
            True if the key was removed, False if it didn't exist
        """
        keys = self.api_keys
        if keys:
            original_length = len(keys)
            keys = [item for item in keys if item.get("slug") != service_slug]
            if len(keys) < original_length:
                self.api_keys = keys  # This triggers the setter which encrypts
                return True
        return False

    def __repr__(self) -> str:
        return f"UserSettings(id={self.id}, user_id={self.user_id})"