From caba2a708fa0a94642d0135c318d0afd0819b70c Mon Sep 17 00:00:00 2001
From: Alex <alexkopylov123@gmail.com>
Date: Sat, 6 Apr 2024 18:11:36 +0400
Subject: [PATCH] fix error with poll

---
 src/data_models/Game.py | 99 +++++++++++++++++++++++++++++++++--------
 src/handlers/save.py    |  5 +--
 2 files changed, 82 insertions(+), 22 deletions(-)

diff --git a/src/data_models/Game.py b/src/data_models/Game.py
index ce011d3..cc5935a 100644
--- a/src/data_models/Game.py
+++ b/src/data_models/Game.py
@@ -1,8 +1,7 @@
-from typing import Literal
+from typing import Literal, Tuple
 
 from pydantic import BaseModel, field_validator
-
-from src import config
+from collections import Counter
 from src.data_models.PollResult import PollResult
 
 
@@ -14,28 +13,92 @@ class Game(BaseModel):
     ]  # Literal["Hitler Chancellor", "Fascist Law", "Hitler Death", "Liberal Law"]
     creator_id: int
 
+    def extract_player_outcomes(cls, results: Tuple[PollResult, ...]) -> Counter:
+        outcomes = [outcome.get_answer_as_text() for outcome in results]
+        return Counter(outcomes)
+
+    def remove_spectators(cls, outcomes_counter: Counter) -> None:
+        outcomes_counter.pop("👀 SPECTATOR | NOT A PLAYER 👀", None)
+
+    def count_player_roles(cls, outcomes_counter: Counter) -> tuple[int, int, int]:
+        total_hitlers = sum(
+            outcomes_counter[role]
+            for role in [
+                "I'm Chancellor Hitler",
+                "I'm Dead Hitler",
+                "I'm Hitler Loser",
+                "I'm Hitler Winner",
+            ]
+        )
+        total_liberals = (
+            outcomes_counter["I'm Liberal Winner"]
+            + outcomes_counter["I'm Liberal Loser"]
+        )
+        total_fascists = (
+            outcomes_counter["I'm Fascistic Winner"]
+            + outcomes_counter["I'm Fascistic Loser"]
+            + total_hitlers
+        )
+        return total_hitlers, total_liberals, total_fascists
+
+    def validate_player_distribution(
+        cls,
+        total_hitlers: int,
+        total_liberals: int,
+        total_fascists: int,
+        max_hitlers: int = 1,
+        max_liberals: int = 6,
+        max_fascists: int = 4,
+    ) -> None:
+        if (
+            total_hitlers > max_hitlers
+            or total_liberals > max_liberals
+            or total_fascists > max_fascists
+        ):
+            raise ValueError("Invalid player distribution according to game rules.")
+
+    def check_mutually_exclusive_victory_conditions(
+        cls, outcomes_counter: Counter
+    ) -> None:
+        liberal_win = (
+            outcomes_counter["I'm Liberal Winner"] > 0
+            or outcomes_counter["I'm Dead Hitler"] > 0
+            or outcomes_counter["I'm Fascistic Loser"] > 0
+            or outcomes_counter["I'm Hitler Loser"] > 0
+        )
+
+        fascist_win = (
+            outcomes_counter["I'm Fascistic Winner"] > 0
+            or outcomes_counter["I'm Hitler Winner"] > 0
+            or outcomes_counter["I'm Chancellor Hitler"] > 0
+            or outcomes_counter["I'm Liberal Loser"] > 0
+        )
+        if liberal_win and fascist_win:
+            raise ValueError("Invalid results: Winners from both teams cannot exist.")
+
     @field_validator("results", mode="after")
-    @classmethod
     def validate_results(
         cls, results: tuple[PollResult]
     ) -> Literal["CH", "DH", "FW", "LW"]:
-        outcomes = set(outcome.get_answer_as_text() for outcome in results)
-        if "I'm Chancellor Hitler" in outcomes:
+        outcomes_counter = cls.extract_player_outcomes(results)
+        cls.remove_spectators(outcomes_counter)
+        total_hitlers, total_liberals, total_fascists = cls.count_player_roles(
+            outcomes_counter
+        )
+        cls.validate_player_distribution(total_hitlers, total_liberals, total_fascists)
+        cls.check_mutually_exclusive_victory_conditions(outcomes_counter)
+
+        # Determine outcome based on specific conditions
+        if outcomes_counter["I'm Chancellor Hitler"] > 0:
             return "CH"
-        if "I'm Dead Hitler" in outcomes:
+        if outcomes_counter["I'm Dead Hitler"] > 0:
             return "DH"
-        if (
-            "I'm Liberal Winner"
-            or "I'm Hitler Looser"
-            or "I'm Fascistic Looser" in outcomes
-        ):
+        if outcomes_counter["I'm Liberal Winner"] > 0:
             return "LW"
         if (
-            "I'm Fascistic Winner"
-            or "I'm Hitler Winner"
-            or "I'm Liberal Looser" in outcomes
+            outcomes_counter["I'm Fascistic Winner"] > 0
+            or outcomes_counter["I'm Hitler Winner"] > 0
         ):
             return "FW"
-        raise ValueError(
-            f"Invalid results '{results}' for Game. Results must be one of {config.GAME_POLL_OUTCOMES}"
-        )
+
+        raise ValueError("Invalid game results: No clear win condition met.")
diff --git a/src/handlers/save.py b/src/handlers/save.py
index 34d1236..199c071 100644
--- a/src/handlers/save.py
+++ b/src/handlers/save.py
@@ -86,12 +86,10 @@ async def save(
 ) -> None:
     msg_with_poll = update.effective_message.reply_to_message
     if await _pass_checks(msg_with_poll=msg_with_poll, update=update, context=context):
-        await context.bot.stop_poll(update.effective_chat.id, msg_with_poll.message_id)
         poll_id = int(msg_with_poll.poll.id)
         poll_data, poll_results = await asyncio.gather(
             fetch_poll_data(poll_id), fetch_poll_results(poll_id)
         )
-
         records = [
             Record(
                 creator_id=poll_data.creator_id,
@@ -102,14 +100,13 @@ async def save(
             )
             for results in poll_results
         ]
-
         game = Game(
             poll_id=poll_data.message_id,
             chat_id=poll_data.chat_id,
             creator_id=poll_data.creator_id,
             results=poll_results,
         )
-
+        await context.bot.stop_poll(update.effective_chat.id, msg_with_poll.message_id)
         # Execute post-game tasks
         await asyncio.gather(
             save_game(game),