banban/talkingq-url/banban/service/im.py

from dataclasses import dataclass
import hashlib
import json
from collections.abc import Mapping
from pathlib import Path
from typing import Any

from fastapi import HTTPException
from services.database_service_base import DatabaseServiceBase
from banban.service.message_audio_storage import MessageAudioStorageService, MessageAudioStorageError

try:
    from banban.dao.im import ImDAO, DeviceIdentity, ConversationMessageCreateResult
    from banban.schemas.im import (
        ChildConversationMessageItem,
        DeviceMessageCreateRequest,
        ParentChildMessageCreateRequest,
    )
except ModuleNotFoundError:
    from banban.dao.im import ImDAO, DeviceIdentity, ConversationMessageCreateResult
    from banban.schemas.im import ChildConversationMessageItem, DeviceMessageCreateRequest, ParentChildMessageCreateRequest


PARENT_PARTICIPANT_TYPE = 1
CHILD_PARTICIPANT_TYPE = 2

CHILD_PEER_CONVERSATION_TYPE = 1
PARENT_CHILD_CONVERSATION_TYPE = 2

CONVERSATION_TYPE_NAMES = {
    CHILD_PEER_CONVERSATION_TYPE: "child_peer",
    PARENT_CHILD_CONVERSATION_TYPE: "parent_child",
}

PARTICIPANT_TYPE_NAMES = {
    PARENT_PARTICIPANT_TYPE: "parent",
    CHILD_PARTICIPANT_TYPE: "child",
}

_AUDIO_CONTENT_TYPE_TO_EXT = {
    "audio/mpeg": "mp3",
    "audio/mp3": "mp3",
    "audio/aac": "aac",
    "audio/wav": "wav",
    "audio/x-wav": "wav",
    "audio/x-m4a": "m4a",
    "audio/mp4": "m4a",
    "audio/webm": "webm",
    "application/octet-stream": "mp3",
}
_AUDIO_EXTENSION_ALIASES = {
    ".mp3": "mp3",
    ".aac": "aac",
    ".m4a": "m4a",
    ".wav": "wav",
    ".webm": "webm",
}


def conversation_type_name(conversation_type: int) -> str:
    return CONVERSATION_TYPE_NAMES.get(conversation_type, f"unknown_{conversation_type}")


def participant_type_name(participant_type: int) -> str:
    return PARTICIPANT_TYPE_NAMES.get(participant_type, f"unknown_{participant_type}")


def build_device_audio_client_msg_id(*, device_id: str, target_device_id: str, audio_url: str) -> str:
    raw = f"{device_id}|{target_device_id}|{audio_url}"
    digest = hashlib.sha256(raw.encode("utf-8")).hexdigest()
    return f"device-audio-{digest[:32]}"


def normalize_content_json(value: Any) -> dict[str, Any] | None:
    if value is None:
        return None
    if isinstance(value, dict):
        return value
    if isinstance(value, str):
        try:
            parsed = json.loads(value)
        except json.JSONDecodeError:
            return None
        if isinstance(parsed, dict):
            return parsed
    return None


def normalize_audio_extension(*, filename: str | None, content_type: str | None) -> tuple[str, str]:
    normalized_content_type = (content_type or "").strip().lower() or "audio/mpeg"
    if normalized_content_type in _AUDIO_CONTENT_TYPE_TO_EXT:
        return _AUDIO_CONTENT_TYPE_TO_EXT[normalized_content_type], normalized_content_type

    suffix = Path(filename or "").suffix.lower()
    if suffix in _AUDIO_EXTENSION_ALIASES:
        extension = _AUDIO_EXTENSION_ALIASES[suffix]
        fallback_content_type = "audio/m4a" if extension == "m4a" else f"audio/{extension}"
        return extension, fallback_content_type

    raise HTTPException(status_code=415, detail="unsupported audio file type")


def row_to_message_item(row: Mapping[str, Any]) -> ChildConversationMessageItem:
    return ChildConversationMessageItem(
        id=int(row["id"]),
        conversation_id=int(row["conversation_id"]),
        seq=int(row["seq"]),
        sender_type=participant_type_name(int(row["sender_type"])),
        sender_id=str(row["sender_id"]),
        receiver_type=participant_type_name(int(row["receiver_type"])),
        receiver_id=str(row["receiver_id"]),
        content_type=int(row["content_type"]),
        content_text=row["content_text"],
        content_json=normalize_content_json(row["content_json"]),
        media_file_key=row["media_file_key"],
        media_duration_ms=row["media_duration_ms"],
        media_mime_type=row["media_mime_type"],
        media_size_bytes=row["media_size_bytes"],
        media_transcript_text=row["media_transcript_text"],
        client_msg_id=row["client_msg_id"],
        sender_name_snapshot=row["sender_name_snapshot"],
        sender_avatar_snapshot=row["sender_avatar_snapshot"],
        receiver_name_snapshot=row["receiver_name_snapshot"],
        receiver_avatar_snapshot=row["receiver_avatar_snapshot"],
        ext_json=normalize_content_json(row["ext_json"]),
        created_at=row["created_at"],
    )


async def present_message_item(
    row: Mapping[str, Any],
    *,
    audio_storage: MessageAudioStorageService,
) -> ChildConversationMessageItem:
    item = row_to_message_item(row)
    if item.content_type == 2 and item.media_file_key:
        try:
            item.media_file_key = await audio_storage.get_audio_url(item.media_file_key)
        except MessageAudioStorageError:
            pass
    return item


class ImService(DatabaseServiceBase):
    def __init__(self):
        super().__init__(service_name="im_service")
        self.audio_storage = MessageAudioStorageService()

    async def assert_parent_child_access(self, *, user_id: int, child_id: int) -> Mapping[str, Any]:
        db_session = await self.get_session()
        try:
            dao = ImDAO(db_session)
            return await dao.assert_parent_child_access(user_id=user_id, child_id=child_id)
        finally:
            await db_session.close()

    async def authenticate_device_identity(self, *, device_id: str, serial_number: str) -> DeviceIdentity:
        db_session = await self.get_session()
        try:
            dao = ImDAO(db_session)
            return await dao.authenticate_device_identity(device_id=device_id, serial_number=serial_number)
        finally:
            await db_session.close()

    async def _build_message_create_result(
        self,
        *,
        dao: ImDAO,
        idempotent: bool,
        conversation_id: int,
        conversation_type: int,
        client_msg_id: str,
    ) -> ConversationMessageCreateResult:
        message_row = await dao._get_message_by_conversation_client_id(
            conversation_id=conversation_id,
            client_msg_id=client_msg_id,
        )
        if not message_row:
            raise RuntimeError("message was not found after insert")

        presented_message = await present_message_item(
            message_row,
            audio_storage=self.audio_storage,
        )

        return ConversationMessageCreateResult(
            idempotent=idempotent,
            conversation_id=conversation_id,
            conversation_type=conversation_type,
            message=presented_message,
        )

    async def create_parent_child_message(
        self,
        *,
        parent_user_id: int,
        child_id: int,
        payload: ParentChildMessageCreateRequest,
    ) -> ConversationMessageCreateResult:
        db_session = await self.get_session()
        try:
            dao = ImDAO(db_session)
            child_row = await dao.assert_parent_child_access(user_id=parent_user_id, child_id=child_id)
            parent_row = await dao._get_parent_row(parent_user_id)
            if not parent_row:
                raise HTTPException(status_code=404, detail="parent not found")

            conversation_id, idempotent = await dao.create_message(
                conversation_type=PARENT_CHILD_CONVERSATION_TYPE,
                participant_a_type=CHILD_PARTICIPANT_TYPE,
                participant_a_id=str(child_id),
                participant_b_type=PARENT_PARTICIPANT_TYPE,
                participant_b_id=str(parent_user_id),
                pair_key=f"{child_id}:{parent_user_id}",
                sender_type=PARENT_PARTICIPANT_TYPE,
                sender_id=str(parent_user_id),
                receiver_type=CHILD_PARTICIPANT_TYPE,
                receiver_id=str(child_id),
                sender_name_snapshot=parent_row["nickname"],
                sender_avatar_snapshot=parent_row["avatar_url"],
                receiver_name_snapshot=child_row["child_name"],
                receiver_avatar_snapshot=None,
                payload=payload,
            )
            return await self._build_message_create_result(
                dao=dao,
                idempotent=idempotent,
                conversation_id=conversation_id,
                conversation_type=PARENT_CHILD_CONVERSATION_TYPE,
                client_msg_id=payload.client_msg_id,
            )
        except Exception:
            await db_session.rollback()
            raise
        finally:
            await db_session.close()

    async def create_parent_child_voice_message(
        self,
        *,
        parent_user_id: int,
        child_id: int,
        filename: str | None,
        content_type: str | None,
        content: bytes,
        media_duration_ms: int | None,
        media_transcript_text: str | None,
        client_msg_id: str,
        ext_json: dict[str, Any] | None = None,
    ) -> ConversationMessageCreateResult:
        if not content:
            raise HTTPException(status_code=400, detail="audio file is empty")

        extension, normalized_content_type = normalize_audio_extension(
            filename=filename,
            content_type=content_type,
        )
        stored_audio = await self.audio_storage.upload_audio(
            device_id=f"parent-{parent_user_id}",
            content=content,
            content_type=normalized_content_type,
            extension=extension,
        )
        payload = ParentChildMessageCreateRequest(
            content_type=2,
            media_file_key=stored_audio.file_key,
            media_duration_ms=media_duration_ms,
            media_mime_type=normalized_content_type,
            media_size_bytes=len(content),
            media_transcript_text=(media_transcript_text or "").strip() or None,
            client_msg_id=client_msg_id,
            ext_json=ext_json,
        )

        try:
            result = await self.create_parent_child_message(
                parent_user_id=parent_user_id,
                child_id=child_id,
                payload=payload,
            )
        except Exception:
            try:
                await self.audio_storage.delete_audio(stored_audio.file_key)
            except MessageAudioStorageError:
                pass
            raise

        return result

    async def _create_device_message_with_payload(
        self,
        *,
        dao: ImDAO,
        device_identity: DeviceIdentity,
        payload: DeviceMessageCreateRequest,
    ) -> ConversationMessageCreateResult:
        if payload.conversation_type == CHILD_PEER_CONVERSATION_TYPE:
            if payload.peer_child_id == device_identity.child_id:
                raise HTTPException(status_code=400, detail="peer_child_id must be different from current child")
            sender_child_row = await dao.assert_child_exists(child_id=device_identity.child_id)
            receiver_child_row = await dao.assert_child_exists(child_id=payload.peer_child_id)
            participant_a_id, participant_b_id, pair_key = dao._build_child_peer_pair(
                device_identity.child_id,
                payload.peer_child_id,
            )

            conversation_id, idempotent = await dao.create_message(
                conversation_type=CHILD_PEER_CONVERSATION_TYPE,
                participant_a_type=CHILD_PARTICIPANT_TYPE,
                participant_a_id=participant_a_id,
                participant_b_type=CHILD_PARTICIPANT_TYPE,
                participant_b_id=participant_b_id,
                pair_key=pair_key,
                sender_type=CHILD_PARTICIPANT_TYPE,
                sender_id=str(device_identity.child_id),
                receiver_type=CHILD_PARTICIPANT_TYPE,
                receiver_id=str(payload.peer_child_id),
                sender_name_snapshot=sender_child_row["child_name"],
                sender_avatar_snapshot=None,
                receiver_name_snapshot=receiver_child_row["child_name"],
                receiver_avatar_snapshot=None,
                payload=payload,
            )
        else:
            sender_child_row = await dao.assert_child_exists(child_id=device_identity.child_id)
            parent_row = await dao._get_parent_row(payload.parent_user_id)
            if not parent_row:
                raise HTTPException(status_code=404, detail="parent not found")
            await dao.assert_parent_child_access(user_id=payload.parent_user_id, child_id=device_identity.child_id)

            conversation_id, idempotent = await dao.create_message(
                conversation_type=PARENT_CHILD_CONVERSATION_TYPE,
                participant_a_type=CHILD_PARTICIPANT_TYPE,
                participant_a_id=str(device_identity.child_id),
                participant_b_type=PARENT_PARTICIPANT_TYPE,
                participant_b_id=str(payload.parent_user_id),
                pair_key=f"{device_identity.child_id}:{payload.parent_user_id}",
                sender_type=CHILD_PARTICIPANT_TYPE,
                sender_id=str(device_identity.child_id),
                receiver_type=PARENT_PARTICIPANT_TYPE,
                receiver_id=str(payload.parent_user_id),
                sender_name_snapshot=sender_child_row["child_name"],
                sender_avatar_snapshot=None,
                receiver_name_snapshot=parent_row["nickname"],
                receiver_avatar_snapshot=parent_row["avatar_url"],
                payload=payload,
            )

        return await self._build_message_create_result(
            dao=dao,
            idempotent=idempotent,
            conversation_id=conversation_id,
            conversation_type=payload.conversation_type,
            client_msg_id=payload.client_msg_id,
        )

    async def create_device_message(
        self,
        *,
        device_id: str,
        serial_number: str,
        payload: DeviceMessageCreateRequest | None = None,
        target_device_id: str | None = None,
        audio_url: str | None = None,
    ) -> tuple[DeviceIdentity, ConversationMessageCreateResult]:
        if payload is None and (not target_device_id or not audio_url):
            raise ValueError("payload or target_device_id/audio_url is required")
        if payload is not None and (target_device_id is not None or audio_url is not None):
            raise ValueError("payload and target_device_id/audio_url cannot be used together")

        db_session = await self.get_session()
        try:
            dao = ImDAO(db_session)
            device_identity = await dao.authenticate_device_identity(
                device_id=device_id,
                serial_number=serial_number,
            )

            resolved_payload = payload
            if resolved_payload is None:
                target_device_identity = await dao.get_device_by_id(device_id=target_device_id)
                resolved_payload = DeviceMessageCreateRequest(
                    conversation_type=CHILD_PEER_CONVERSATION_TYPE,
                    peer_child_id=target_device_identity.child_id,
                    content_type=2,
                    media_file_key=audio_url,
                    media_mime_type="audio/mpeg",
                    client_msg_id=build_device_audio_client_msg_id(
                        device_id=device_id,
                        target_device_id=target_device_id,
                        audio_url=audio_url,
                    ),
                    ext_json={
                        "source": "device_audio_message",
                        "source_device_id": device_id,
                        "target_device_id": target_device_id,
                    },
                )

            result = await self._create_device_message_with_payload(
                dao=dao,
                device_identity=device_identity,
                payload=resolved_payload,
            )
            return device_identity, result
        except Exception:
            await db_session.rollback()
            raise
        finally:
            await db_session.close()

    async def assert_child_exists(self, *, child_id: int) -> Mapping[str, Any]:
        db_session = await self.get_session()
        try:
            dao = ImDAO(db_session)
            return await dao.assert_child_exists(child_id=child_id)
        finally:
            await db_session.close()


im_service = ImService()