From ca9ccbe8a37b6eef5c78bfabae7ae69d5faff841 Mon Sep 17 00:00:00 2001
From: Kevin Trocolli <pitok236@gmail.com>
Date: Fri, 21 Jun 2024 23:00:35 -0400
Subject: [PATCH] chuni: add missing net battle funcs

---
 .../1e150d16ab6b_chuni_add_net_battle_uk.py   | 27 ++++++++++++
 titles/chuni/schema/item.py                   |  2 +-
 titles/chuni/schema/profile.py                | 42 ++++++++++++++++++-
 3 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 core/data/alembic/versions/1e150d16ab6b_chuni_add_net_battle_uk.py

diff --git a/core/data/alembic/versions/1e150d16ab6b_chuni_add_net_battle_uk.py b/core/data/alembic/versions/1e150d16ab6b_chuni_add_net_battle_uk.py
new file mode 100644
index 0000000..5fcd5a8
--- /dev/null
+++ b/core/data/alembic/versions/1e150d16ab6b_chuni_add_net_battle_uk.py
@@ -0,0 +1,27 @@
+"""chuni_add_net_battle_uk
+
+Revision ID: 1e150d16ab6b
+Revises: b23f985100ba
+Create Date: 2024-06-21 22:57:18.418488
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import mysql
+
+# revision identifiers, used by Alembic.
+revision = '1e150d16ab6b'
+down_revision = 'b23f985100ba'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_unique_constraint(None, 'chuni_profile_net_battle', ['user'])
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    op.drop_constraint(None, 'chuni_profile_net_battle', type_='unique')
+    # ### end Alembic commands ###
diff --git a/titles/chuni/schema/item.py b/titles/chuni/schema/item.py
index 2f386ef..30db4b8 100644
--- a/titles/chuni/schema/item.py
+++ b/titles/chuni/schema/item.py
@@ -262,7 +262,7 @@ cmission_progress = Table(
     "chuni_item_cmission_progress",
     metadata,
     Column("id", Integer, primary_key=True, nullable=False),
-    Column("user", ForeignKey("aime_user.id", ondelete="cascade"), nullable=False),
+    Column("user", ForeignKey("aime_user.id", ondelete="cascade", onupdate="cascade"), nullable=False),
     Column("missionId", Integer, nullable=False),
     Column("order", Integer),
     Column("stage", Integer),
diff --git a/titles/chuni/schema/profile.py b/titles/chuni/schema/profile.py
index 2f8bce3..7ee0164 100644
--- a/titles/chuni/schema/profile.py
+++ b/titles/chuni/schema/profile.py
@@ -412,6 +412,18 @@ rating = Table(
     mysql_charset="utf8mb4",
 )
 
+net_battle = Table(
+    "chuni_profile_net_battle",
+    metadata,
+    Column("id", Integer, primary_key=True, nullable=False),
+    Column("user", Integer, ForeignKey("aime_user.id", ondelete="cascade", onupdate="cascade"), nullable=False, unique=True),
+    Column("isRankUpChallengeFailed", Boolean),
+    Column("highestBattleRankId", Integer),
+    Column("battleIconId", Integer),
+    Column("battleIconNum", Integer),
+    Column("avatarEffectPoint", Integer),
+    mysql_charset="utf8mb4",
+)
 
 class ChuniProfileData(BaseData):
     async def update_name(self, user_id: int, new_name: str) -> bool:
@@ -779,4 +791,32 @@ class ChuniProfileData(BaseData):
         else:
             versions_raw = result.fetchall()
             versions = [row[0] for row in versions_raw]
-        return sorted(versions, reverse=True)
\ No newline at end of file
+        return sorted(versions, reverse=True)
+
+    async def put_net_battle(self, user_id: int, net_battle_data: Dict) -> Optional[int]:
+        sql = insert(net_battle).values(
+            user=user_id,
+            isRankUpChallengeFailed=net_battle_data['isRankUpChallengeFailed'],
+            highestBattleRankId=net_battle_data['highestBattleRankId'],
+            battleIconId=net_battle_data['battleIconId'],
+            battleIconNum=net_battle_data['battleIconNum'],
+            avatarEffectPoint=net_battle_data['avatarEffectPoint'],
+        )
+
+        conflict = sql.on_duplicate_key_update(
+            isRankUpChallengeFailed=net_battle_data['isRankUpChallengeFailed'],
+            highestBattleRankId=net_battle_data['highestBattleRankId'],
+            battleIconId=net_battle_data['battleIconId'],
+            battleIconNum=net_battle_data['battleIconNum'],
+            avatarEffectPoint=net_battle_data['avatarEffectPoint'],
+        )
+
+        result = await self.execute(conflict)
+        if result:
+            return result.inserted_primary_key['id']
+        self.logger.error(f"Failed to put net battle data for user {user_id}")
+
+    async def get_net_battle(self, user_id: int, net_battle_data: Dict) -> Optional[Row]:
+        result = await self.execute(net_battle.select(net_battle.c.user == user_id))
+        if result:
+            return result.fetchone()