# common/src/common/dao/stream_key.py
"""
Data Access Object for stream keys in PostgreSQL.

Stream keys are encrypted at rest using AWS KMS via the Crypto utility.
"""

import base64
import datetime
import hashlib
from typing import Optional

from common.database import DBConfig, Db
from common.logging import get_logger, span
from common.models.stream_key import StreamKey, StreamKeyUpdate
from common.utils.crypto import Crypto, CryptoConfig

logger = get_logger(__name__)

SELECT_STREAM_KEY_QUERY = """
                          SELECT id,
                                 user_id,
                                 name,
                                 description,
                                 encrypted_key,
                                 is_active,
                                 consumer_active,
                                 created_at,
                                 updated_at
                          FROM stream_keys
                          """

SELECT_STREAM_KEY_BY_ID_QUERY = (
    SELECT_STREAM_KEY_QUERY + "WHERE id = %s AND is_active = TRUE"
)

SELECT_STREAM_KEY_BY_ID_INCLUDE_INACTIVE_QUERY = (
    SELECT_STREAM_KEY_QUERY + "WHERE id = %s"
)

SELECT_STREAM_KEYS_BY_USER_QUERY = (
    SELECT_STREAM_KEY_QUERY
    + "WHERE user_id = %s AND is_active = TRUE ORDER BY created_at DESC"
)

SELECT_STREAM_KEY_BY_KEY_VALUE_QUERY = """
                                        SELECT id,
                                               user_id,
                                               name,
                                               description,
                                               encrypted_key,
                                               is_active,
                                               consumer_active,
                                               created_at,
                                               updated_at
                                        FROM stream_keys
                                        WHERE encrypted_key = %s
                                        """


class StreamKeyNotFoundError(Exception):
    """Exception raised when a stream key is not found."""

    def __init__(self, message="Stream key not found"):
        self.message = message
        super().__init__(self.message)


class StreamKeyDAO:
    """Data Access Object for stream keys in PostgreSQL."""

    def __init__(self, db: Optional[Db] = None, crypto: Optional[Crypto] = None):
        """Initialize with an existing DB connection or create a new one."""
        self.db = db if db else Db(DBConfig())
        self.crypto = crypto if crypto else Crypto(config=CryptoConfig())
        logger.debug("StreamKeyDAO initialized", extra={"component": "StreamKeyDAO"})

    async def create_stream_key(
        self, user_id: int, key: str, name: str | None, description: str | None
    ) -> StreamKey:
        """Create a new stream key.

        The key is encrypted before storage.
        """
        with span(
            logger,
            "create_stream_key",
            {"user_id": user_id},
            log_entry=False,
            log_exit=False,
        ):
            try:
                # Encrypt the stream key and URL-safe base64 encode for string storage
                # URL-safe encoding uses - and _ instead of + and / to avoid issues in RTMP URLs
                encrypted_bytes = self.crypto.encrypt(key.encode("utf-8"))
                encrypted_key = base64.urlsafe_b64encode(encrypted_bytes).decode(
                    "utf-8"
                )

                # SHA-256 hash for O(1) lookup during RTMP authentication
                key_hash = hashlib.sha256(key.encode("utf-8")).hexdigest()

                columns = [
                    "user_id",
                    "name",
                    "description",
                    "encrypted_key",
                    "key_hash",
                    "is_active",
                    "consumer_active",
                ]
                values = [
                    user_id,
                    name,
                    description,
                    encrypted_key,
                    key_hash,
                    True,
                    False,
                ]

                result = await self.db.insert(
                    table="stream_keys", columns=columns, values=tuple(values)
                )

                if result:
                    # Fetch the created stream key to get the full record
                    created_key = await self.get_stream_key_by_id(
                        result, include_inactive=True
                    )
                    logger.info(
                        f"Created stream key with ID {result}",
                        extra={"stream_key_id": result, "user_id": user_id},
                    )
                    return created_key

                logger.error(
                    f"Failed to create stream key for user {user_id}",
                    extra={"user_id": user_id},
                )
                raise Exception("Failed to create stream key")
            except Exception as e:
                logger.error(
                    f"Error creating stream key: {e}",
                    extra={"user_id": user_id, "error": str(e)},
                )
                raise

    async def get_stream_key_by_id(
        self, stream_key_id: int, include_inactive: bool = False
    ) -> StreamKey:
        """Get a stream key by its ID.

        Args:
            stream_key_id: The database ID of the stream key
            include_inactive: If True, include inactive stream keys

        Returns:
            StreamKey: The stream key (key value is not decrypted in this method)

        Raises:
            StreamKeyNotFoundError: If the stream key is not found
        """
        with span(
            logger,
            "get_stream_key_by_id",
            {"stream_key_id": stream_key_id},
            log_entry=False,
            log_exit=False,
        ):
            try:
                query = (
                    SELECT_STREAM_KEY_BY_ID_INCLUDE_INACTIVE_QUERY
                    if include_inactive
                    else SELECT_STREAM_KEY_BY_ID_QUERY
                )
                row = await self.db.fetch_one(query, (stream_key_id,))
                if not row:
                    logger.warning(
                        f"Stream key with ID {stream_key_id} not found",
                        extra={"stream_key_id": stream_key_id},
                    )
                    raise StreamKeyNotFoundError(
                        f"Stream key with ID {stream_key_id} not found."
                    )

                stream_key = self._row_to_stream_key(row)
                logger.debug(
                    f"Found stream key with ID {stream_key_id}",
                    extra={
                        "stream_key_id": stream_key_id,
                        "user_id": stream_key.user_id,
                    },
                )
                return stream_key
            except StreamKeyNotFoundError:
                raise
            except Exception as e:
                logger.error(
                    f"Error retrieving stream key with ID {stream_key_id}: {e}",
                    extra={"stream_key_id": stream_key_id, "error": str(e)},
                )
                raise

    async def get_stream_keys_by_user(self, user_id: int) -> list[StreamKey]:
        """Get all active stream keys for a user.

        Args:
            user_id: The user ID

        Returns:
            List of StreamKey objects
        """
        with span(
            logger,
            "get_stream_keys_by_user",
            {"user_id": user_id},
            log_entry=False,
            log_exit=False,
        ):
            try:
                rows = await self.db.fetch_all(
                    SELECT_STREAM_KEYS_BY_USER_QUERY, (user_id,)
                )
                if not rows:
                    logger.debug(
                        f"No stream keys found for user {user_id}",
                        extra={"user_id": user_id},
                    )
                    return []

                stream_keys = [self._row_to_stream_key(row) for row in rows]
                logger.debug(
                    f"Found {len(stream_keys)} stream keys for user {user_id}",
                    extra={"user_id": user_id, "count": len(stream_keys)},
                )
                return stream_keys
            except Exception as e:
                logger.error(
                    f"Error retrieving stream keys for user {user_id}: {e}",
                    extra={"user_id": user_id, "error": str(e)},
                )
                raise

    async def get_stream_key_by_encrypted_key(
        self, encrypted_key: str
    ) -> Optional[StreamKey]:
        """Get a stream key by its encrypted key value.

        This is useful for validating stream keys during ingestion.

        Args:
            encrypted_key: The encrypted key str

        Returns:
            StreamKey if found, None otherwise
        """
        with span(
            logger,
            "get_stream_key_by_encrypted_key",
            log_entry=False,
            log_exit=False,
        ):
            try:
                row = await self.db.fetch_one(
                    SELECT_STREAM_KEY_BY_KEY_VALUE_QUERY, (encrypted_key,)
                )
                if not row:
                    logger.debug("Stream key not found by encrypted key")
                    return None

                stream_key = self._row_to_stream_key(row)
                logger.debug(
                    "Found stream key by encrypted key",
                    extra={
                        "stream_key_id": stream_key.id,
                        "user_id": stream_key.user_id,
                    },
                )
                return stream_key
            except Exception as e:
                logger.error(
                    f"Error retrieving stream key by encrypted key: {e}",
                    extra={"error": str(e)},
                )
                raise

    async def get_stream_key_by_plaintext_key(
        self, plaintext_key: str
    ) -> Optional[StreamKey]:
        """Get a stream key by its plaintext key value.

        Hashes the plaintext key and looks up by key_hash column.
        This is the preferred method for RTMP authentication.

        Args:
            plaintext_key: The plaintext stream key (what the user enters in OBS)

        Returns:
            StreamKey if found, None otherwise
        """
        with span(
            logger,
            "get_stream_key_by_plaintext_key",
            log_entry=False,
            log_exit=False,
        ):
            try:
                # Hash the plaintext key
                key_hash = hashlib.sha256(plaintext_key.encode("utf-8")).hexdigest()

                row = await self.db.fetch_one(
                    SELECT_STREAM_KEY_QUERY
                    + "WHERE key_hash = %s AND is_active = TRUE",
                    (key_hash,),
                )
                if not row:
                    logger.debug("Stream key not found by key hash")
                    return None

                stream_key = self._row_to_stream_key(row)
                logger.debug(
                    "Found stream key by key hash",
                    extra={
                        "stream_key_id": stream_key.id,
                        "user_id": stream_key.user_id,
                    },
                )
                return stream_key
            except Exception as e:
                logger.error(
                    f"Error retrieving stream key by plaintext key: {e}",
                    extra={"error": str(e)},
                )
                raise

    async def update_stream_key(
        self, stream_key_id: int, update: StreamKeyUpdate
    ) -> StreamKey:
        """Update a stream key's metadata.

        Args:
            stream_key_id: The database ID of the stream key
            update: The update data

        Returns:
            Updated StreamKey

        Raises:
            StreamKeyNotFoundError: If the stream key is not found
        """
        with span(
            logger,
            "update_stream_key",
            {"stream_key_id": stream_key_id},
            log_entry=False,
            log_exit=False,
        ):
            try:
                # Verify stream key exists
                await self.get_stream_key_by_id(stream_key_id, include_inactive=True)

                update_dict = update.dict(exclude_none=True)
                if not update_dict:
                    logger.debug(
                        f"No updates provided for stream key {stream_key_id}, skipping",
                        extra={"stream_key_id": stream_key_id},
                    )
                    return await self.get_stream_key_by_id(
                        stream_key_id, include_inactive=True
                    )

                set_clauses = []
                values = []

                for key, value in update_dict.items():
                    set_clauses.append(key)
                    values.append(value)

                # Always update the updated_at timestamp
                set_clauses.append("updated_at")
                values.append(datetime.datetime.utcnow())

                await self.db.update(
                    table="stream_keys",
                    set_columns=set_clauses,
                    set_values=values,
                    where_clause="id = %s",
                    where_params=(stream_key_id,),
                )
                logger.info(
                    f"Updated stream key {stream_key_id} with fields: {list(update_dict.keys())}",
                    extra={
                        "stream_key_id": stream_key_id,
                        "fields": list(update_dict.keys()),
                    },
                )
                return await self.get_stream_key_by_id(
                    stream_key_id, include_inactive=True
                )
            except StreamKeyNotFoundError:
                raise
            except Exception as e:
                logger.error(
                    f"Error updating stream key {stream_key_id}: {e}",
                    extra={"stream_key_id": stream_key_id, "error": str(e)},
                )
                raise

    async def update_consumer_status(
        self, stream_key_id: int, active: bool
    ) -> StreamKey:
        """Update the consumer active status for a stream key.

        Args:
            stream_key_id: The database ID of the stream key
            active: Whether the consumer is active

        Returns:
            Updated StreamKey

        Raises:
            StreamKeyNotFoundError: If the stream key is not found
        """
        with span(
            logger,
            "update_consumer_status",
            {"stream_key_id": stream_key_id, "active": active},
            log_entry=False,
            log_exit=False,
        ):
            try:
                await self.db.update(
                    table="stream_keys",
                    set_columns=["consumer_active", "updated_at"],
                    set_values=[active, datetime.datetime.utcnow()],
                    where_clause="id = %s",
                    where_params=(stream_key_id,),
                )
                logger.info(
                    f"Updated consumer status for stream key {stream_key_id} to {active}",
                    extra={"stream_key_id": stream_key_id, "consumer_active": active},
                )
                return await self.get_stream_key_by_id(
                    stream_key_id, include_inactive=True
                )
            except Exception as e:
                logger.error(
                    f"Error updating consumer status for stream key {stream_key_id}: {e}",
                    extra={"stream_key_id": stream_key_id, "error": str(e)},
                )
                raise

    async def delete_stream_key(self, stream_key_id: int) -> None:
        """Delete (soft delete) a stream key by setting is_active to False.

        Args:
            stream_key_id: The database ID of the stream key

        Raises:
            StreamKeyNotFoundError: If the stream key is not found
        """
        with span(
            logger,
            "delete_stream_key",
            {"stream_key_id": stream_key_id},
            log_entry=False,
            log_exit=False,
        ):
            try:
                # Verify stream key exists
                await self.get_stream_key_by_id(stream_key_id, include_inactive=True)

                # Soft delete by setting is_active to False
                await self.db.update(
                    table="stream_keys",
                    set_columns=["is_active", "updated_at"],
                    set_values=[False, datetime.datetime.utcnow()],
                    where_clause="id = %s",
                    where_params=(stream_key_id,),
                )
                logger.info(
                    f"Deleted (deactivated) stream key {stream_key_id}",
                    extra={"stream_key_id": stream_key_id},
                )
            except StreamKeyNotFoundError:
                raise
            except Exception as e:
                logger.error(
                    f"Error deleting stream key {stream_key_id}: {e}",
                    extra={"stream_key_id": stream_key_id, "error": str(e)},
                )
                raise

    def decrypt_stream_key(self, stream_key: StreamKey) -> str:
        """Decrypt the encrypted_key field to get the plaintext stream key.

        This is needed when we have a StreamKey object but need the original
        plaintext key for constructing RTMP URLs.

        Args:
            stream_key: The StreamKey object with encrypted_key field

        Returns:
            The plaintext stream key string
        """
        with span(
            logger,
            "decrypt_stream_key",
            {"stream_key_id": stream_key.id},
            log_entry=False,
            log_exit=False,
        ):
            try:
                # Decode from URL-safe base64 back to bytes
                encrypted_bytes = base64.urlsafe_b64decode(stream_key.encrypted_key)
                # Decrypt using AWS KMS
                plaintext_bytes = self.crypto.decrypt(encrypted_bytes)
                return plaintext_bytes.decode("utf-8")
            except Exception as e:
                logger.error(
                    f"Error decrypting stream key {stream_key.id}: {e}",
                    extra={"stream_key_id": stream_key.id, "error": str(e)},
                )
                raise

    def _row_to_stream_key(self, row) -> StreamKey:
        """Convert a database row to a StreamKey model."""
        if not row:
            raise ValueError("Row is empty")
        if len(row) < 9:
            raise ValueError(
                "Row does not contain enough data to create a StreamKey object."
            )

        return StreamKey(
            id=row[0],  # Database ID
            user_id=row[1],
            name=row[2],
            description=row[3],
            encrypted_key=row[4],
            is_active=row[5],
            consumer_active=row[6],
            created_at=row[7],
            updated_at=row[8],
        )
