diff --git a/PresenceTracker.py b/PresenceTracker.py new file mode 100644 index 0000000..9aea581 --- /dev/null +++ b/PresenceTracker.py @@ -0,0 +1,118 @@ +import sqlite3 +from datetime import datetime +from loguru import logger +from datetime_utils import to_unix + + +def get_db() -> sqlite3.Connection: + SCHEMA = """ + create table if not exists activities ( + id integer primary key autoincrement, + name text not null unique + ); + + create table if not exists users_cache ( + user_id integer primary key, + nickname text not null + ); + + create table if not exists presence_journal ( + user_id integer not null, + start_time integer not null, -- In unix timestamp + end_time integer not null, + activity_name_id integer not null, + primary key (user_id, start_time), + foreign key (activity_name_id) references activities(id) + );""" + + db = sqlite3.connect('presence-tracker.sqlite') + db.executescript(SCHEMA) + return db + + +class PresenceTracker: + def __init__(self, db: sqlite3.Connection = get_db()): + self.db = db + + def log_activity(self, user_id: int, activity_name: str, start_time: datetime, end_time: datetime | None = None): + """Create record or update uncommited end_time""" + + # end_time behaves like "last seen with this activity" + + logger.info(f'Logging {user_id=}, {activity_name=}, {start_time.isoformat()=}, {end_time=}') + + activity_id = self.get_activity_id(activity_name) + _end_time = self.to_unix_default(end_time) # Defaulted end_time + + unix_start_time = to_unix(start_time) + + if unix_start_time > _end_time: + logger.warning(f'Got {unix_start_time=} > {_end_time}, {_end_time - unix_start_time=}, {user_id=}, {activity_name=}') + + + with self.db: + self.db.execute( + """insert into presence_journal (user_id, start_time, activity_name_id, end_time) + values (:user_id, :start_time, :activity_name_id, :end_time) + on conflict do update set end_time = :end_time;""", + { + 'user_id': user_id, + 'start_time': to_unix(start_time), + 'activity_name_id': activity_id, + 'end_time': _end_time + } + ) + + def get_activity_id(self, activity_name: str) -> int: + # Try to find first, insert after + params = (activity_name,) + query = self.db.execute('select id from activities where name = ?;', params).fetchone() + if query is not None: + logger.trace(f'Found record for {activity_name!r} id={query}') + return query[0] + + with self.db: + query = self.db.execute('insert into activities (name) values (?) returning id;', params).fetchone() + + logger.debug(f'Inserted record for {activity_name!r} {query}') + return query[0] + + def saturate_users_cache(self, user_id: int, user_name: str): + logger.trace(f'Adding {user_name!r} to cache') + with self.db: + self.db.execute( + 'insert or ignore into users_cache (user_id, nickname) VALUES (?, ?);', + (user_id, user_name) + ) + + def user_breakdown(self, user_id: int) -> dict[str, int]: + """Return dict with activities names as keys and spent time in seconds as values""" + res = self.db.execute( + 'select a.name, round(sum(end_time - start_time) / 3600.0, 1) as total ' + 'from presence_journal ' + 'left join activities a on a.id = presence_journal.activity_name_id ' + 'where user_id = ? ' + 'group by activity_name_id ' + 'order by total desc;', + (user_id,) + ).fetchall() + + return {row[0]: row[1] for row in res} + + def top_users(self, last_days: int | None) -> dict[int, tuple[str | None, int]]: + """Returns keys - ids, values - tuple(nickname?, seconds spent)""" # TODO: implement last_days + res = self.db.execute( + """select uc.user_id, uc.nickname, round(sum(end_time - presence_journal.start_time) / 3600.0, 1) as total + from presence_journal left join users_cache uc on presence_journal.user_id = uc.user_id + group by uc.user_id + order by total desc + limit 50;""").fetchall() + + return {row[0]: row[1:3] for row in res} + + @staticmethod + def default_end_time(end_time: datetime | None) -> datetime: + return datetime.now() if end_time is None else end_time + + def to_unix_default(self, end_time: datetime | None) -> int: + return to_unix(self.default_end_time(end_time)) diff --git a/datetime_utils.py b/datetime_utils.py new file mode 100644 index 0000000..40632bd --- /dev/null +++ b/datetime_utils.py @@ -0,0 +1,6 @@ +from datetime import datetime + +def to_unix(timestamp: datetime) -> int: + """Convert datetime.datetime to unix timestamp as int (since we store all timestamps in unix format in DB""" + return int(datetime.timestamp(timestamp)) + diff --git a/extensions/Frontend.py b/extensions/Frontend.py new file mode 100644 index 0000000..6825a31 --- /dev/null +++ b/extensions/Frontend.py @@ -0,0 +1,77 @@ +from typing import TYPE_CHECKING, Sequence + +import discord +import discord.ext.commands as commands +from loguru import logger + + +if TYPE_CHECKING: + from main import Bot + + +def build_table_embed(*columns: Sequence[object], name: str, description: str, **kwargs) -> discord.Embed: + """First value of one column is a header""" + + embed = discord.Embed( + title=name, + description=description, + **kwargs + ) + + for column in columns: + column_iter = iter(column) + embed.add_field(name=next(column_iter), value='\n'.join(str(cell) for cell in column_iter), inline=True) + + return embed + + +class Frontend(commands.Cog): + def __init__(self, bot: 'Bot'): + self.bot = bot + + @discord.app_commands.command() + async def stats(self, interaction: discord.Interaction, user: discord.User | None): + """ + Show how many hours a user spent in games per game (default - You) + + :param interaction: + :param user: User for who we will get stats, You by default + """ + + _user = interaction.user if user is None else user + + logger.trace(f'Starting my_stats for {_user}') + stats: dict[str, int] = self.bot.activity_tracker.user_breakdown(_user.id) + embed = build_table_embed( + ('Games', *stats.keys()), + ('Hours spent', *stats.values()), + name=f'Gaming stats for {str(_user)}', + description=f'Totally played {sum(stats.values())} hours' + ) + await interaction.response.send_message(embed=embed) + + @discord.app_commands.command() + async def top_users( + self, + interaction: discord.Interaction + # ,last_days_count: discord.app_commands.Range[int, 1, None] + ): + """ + Show top of players by spent in games hours + + :param interaction: + :param last_days_count: Show stats for last number of days, i.e. for last 14 days (all time for default) + :return: + """ + data = self.bot.activity_tracker.top_users(0) + embed = build_table_embed( + ('Username', *(value[0] for value in data.values())), + ('Spent hours', *(value[1] for value in data.values())), + name='Top users by spent in games', + description='Limited up to 50 entries' + ) + await interaction.response.send_message(embed=embed) + + +async def setup(bot: 'Bot'): + await bot.add_cog(Frontend(bot)) diff --git a/main.py b/main.py new file mode 100644 index 0000000..31683ba --- /dev/null +++ b/main.py @@ -0,0 +1,122 @@ +import asyncio +import os +import signal +import logging +import sys +from typing import Any +from itertools import chain + +from loguru import logger +from PresenceTracker import PresenceTracker +import discord +from discord.ext.commands import Bot as BotBase +from pathlib import Path + +""" +Cases should be considered: +Case 1 None -> Playing +Case 2 Playing -> None +Case 3 Playing -> bot shutdown -> Already stopped playing (just commit last end_time?) +Case 4 None -> bot shutdown -> Already playing (catch on stop playing, so Case 2) +Case 5 Playing -> Gone invisible -> Gone visible with the same activity (handles by Case 1 and Case 2) +Case 6: Playing -> Still playing (for end_time updating) + +TODO: Handle negative diff +""" + + +class InterceptHandler(logging.Handler): + def emit(self, record): + # Get corresponding Loguru level if it exists + try: + level = logger.level(record.levelname).name + except ValueError: + level = record.levelno + + # Find caller from where originated the logged message + frame, depth = logging.currentframe(), 2 + while frame.f_code.co_filename == logging.__file__: + frame = frame.f_back + depth += 1 + + logger.opt(depth=depth, exception=record.exc_info).log(level, record.getMessage()) + +logging.basicConfig(handlers=[InterceptHandler()], level=logging.INFO) + + +class Bot(BotBase): + def __init__(self, *, intents: discord.Intents, activity_tracker: PresenceTracker, **options: Any): + super().__init__(intents=intents, **options) + self.activity_tracker = activity_tracker + + async def setup_hook(self) -> None: + logger.info(f'Invite URL: {"No self.user" if self.user is None else discord.utils.oauth_url(self.user.id)}') + + for filepath in Path('extensions').glob('*.py'): + filename = filepath.name + logger.info(f'Loading extension {filename}') + await self.load_extension(f"extensions.{filename[:-3]}") + + # self.tree.clear_commands(guild=discord.Object(648268554386407432)) + # await self.tree.sync(guild=discord.Object(648268554386407432)) + # + # self.tree.copy_global_to(guild=discord.Object(648268554386407432)) + # await self.tree.sync(guild=discord.Object(648268554386407432)) + + # await self.tree.sync() + + + def signal_handler(self, signame, _): + logger.info(f'Got {signame} signal, shutting down') + self.loop.create_task(self.close()) + + async def on_ready(self): + logger.info(f'Ready {self.user}') + + async def on_presence_update(self, before: discord.Member, after: discord.Member): + if before.bot: + return + + for activity in chain(before.activities, after.activities): + if activity.type == discord.ActivityType.playing: + if isinstance(activity, (discord.Game, discord.Activity)): + if activity.name is not None and activity.start is not None: + self.activity_tracker.log_activity(after.id, activity.name, activity.start, activity.end) + self.activity_tracker.saturate_users_cache(after.id, after.name + '#' + after.discriminator) + + else: + logger.warning(f'Got activity with missing name or start: {activity.to_dict()!r}') + + else: + logger.warning( + f'Got discord.ActivityType.playing with unusual type: {type(activity)}; {activity.to_dict()}') + + else: + logger.trace(f'Got not playing activity: {activity.to_dict()}') + + +async def async_main(): + # logger.disable('discord') + logger.remove() + logger.add(sink=sys.stderr, level='TRACE') + + intents = discord.Intents.default() + intents.presences = True + intents.members = True + + activity_tracker = PresenceTracker() + bot = Bot(intents=intents, activity_tracker=activity_tracker, command_prefix='') + + for sig in (signal.SIGTERM, signal.SIGINT): + signal.signal(sig, bot.signal_handler) + + await bot.start(os.environ['TOKEN']) + + +def main(): + loop = asyncio.new_event_loop() + loop.run_until_complete(async_main()) + + +if __name__ == '__main__': + main() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..fb21c91 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +discord==2.0.0 +loguru==0.6.0