From 2d2ed5b0d62b7fd639dfcb547282a3acd4b43ea3 Mon Sep 17 00:00:00 2001
From: Vadym <62544653+icenfame@users.noreply.github.com>
Date: Sat, 12 Mar 2022 22:17:50 +0200
Subject: [PATCH] feat: auto-play game support

---
 bot.py | 80 +++++++++++++++++++++++++++++++++++++++++-----------------
 1 file changed, 57 insertions(+), 23 deletions(-)

diff --git a/bot.py b/bot.py
index 47cd450..06132c0 100644
--- a/bot.py
+++ b/bot.py
@@ -1,4 +1,6 @@
-"""Automation module"""
+"""Bot for Telegram game LumberJack"""
+import sys
+from time import sleep
 import pyautogui
 import numpy as np
 import keyboard
@@ -14,6 +16,10 @@ def detect_game():
         pyautogui.locateOnScreen('regions/buttons/left.png')
     )
 
+    if not button:
+        sys.exit(
+            "[ERROR]: Cannot find game. Make sure you have opened game on your screen")
+
     canvas_width = 600
     white = (255, 255, 255)
     background = (51, 51, 51)
@@ -46,10 +52,9 @@ def detect_game():
 
 
 def screenshot(region):
-    """Make screenshot of specific region"""
+    """Take a screenshot of a specific region"""
 
-    img = pyautogui.screenshot(region=region)
-    return img
+    return pyautogui.screenshot(region=region)
 
 
 def check_player(region, img):
@@ -61,58 +66,87 @@ def check_player(region, img):
 
     if img.getpixel((head_x, region[3] - head_y)) == head_color:
         return "left"
-    elif img.getpixel((region[2] - head_x, region[3] - head_y)) == head_color:
+
+    if img.getpixel((region[2] - head_x, region[3] - head_y)) == head_color:
         return "right"
-    else:
-        exit("Player not found")
+
+    return None
 
 
-def check_wood(region, img, player):
-    """Check wood over the player head"""
+def check_branch(region, img, player):
+    """Check branch over the player head"""
 
-    wood_y = region[3] - 498
-    wood_w = 32
-    wood_h = 100
+    branch_y = region[3] - 498
+    branch_w = 32
+    branch_h = 100
+    branch_color = (126, 173, 79)
 
     if player == "left":
-        wood_x = 230
+        branch_x = 230
     elif player == "right":
-        wood_x = region[2] - 230 - wood_w
+        branch_x = region[2] - 230 - branch_w
 
-    # TODO find all woods by pixel
-
-    img = img.crop((wood_x, wood_y, wood_x + wood_w, wood_y + wood_h))
+    img = img.crop((branch_x, branch_y, branch_x +
+                    branch_w, branch_y + branch_h))
     arr = np.array(img)
     arr = arr.reshape((-1, 3))
 
-    if [126, 173, 79] in arr:
+    if branch_color in arr:
         return player
 
+    return None
+
+
+def play(region):
+    """Auto-play game"""
+
+    play_button = (
+        pyautogui.locateOnScreen('regions/buttons/play.png') or
+        pyautogui.locateOnScreen('regions/buttons/replay.png')
+    )
+
+    if play_button:
+        pyautogui.click(play_button)
+
+    pyautogui.click(region)
+
 
 def main():
     """Main function"""
+
     region = detect_game()
+    score = 0
 
-    # TODO auto start game & end if lost
+    play(region)
+    sleep(1)
 
     while True:
         if keyboard.is_pressed('esc'):
-            break
+            sys.exit("[INFO]: User ESC exit")
 
         img = screenshot(region)
         player = check_player(region, img)
-        wood = check_wood(region, img, player)
 
-        if player == wood == "left":
+        if not player:
+            if score > 0:
+                print("[INFO]: Score", score - 2)
+
+            break
+
+        branch = check_branch(region, img, player)
+
+        if player == branch == "left":
             pyautogui.press("right")
             pyautogui.press("right")
-        elif player == wood == "right":
+        elif player == branch == "right":
             pyautogui.press("left")
             pyautogui.press("left")
         else:
             pyautogui.press(player)
             pyautogui.press(player)
 
+        score += 2
+
 
 if __name__ == "__main__":
     main()