TgToFileSystem/backend/UserManager.py

import os
from enum import Enum, IntEnum, unique, auto
import sqlite3
import logging
import datetime

from pydantic import BaseModel
from telethon import types

logger = logging.getLogger(__file__.split("/")[-1])

class UserUpdateParam(BaseModel):
    client_id: str
    username: str
    phone: str
    tg_user_id: int
    last_login_time: int


class MessageUpdateParam(BaseModel):
    unique_id: str
    user_id: int
    chat_id: int
    msg_id: int
    msg_type: str
    msg_ctx: str
    file_type: str
    file_name: str
    msg_js: str


class UserManager(object):
    def __init__(self) -> None:
        self.con = sqlite3.connect(f"{os.path.dirname(__file__)}/db/user.db")
        self.cur = self.con.cursor()
        if not self._table_has_been_inited():
            self._first_runtime_run_once()

    def __del__(self) -> None:
        self.con.commit()
        self.con.close()

    def update_user(self) -> None:
        raise NotImplementedError

    def update_message(self) -> None:
        raise NotImplementedError
    
    def generate_unique_id_by_msg(self, me: types.User, msg: types.Message) -> str:
        user_id = me.id
        chat_id = msg.chat_id
        msg_id = msg.id
        unique_id = str(user_id) + str(chat_id) + str(msg_id)
        return unique_id

    def get_all_msg_by_chat_id(self, chat_id: int) -> list[any]:
        res = self.cur.execute(
            "SELECT * FROM message WHERE chat_id == ? ORDER BY date_time DESC", (chat_id,))
        return res.fetchall()

    def get_msg_by_chat_id_and_keyword(self, chat_id: int, keyword: str, limit: int = 10, offset: int = 0, inc: bool = False, ignore_case: bool = False) -> list[any]:
        keyword_condition = "msg_ctx LIKE '%{key}%' OR file_name LIKE '%{key}%'"
        if ignore_case:
            keyword_condition = "LOWER(msg_ctx) LIKE LOWER('%{key}%') OR LOWER(file_name) LIKE LOWER('%{key}%')"
        keyword_condition = keyword_condition.format(key=keyword)
        execute_script = f"SELECT * FROM message WHERE chat_id == {chat_id} AND ({keyword_condition}) ORDER BY date_time {'' if inc else 'DESC '}LIMIT {limit} OFFSET {offset}"
        res = self.cur.execute(execute_script)
        return res

    def get_oldest_msg_by_chat_id(self, chat_id: int) -> list[any]:
        res = self.cur.execute(
            "SELECT * FROM message WHERE chat_id == ? ORDER BY date_time LIMIT 1", (chat_id,))
        return res.fetchall()

    def get_newest_msg_by_chat_id(self, chat_id: int) -> list[any]:
        res = self.cur.execute(
            "SELECT * FROM message WHERE chat_id == ? ORDER BY date_time DESC LIMIT 1", (chat_id,))
        return res.fetchall()

    def get_msg_by_unique_id(self, unique_id: str) -> list[any]:
        res = self.cur.execute(
            "SELECT * FROM message WHERE unique_id == ? ORDER BY date_time DESC LIMIT 1", (unique_id,))
        return res.fetchall()

    @unique
    class MessageTypeEnum(Enum):
        OTHERS = "others"
        TEXT = "text"
        PHOTO = "photo"
        FILE = "file"

    @unique
    class ColumnStrEnum(Enum):
        UNIQUE_ID = "unique_id"
        USER_ID = "user_id"
        CHAT_ID = "chat_id"
        MSG_ID = "msg_id"
        MSG_TYPE = "msg_type"
        MSG_CTX = "msg_ctx"
        MIME_TYPE = "mime_type"
        FILE_NAME = "file_name"
        MSG_JS = "msg_js"
        DATE_TIME = "date_time"

    def insert_by_message(self, me: types.User, msg: types.Message):
        user_id = me.id
        chat_id = msg.chat_id
        msg_id = msg.id
        unique_id = str(user_id) + str(chat_id) + str(msg_id)
        msg_type = UserManager.MessageTypeEnum.OTHERS.value
        mime_type = ""
        file_name = ""
        msg_ctx = msg.message
        msg_js = msg.to_json()
        date_time = int(msg.date.timestamp() * 1_000) * 1_000_000
        try:
            if msg.media is None:
                msg_type = UserManager.MessageTypeEnum.TEXT.value
            elif isinstance(msg.media, types.MessageMediaPhoto):
                msg_type = UserManager.MessageTypeEnum.PHOTO.value
            elif isinstance(msg.media, types.MessageMediaDocument):
                document = msg.media.document
                mime_type = document.mime_type
                for attr in document.attributes:
                    if isinstance(attr, types.DocumentAttributeFilename):
                        file_name = attr.file_name
                msg_type = UserManager.MessageTypeEnum.FILE.value
        except Exception as err:
            logger.error(f"{err=}")
        insert_data = (unique_id, user_id, chat_id, msg_id,
                       msg_type, msg_ctx, mime_type, file_name, msg_js, date_time)
        execute_script = "INSERT INTO message (unique_id, user_id, chat_id, msg_id, msg_type, msg_ctx, mime_type, file_name, msg_js, date_time) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
        try:
            self.cur.execute(execute_script, insert_data)
            self.con.commit()
        except Exception as err:
            logger.error(f"{err=}")

    @unique
    class ColumnEnum(IntEnum):
        UNIQUE_ID = 0
        USER_ID = auto()
        CHAT_ID = auto()
        MSG_ID = auto()
        MSG_TYPE = auto()
        MSG_CTX = auto()
        MIME_TYPE = auto()
        FILE_NAME = auto()
        MSG_JS = auto()
        DATE_TIME = auto()
        COLUMN_LEN = auto()

    def get_column_by_enum(self, column: tuple[any], index: ColumnEnum) -> any:
        if len(column) == UserManager.ColumnEnum.COLUMN_LEN:
            return column[index]
        return None

    def get_column_msg_id(self, column: tuple[any]) -> int | None:
        if len(column) == UserManager.ColumnEnum.COLUMN_LEN:
            return column[UserManager.ColumnEnum.MSG_ID]
        return None

    def get_column_msg_js(self, column: tuple[any]) -> str | None:
        if len(column) == UserManager.ColumnEnum.COLUMN_LEN:
            return column[UserManager.ColumnEnum.MSG_JS]
        return None

    def get_user_info() -> None:
        raise NotImplementedError

    def _table_has_been_inited(self) -> bool:
        res = self.cur.execute("SELECT name FROM sqlite_master")
        return len(res.fetchall()) != 0

    def _first_runtime_run_once(self) -> None:
        if len(self.cur.execute("SELECT name FROM sqlite_master WHERE name='user'").fetchall()) == 0:
            self.cur.execute(
                "CREATE TABLE user(client_id primary key, username, phone, tg_user_id, last_login_time)")
        if len(self.cur.execute("SELECT name FROM sqlite_master WHERE name='message'").fetchall()) == 0:
            self.cur.execute(
                "CREATE TABLE message(unique_id varchar(64) primary key, user_id int NOT NULL, chat_id int NOT NULL, msg_id int NOT NULL, msg_type varchar(64), msg_ctx, mime_type, file_name, msg_js, date_time int)")


if __name__ == "__main__":
    db = UserManager()
    # db.cur.execute(
    #     "UPDATE user SET (client_id, username, phone) = (123, 'hehe', 66666) WHERE client_id == 123")
    # res = db.cur.execute("SELECT name FROM sqlite_master")
    # print(res.fetchall())
    # res = db.cur.execute(
    #     "SELECT msg_id, msg_ctx, file_name FROM message WHERE chat_id == -1001216816802")
    # # res.execute("SELECT * FROM message WHERE chat_id == ? ORDER BY msg_id DESC LIMIT 1", (-1001216816802,))
    # # print(res.fetchall())
    # # print("\n\n\n\n\n\n")
    # res.execute("SELECT COUNT(msg_id) FROM message")
    # # res = db.cur.execute("SELECT DISTINCT chat_id FROM message")
    # print(res.fetchall())
    # db.cur.execute("SELECT * FROM")