nizika_broadcast/test.py

from __future__ import annotations

import math
import os
import random
import sys
import wave
from datetime import datetime, timedelta
from enum import Enum, auto
from typing import Callable, TypedDict

import cv2
import emoji
import ephem
import pygame
import pygame.gfxdraw
import pytchat
import requests
from cv2 import VideoCapture
from ephem import Moon, Observer, Sun
from pygame import Rect, Surface
from pygame.font import Font
from pygame.mixer import Sound
from pygame.time import Clock
from pytchat.core.pytchat import PytchatCore
from pytchat.processors.default.processor import Chat

from aques import Aques
from common_module import CommonModule
from nizika_ai.config import DB
from nizika_ai.consts import (AnswerType, Character, GPTModel, Platform,
                              QueryType)
from nizika_ai.models import Answer, AnsweredFlag, Query, User

pygame.init ()

FPS = 30

SYSTEM_FONT = pygame.font.SysFont ('notosanscjkjp', 24, bold = True)
USER_FONT = pygame.font.SysFont ('notosanscjkjp', 32, italic = True)
DEERJIKA_FONT = pygame.font.SysFont ('07nikumarufont', 50)


def main (
) -> None:
    game = Game ()
    Bg (game)
    balloon = Balloon (game)
    deerjika = Deerjika (game, DeerjikaPattern.RELAXED,
                         x = CWindow.WIDTH * 3 / 4,
                         y = CWindow.HEIGHT - 120,
                         balloon = balloon)
    CurrentTime (game, SYSTEM_FONT)
    try:
        broadcast = Broadcast (os.environ['BROADCAST_CODE'])
    except Exception:
        pass
    while True:
        for event in pygame.event.get ():
            if event.type == pygame.QUIT:
                pygame.quit ()
                sys.exit ()
        if not balloon.enabled:
            try:
                DB.begin_transaction ()
                answer_flags = (AnsweredFlag.where ('platform', Platform.YOUTUBE.value)
                                            .where ('answered', False)
                                            .get ())
                if answer_flags:
                    answer_flag = random.choice (answer_flags)
                    answer = Answer.find (answer_flag.answer_id)
                    if answer.answer_type == AnswerType.YOUTUBE_REPLY.value:
                        query = Query.find (answer.query_id)
                        deerjika.talk (query.content, answer.content)
                        answer_flag.answered = True
                        answer_flag.save ()
                DB.commit ()
                add_query (broadcast)
            except Exception:
                pass
        game.redraw ()


class Bg:
    """
    背景オブゼクト管理用クラス

    Attributes:
        base (BgBase):      最背面
        grass (BgGrass):    草原部分
        jojoko (Jojoko):    大月ヨヨコ
        kita (KitaSun):     き太く陽
    """

    base:   BgBase
    grass:  BgGrass
    jojoko: Jojoko
    kita:   KitaSun

    def __init__ (
            self,
            game:   Game,
    ):
        self.base = BgBase (game)
        self.jojoko = Jojoko (game)
        self.kita = KitaSun (game)
        self.grass = BgGrass (game)


class DeerjikaPattern (Enum):
    """
    ニジカの状態

    Members:
        NORMAL:     通常
        RELAXED:    足パタパタ
        SLEEPING:   寝ニジカ
        DANCING:    ダンシング・ニジカ
    """

    NORMAL      = auto ()
    RELAXED     = auto ()
    SLEEPING    = auto ()
    DANCING     = auto ()


class Direction (Enum):
    """
    クリーチャの向き

    Members:
        LEFT:   左向き
        RIGHT:  右向き
    """

    LEFT    = auto ()
    RIGHT   = auto ()


class Game:
    """
    ゲーム・クラス

    Attributes:
        clock (Clock):                  Clock オブゼクト
        frame (int):                    フレーム・カウンタ
        last_answered_at (datetime):    最後に回答した時刻
        now (datetime):                 基準日時
        objects (list[GameObject]):     再描画するクラスのリスト
        screen (Surface):               基底スクリーン
        sky (Sky):                      天体情報
    """

    bgm:                Sound
    clock:              Clock
    frame:              int
    last_answered_at:   datetime
    now:                datetime
    objects:            list[GameObject]
    screen:             Surface
    sky:                Sky

    def __init__ (
            self,
    ):
        self.now = datetime.now ()
        self.screen = pygame.display.set_mode ((CWindow.WIDTH, CWindow.HEIGHT))
        self.clock = Clock ()
        self.frame = 0
        self.objects = []
        self.bgm = Sound ('assets/bgm.mp3')
        self.bgm.set_volume (.15)
        self.bgm.play (loops = -1)
        self._create_sky ()

    def redraw (
            self,
    ) -> None:
        self.now = datetime.now ()
        self.sky.observer.date = self.now - timedelta (hours = 9)
        for obj in sorted (self.objects, key = lambda obj: obj.layer):
            if obj.enabled:
                obj.redraw ()
        pygame.display.update ()
        self.clock.tick (FPS)

    def _create_sky (
            self,
    ) -> None:
        self.sky = Sky ()
        self.sky.observer = Observer ()
        self.sky.observer.lat = '35'
        self.sky.observer.lon = '139'


class GameObject:
    """
    各ゲーム・オブゼクトの基底クラス

    Attributes:
        arg (float):    回転角度 (rad)
        ax (float):     X 軸に対する加速度 (px/frame^2)
        ay (float):     y 軸に対する加速度 (px/frame^2)
        enabled (bool): オブゼクトの表示可否
        frame (int):    フレーム・カウンタ
        game (Game):    ゲーム基盤
        height (int):   高さ (px)
        vx (float):     x 軸に対する速度 (px/frame)
        vy (float):     y 軸に対する速度 (px/frame)
        width (int):    幅 (px)
        x (float):      X 座標 (px)
        y (float):      Y 座標 (px)
    """

    arg:        float   = 0
    ax:         float   = 0
    ay:         float   = 0
    enabled:    bool    = True
    frame:      int
    game:       Game
    height:     int
    layer:      float
    vx:         float   = 0
    vy:         float   = 0
    width:      int
    x:          float
    y:          float

    def __init__ (
            self,
            game:       Game,
            layer:      float | None    = None,
            enabled:    bool            = True,
            x:          float           = 0,
            y:          float           = 0,
    ):
        self.game = game
        self.enabled = enabled
        self.frame = 0
        if layer is None:
            if self.game.objects:
                layer = max (obj.layer for obj in self.game.objects) + 10
            else:
                layer = 0
        self.layer = layer
        self.x = x
        self.y = y
        self.game.objects.append (self)

    def redraw (
            self,
    ) -> None:
        self.x += self.vx
        self.y += self.vy
        self.vx += self.ax
        self.vy += self.ay
        self.frame += 1


class BgBase (GameObject):
    """
    背景

    Attributes:
        surface (Surface):  背景 Surface
    """

    bg:         Surface
    bg_evening: Surface
    bg_grass:   Surface
    bg_night:   Surface
    surface:    Surface

    def __init__ (
            self,
            game:   Game,
    ):
        super ().__init__ (game)
        self.bg = pygame.image.load ('assets/bg.jpg')
        self.bg_evening = pygame.image.load ('assets/bg-evening.jpg')
        self.bg_grass = pygame.image.load ('assets/bg-grass.png')
        self.bg_night = pygame.image.load ('assets/bg-night.jpg')
        self.surface = pygame.transform.scale (self.bg, (CWindow.WIDTH, CWindow.HEIGHT))

    def redraw (
            self,
    ) -> None:
        self.game.screen.blit (self.surface, (self.x, self.y))
        super ().redraw ()


class BgGrass (GameObject):
    """
    背景の草原部分

    Attributes:
        surface (Surface):  草原 Surface
    """

    surface:    Surface

    def __init__ (
            self,
            game:   Game,
    ):
        super ().__init__ (game)
        self.game = game
        self.surface = pygame.image.load ('assets/bg-grass.png')
        self.surface = pygame.transform.scale (self.surface, (CWindow.WIDTH, CWindow.HEIGHT))

    def redraw (
            self,
    ) -> None:
        self.game.screen.blit (self.surface, (self.x, self.y))
        super ().redraw ()


class Creature (GameObject):
    sound:  Sound

    def bell (
        self,
    ) -> None:
        self.sound.play ()


class Deerjika (Creature):
    """
    伊地知ニジカ

    Attributes:
        height (int):               高さ (px)
        scale (float):              拡大率
        surfaces (list[Surface]):   ニジカの各フレームを Surface にしたリスト
        width (int):                幅 (px)
    """

    height:     int
    scale:      float               = .8
    surfaces:   list[Surface]
    width:      int
    talking:    bool                = False
    wav:        bytearray | None    = None
    balloon:    Balloon

    def __init__ (
            self,
            game:       Game,
            pattern:    DeerjikaPattern = DeerjikaPattern.NORMAL,
            direction:  Direction       = Direction.LEFT,
            layer:      float | None    = None,
            x:          float           = 0,
            y:          float           = 0,
            balloon:    Balloon | None  = None,
    ):
        if balloon is None:
            raise Exception
        super ().__init__ (game, layer, x = x, y = y)
        self.pattern = pattern
        self.direction = direction
        self.balloon = balloon
        match pattern:
            case DeerjikaPattern.NORMAL:
                ...
            case DeerjikaPattern.RELAXED:
                match direction:
                    case Direction.LEFT:
                        self.width = 1280
                        self.height = 720
                        surface = pygame.image.load ('assets/deerjika_relax_left.png')
                        self.surfaces = []
                        for x in range (0, surface.get_width (), self.width):
                            self.surfaces.append (
                                    surface.subsurface (x, 0, self.width, self.height))
                    case Direction.RIGHT:
                        ...
        self.sound = Sound ('assets/noon.wav')

    def redraw (
            self,
    ) -> None:
        surface = pygame.transform.scale (self.surfaces[self.frame % len (self.surfaces)],
                                          (self.width * self.scale, self.height * self.scale))
        self.game.screen.blit (surface, surface.get_rect (center = (self.x, self.y)))
        super ().redraw ()
        if (not self.balloon.enabled) and self.talking:
            self.talking = False
        if (self.balloon.enabled and self.balloon.frame >= FPS * 1.5
                and not self.talking):
            self.read_out ()

    def talk (
            self,
            query:      str,
            answer:     str,
    ) -> None:
        self.bell ()
        self._create_wav (answer)
        length = 300
        if self.wav is not None:
            with wave.open ('./nizika_talking.wav', 'rb') as f:
                length = int (FPS * (f.getnframes () / f.getframerate () + 4))
        self.balloon.talk (query, answer, length = length)

    def read_out (
            self,
    ) -> None:
        Sound ('./nizika_talking.wav').play ()
        self.talking = True

    def _create_wav (
            self,
            message:    str,
    ) -> None:
        try:
            self.wav = Aques.main (message, False)
        except:
            self.wav = None
        if self.wav is None:
            return
        with open ('./nizika_talking.wav', 'wb') as f:
            f.write (self.wav)


class CurrentTime (GameObject):
    """
    現在日時表示

    Attributes:
        font (Font):    フォント
    """

    font:   Font

    def __init__ (
            self,
            game:   Game,
            font:   Font,
    ):
        super ().__init__ (game)
        self.font = font

    def redraw (
            self,
    ) -> None:
        for i in range (4):
            self.game.screen.blit (
                    self.font.render (str (self.game.now), True, (0, 0, 0)),
                    (i % 2, i // 2 * 2))
        super ().redraw ()


class Balloon (GameObject):
    """
    吹出し

    Attributes:
        answer (str):           回答テキスト
        image_url (str, None):  画像 URL
        length (int):           表示する時間 (frame)
        query (str):            質問テキスト
        surface (Surface):      吹出し Surface
        x_flip (bool):          左右反転フラグ
        y_flip (bool):          上下反転フラグ
    """

    answer:     str         = ''
    image_url:  str | None  = None
    length:     int         = 300
    query:      str         = ''
    surface:    Surface
    x_flip:     bool        = False
    y_flip:     bool        = False

    def __init__ (
            self,
            game:   Game,
            x_flip: bool    = False,
            y_flip: bool    = False,
    ):
        super ().__init__ (game, enabled = False)
        self.x_flip = x_flip
        self.y_flip = y_flip
        self.surface = pygame.transform.scale (pygame.image.load ('assets/balloon.png'),
                                               (CWindow.WIDTH, CWindow.HEIGHT / 2))
        self.surface = pygame.transform.flip (self.surface, self.x_flip, self.y_flip)

    def redraw (
            self,
    ) -> None:
        if self.frame >= self.length:
            self.enabled = False
            self.game.last_answered_at = self.game.now
            return
        query = self.query
        if CommonModule.len_by_full (query) > 21:
            query = CommonModule.mid_by_full (query, 0, 19.5) + '...'
        answer = Surface ((800, ((CommonModule.len_by_full (self.answer) - 1) // 16 + 1) * 50),
                          pygame.SRCALPHA)
        for i in range (int (CommonModule.len_by_full (self.answer) - 1) // 16 + 1):
            answer.blit (DEERJIKA_FONT.render (
                             CommonModule.mid_by_full (self.answer, 16 * i, 16), True, (192, 0, 0)),
                         (0, 50 * i))
        surface = self.surface.copy ()
        surface.blit (USER_FONT.render ('>' + query, True, (0, 0, 0)), (120, 70))
        y:  int
        if self.frame < 30:
            y = 0
        elif self.frame >= self.length - 90:
            y = answer.get_height () - 100
        else:
            y = int ((answer.get_height () - 100) * (self.frame - 30) / (self.length - 120))
        surface.blit (answer, (100, 150), Rect (0, y, 800, 100))
        self.game.screen.blit (surface, (0, 0))
        super ().redraw ()

    def talk (
            self,
            query:      str,
            answer:     str,
            image_url:  str | None  = None,
            length:     int         = 300,
    ) -> None:
        self.query = query
        self.answer = answer
        self.image_url = image_url
        self.length = length
        self.frame = 0
        self.enabled = True


class KitaSun (GameObject):
    """
    き太く陽

    Attributes:
        sun (Sun):          ephem の太陽オブゼクト
        surface (Surface):  き太く陽 Surface
    """

    alt:        float
    az:         float
    sun:        Sun
    surface:    Surface

    def __init__ (
            self,
            game:   Game,
    ):
        super ().__init__ (game)
        self.surface = pygame.transform.scale (pygame.image.load ('assets/sun.png'), (200, 200))
        self.sun = Sun ()

    def redraw (
            self,
    ) -> None:
        surface = pygame.transform.rotate (self.surface, -(90 + math.degrees (self.arg)))
        self.game.screen.blit (surface, surface.get_rect (center = (self.x, self.y)))
        super ().redraw ()
        self.sun.compute (self.game.sky.observer)
        self.alt = self.sun.alt
        self.az = self.sun.az
        if abs (self.new_arg - self.arg) > math.radians (15):
            self.arg = self.new_arg
        self.x = self.new_x
        self.y = self.new_y

    @property
    def new_x (
            self,
    ) -> float:
        return CWindow.WIDTH * (math.degrees (self.az) - 80) / 120

    @property
    def new_y (
            self,
    ) -> float:
        return ((CWindow.HEIGHT / 2)
                - ((CWindow.HEIGHT / 2 + 100) * math.sin (self.alt)
                   / math.sin (math.radians (60))))

    @property
    def new_arg (
            self,
    ) -> float:
        return math.atan2 (self.new_y - self.y, self.new_x - self.x)


class Jojoko (GameObject):
    """
    大月ヨヨコ

    Attributes:
        base (Surface):     満月ヨヨコ Surface
        moon (Moon):        ephem の月オブゼクト
        surface (Surface):  缺けたヨヨコ
    """

    alt:        float
    az:         float
    base:       Surface
    moon:       Moon
    surface:    Surface

    def __init__ (
            self,
            game:   Game,
    ):
        super ().__init__ (game)
        self.base = pygame.transform.scale (pygame.image.load ('assets/moon.png'), (200, 200))
        self.moon = Moon ()
        self.surface = self._get_surface ()

    def redraw (
            self,
    ) -> None:
        if self.frame % (FPS * 3600) == 0:
            self.surface = self._get_surface ()
        surface = pygame.transform.rotate (self.surface, -(90 + math.degrees (self.arg)))
        surface.set_colorkey ((0, 255, 0))
        self.game.screen.blit (surface, surface.get_rect (center = (self.x, self.y)))
        super ().redraw ()
        self.moon.compute (self.game.sky.observer)
        self.alt = self.moon.alt
        self.az = self.moon.az
        if abs (self.new_arg - self.arg) > math.radians (15):
            self.arg = self.new_arg
        self.x = self.new_x
        self.y = self.new_y

    @property
    def phase (
            self,
    ) -> float:
        dt: datetime = ephem.localtime (ephem.previous_new_moon (self.game.sky.observer.date))
        return (self.game.now - dt).total_seconds () / 60 / 60 / 24

    def _get_surface (
            self,
    ) -> Surface:
        """
        ヨヨコを月齢に応じて缺かす．

        Returns:
            Surface:    缺けたヨヨコ
        """
        jojoko = self.base.copy ()
        for i in range (200):
            if 1 <= self.phase < 15:
                pygame.gfxdraw.bezier (jojoko, ((0, 100 + i), (100, 180 * self.phase / 7 - 80 + i), (200, 100 + i)), 3, (0, 255, 0))
            elif self.phase < 16:
                pass
            elif self.phase < 30:
                pygame.gfxdraw.bezier (jojoko, ((0, 100 - i), (100, 180 * (self.phase - 15) / 7 - 80 - i), (200, 100 - i)), 3, (0, 255, 0))
            else:
                jojoko.fill ((0, 255, 0))
        return jojoko

    @property
    def new_x (
            self,
    ) -> float:
        return CWindow.WIDTH * (math.degrees (self.az) - 80) / 120

    @property
    def new_y (
            self,
    ) -> float:
        return ((CWindow.HEIGHT / 2)
                - ((CWindow.HEIGHT / 2 + 100) * math.sin (self.alt)
                   / math.sin (math.radians (60))))

    @property
    def new_arg (
            self,
    ) -> float:
        return math.atan2 (self.new_y - self.y, self.new_x - self.x)


class Sky:
    """
    天体に関する情報を保持するクラス

    Attributes:
        observer (Observer):    観測値
    """

    observer:   Observer


class CWindow:
    """
    ウィンドゥに関する定数クラス

    Attributes:
        WIDTH (int):    ウィンドゥ幅
        HEIGHT (int):   ウィンドゥ高さ
    """

    WIDTH   = 1024
    HEIGHT  =  768


def get_surfaces_from_video (
        video_path: str,
) -> list[Surface]:
    cap = VideoCapture (video_path)
    if not cap.isOpened ():
        return []

    fps = cap.get (cv2.CAP_PROP_FPS)

    surfaces: list[Surface] = []
    while cap.isOpened ():
        (ret, frame) = cap.read ()
        if not ret:
            break
        frame = cv2.cvtColor (frame, cv2.COLOR_BGR2RGB)
        frame_surface = pygame.surfarray.make_surface (frame)
        frame_surface = pygame.transform.rotate (frame_surface, -90)
        surfaces.append (pygame.transform.flip (frame_surface, True, False))

    cap.release ()

    return surfaces


class Broadcast:
    chat:   PytchatCore

    def __init__ (
            self,
            broadcast_code,
    ):
        self.chat = pytchat.create (broadcast_code)

    def fetch_chat (
            self,
    ) -> Chat | None:
        if not self.chat.is_alive ():
            return None
        chats = self.chat.get ().items
        print (chats)
        if not chats:
            return None
        return random.choice (chats)


class Log:
    ...


def fetch_bytes_from_url (
        url:    str,
) -> bytes | None:
    res = requests.get (url, timeout = 60)
    if res.status_code != 200:
        return None
    return res.content


def add_query (
        broadcast:  Broadcast,
) -> None:
    chat = broadcast.fetch_chat ()
    if chat is None:
        return
    DB.begin_transaction ()
    chat.message = emoji.emojize (chat.message)
    message: str = chat.message
    user = (User.where ('platform', Platform.YOUTUBE.value)
                .where ('code', chat.author.channelId)
                .first ())
    if user is None:
        user = User ()
        user.platform = Platform.YOUTUBE.value
        user.code = chat.author.channelId
    user.name = chat.author.name
    user.icon = fetch_bytes_from_url (chat.author.imageUrl)
    user.save ()
    query = Query ()
    query.user_id = user.id
    query.target_character = Character.DEERJIKA.value
    query.content = chat.message
    query.query_type = QueryType.YOUTUBE_COMMENT.value
    query.model = GPTModel.GPT3_TURBO.value
    query.sent_at = datetime.now ()
    query.answered = False
    query.save ()
    DB.commit ()


if __name__ == '__main__':
    main ()