LongArticleTaskServer/app/ab_test/get_cover.py

from app.infra.shared.response import Response
from app.infra.mapper import LongVideoDatabaseMapper, AigcDatabaseMapper


class GetCoverService(Response):
    def __init__(self, pool, params):
        self.pool = pool
        self.params = params
        self.aigc_prefix = "http://res.cybertogether.net/"
        self.pq_prefix = "https://rescdn.yishihui.com/"
        self.suffix = "?x-oss-process=image/resize,m_fill,w_450,h_360/format,jpg"

    async def montage_cover(self, oss_key, pool_name):
        match pool_name:
            case "aigc":
                if oss_key.startswith("http"):
                    return oss_key + self.suffix

                return self.aigc_prefix + oss_key + self.suffix
            case "long_video":
                if oss_key.startswith("http"):
                    return oss_key + self.suffix
                return self.pq_prefix + oss_key + self.suffix
            case _:
                return None

    async def fetch_cover_info(self, pool_name, channel_content_id: str):
        match pool_name:
            case "aigc":
                fetch_response = await AigcDatabaseMapper.fetch_aigc_cover(self.pool, channel_content_id)
                if fetch_response:
                    image_oss = fetch_response[0]["oss_object_key"]
                    if image_oss:
                        cover = await self.montage_cover(image_oss, "aigc")
                    else:
                        image_url = fetch_response[0]["image_url"]
                        if not image_url:
                            return await self.fetch_cover_info(
                                "long_video", channel_content_id
                            )
                        else:
                            cover = await self.montage_cover(image_url, "aigc")
                else:
                    return await self.fetch_cover_info("long_video", channel_content_id)
            case "long_video":
                fetch_response = await LongVideoDatabaseMapper.fetch_long_video_cover(
                    self.pool, channel_content_id
                )
                if fetch_response:
                    image_oss = fetch_response[1]["image_path"]
                    cover = await self.montage_cover(image_oss, "long_video")
                else:
                    return self.error_response(
                        error_code="402",
                        error_message="can't find cover in pq",
                    )
            case _:
                return self.error_response(
                    error_code="403",
                    error_message="pool_name is not supported",
                )

        return self.success_response(data={"cover": cover})

    async def get_cover(self, content_id: str, video_index: int, seed_video_id: str):
        if video_index == 2:
            return await self.fetch_cover_info("long_video", seed_video_id)

        channel_info = await AigcDatabaseMapper.fetch_channel_info(self.pool, content_id)
        if not channel_info:
            return self.error_response(
                error_code="402",
                error_message="can't find content info in aigc system",
            )

        channel_content_id = channel_info[0]["channel_content_id"]
        channel_type = channel_info[0]["channel"]
        match channel_type:
            case 5:
                return await self.fetch_cover_info("aigc", channel_content_id)
            case 6:
                return await self.fetch_cover_info("aigc", channel_content_id)
            case 10:
                return await self.fetch_cover_info("long_video", channel_content_id)
            case _:
                return self.error_response(
                    error_code="403",
                    error_message="channel_type is not supported",
                )

    async def deal(self):
        content_id = self.params.get("content_id")
        video_index = self.params.get("video_index")
        seed_video_id = self.params.get("seed_video_id")
        if not content_id:
            return self.error_response(
                error_code="401",
                error_message="content_id is required",
            )
        elif not video_index:
            return self.error_response(
                error_code="401",
                error_message="video_index is required",
            )
        elif not seed_video_id:
            return self.error_response(
                error_code="401",
                error_message="seed_video_id is required",
            )
        else:
            return await self.get_cover(content_id, int(video_index), seed_video_id)