main.py

import pyautogui
import numpy as np
from ultralytics import YOLO
import time, random, keyboard, cv2, configparser, mss, os, threading, sys, tkinter as tk
from tkinter import ttk, messagebox
from pystray import MenuItem as item, Icon as icon, Menu
from PIL import Image, ImageDraw

# --- Класс для управления конфигурацией ---
class Config:
    def __init__(self, filename='config.ini'):
        self.filename = filename
        self.parser = configparser.ConfigParser()
        if not os.path.exists(filename):
            sys.exit(f"Ошибка: Файл конфигурации {filename} не найден.")
        self.load()

    def load(self):
        self.parser.read(self.filename)
        # --- BOT_CONFIG ---
        bot_cfg = self.parser['BOT_CONFIG']
        self.model_path = bot_cfg.get('model_path', 'best.pt')
        self.fishing_timeout = (bot_cfg.getint('fishing_timeout_base', 17), bot_cfg.getint('fishing_timeout_scatter', 2))
        self.canny_threshold_multiplier = bot_cfg.getfloat('canny_threshold_multiplier', 1.3)
        self.canny_avg_window_size = bot_cfg.getint('canny_avg_window_size', 30)
        self.max_missed_detections = bot_cfg.getint('max_missed_detections', 1)
        self.skip_chance = (bot_cfg.getint('chance_to_skip_catch_base', 5), bot_cfg.getint('chance_to_skip_catch_scatter', 2))
        self.recast_chance = (bot_cfg.getint('chance_to_recast_base', 6), bot_cfg.getint('chance_to_recast_scatter', 3))
        self.recast_check_interval = bot_cfg.getint('recast_check_interval', 5)
        # Раздельные пороги для зон безопасности
        self.minimap_pixel_threshold = bot_cfg.getfloat('minimap_pixel_threshold', 0.05)
        self.hpbar_pixel_threshold = bot_cfg.getfloat('hpbar_pixel_threshold', 0.05)

        # --- COORDINATES ---
        coords = self.parser['COORDINATES']
        self.scan_zone = (coords.getint('x', 200), coords.getint('y', 300), coords.getint('width', 1520), coords.getint('height', 480))
        self.center_x = self.scan_zone[0] + self.scan_zone[2] // 2
        self.center_y = self.scan_zone[1] + self.scan_zone[3] // 2

        # --- SAFETY_ZONES ---
        safety = self.parser['SAFETY_ZONES']
        self.minimap_zone = (safety.getint('minimap_x', 1750), safety.getint('minimap_y', 50), safety.getint('minimap_w', 150), safety.getint('minimap_h', 150))
        self.hpbar_zone = (safety.getint('hpbar_x', 50), safety.getint('hpbar_y', 50), safety.getint('hpbar_w', 300), safety.getint('hpbar_h', 50))

    def save_setting(self, section, key, value):
        if not self.parser.has_section(section): self.parser.add_section(section)
        self.parser.set(section, str(key), str(value))
        with open(self.filename, 'w') as configfile: self.parser.write(configfile)

# --- Класс для управления состоянием бота ---
class State:
    def __init__(self):
        self.running = False
        self.show_hud = False
        self.is_child_gui_active = False
        self.fishing_timer_start = 0
        self.current_fishing_timeout = 0
        self.edge_counts = []
        self.target_roi = None
        self.bobber_detections = 0
        self.missed_detections = 0
        self.last_bobber_center = None
        self.last_safety_check_time = 0.0
        self.last_recast_check_time = 0.0
        self.bot_start_time = 0.0
        self.mode_start_time = 0.0
        self.counters = {'recasts': 0, 'skips': 0, 'fish_caught': 0}
        self.last_frames = {'minimap': None, 'hpbar': None}
        self.stop_event = threading.Event()
        self.tray_icon = None
        self.safety_stop_reason = None
        self.is_skipping_this_cast = False # [ИЗМЕНЕНО] Флаг, указывающий на пропуск поклёвки в текущем цикле

    def set_running(self, new_state):
        if new_state and not self.running:
            if self.bot_start_time == 0.0: self.bot_start_time = time.time()
            self.mode_start_time = time.time()
            self.safety_stop_reason = None
        self.running = new_state
        if self.tray_icon:
            self.tray_icon.icon = create_icon('green' if self.running else 'red')
            self.tray_icon.update_menu()

    def reset_fishing_cycle(self):
        self.target_roi = None
        self.bobber_detections = 0
        self.missed_detections = 0
        self.edge_counts.clear()
        self.is_skipping_this_cast = False # [ИЗМЕНЕНО] Сброс флага при новом забросе
        min_t, max_t = max(1, cfg.fishing_timeout[0] - cfg.fishing_timeout[1]), cfg.fishing_timeout[0] + cfg.fishing_timeout[1]
        self.current_fishing_timeout = random.randint(min_t, max_t)
        pyautogui.moveTo(cfg.center_x, cfg.center_y, duration=random.uniform(0.2, 0.4))
        pyautogui.rightClick()
        non_blocking_sleep(random.uniform(0.04, 0.09))
        pyautogui.rightClick()
        self.fishing_timer_start = time.time()
        self.mode_start_time = time.time()
        pyautogui.moveTo(30, 30, duration=random.uniform(1.0, 2.0))

# --- Глобальные экземпляры ---
cfg = Config()
state = State()
model = None

# --- Вспомогательные Функции ---
def non_blocking_sleep(duration):
    end_time = time.time() + duration
    while time.time() < end_time:
        if state.stop_event.is_set(): break
        time.sleep(0.01)

def get_dynamic_padding(box_width, box_height, min_pad=20, max_pad=150):
    size = max(box_width, box_height)
    padding = int(max_pad / (size**0.5 + 1))
    return max(min_pad, min(padding, max_pad))

def compare_pixels(old_frame, new_frame, threshold=10):
    if old_frame is None or new_frame is None or old_frame.shape != new_frame.shape: return 1.0
    diff = cv2.absdiff(old_frame, new_frame)
    gray_diff = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
    change_mask = gray_diff > threshold
    return np.sum(change_mask) / (new_frame.size // new_frame.shape[2])

def create_hud_image(stats, preview_img):
    hud_h, hud_w = 200, 600
    hud_image = np.zeros((hud_h, hud_w, 3), dtype=np.uint8)
    font, font_scale = cv2.FONT_HERSHEY_SIMPLEX, 0.6
    status_color = (0, 255, 0)
    if "PAUSED" in stats['status']: status_color = (0, 0, 255)
    elif "STOP" in stats['status']: status_color = (255, 0, 255)

    text_lines = [
        (f"Status: {stats['status']}", status_color, 2),
        (f"Uptime: {stats['uptime']}", (255, 255, 255), 1),
        (f"Mode: {stats['mode']} ({stats['mode_time']})", (255, 255, 255), 1),
        (f"Edges: {stats['canny_count']} | Avg: {stats['canny_avg']:.1f} | Thr: {stats['canny_thr']:.1f}", (100, 255, 100), 1),
        (f"Recasts: {stats['recasts']} | Skips: {stats['skips']}", (200, 200, 0), 1),
        (f"Fish Caught: {stats['fish_caught']}", (0, 165, 255), 1)
    ]
    for i, (text, color, thickness) in enumerate(text_lines):
        cv2.putText(hud_image, text, (10, 25 + i * 30), font, font_scale, color, thickness)
    if preview_img is not None:
        try:
            preview_size = hud_h - 20
            preview_resized = cv2.resize(preview_img, (preview_size, preview_size))
            if len(preview_resized.shape) == 2: preview_resized = cv2.cvtColor(preview_resized, cv2.COLOR_GRAY2BGR)
            hud_image[10:10+preview_size, hud_w-preview_size-10:hud_w-10] = preview_resized
            cv2.rectangle(hud_image, (hud_w-preview_size-10, 10), (hud_w-10, 10+preview_size), (100, 100, 100), 1)
        except Exception: pass
    return hud_image

# --- Функции выбора зон ---
def zone_selector(target_section='COORDINATES', prefix=''):
    state.is_child_gui_active = True
    screenshot = pyautogui.screenshot()
    image = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
    clone = image.copy()
    ref_point, cropping = [], False
    
    def click_and_crop(event, x, y, flags, param):
        nonlocal ref_point, cropping
        if event == cv2.EVENT_LBUTTONDOWN:
            ref_point, cropping = [(x, y)], True
        elif event == cv2.EVENT_LBUTTONUP:
            ref_point.append((x, y))
            cropping = False

    window_name = "Zone Selector (Press 'q' to quit, 'r' to reset)"
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)
    cv2.setWindowProperty(window_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
    cv2.setMouseCallback(window_name, click_and_crop)

    while not state.stop_event.is_set():
        display_image = clone.copy()
        if len(ref_point) == 2: cv2.rectangle(display_image, ref_point[0], ref_point[1], (0, 255, 0), 2)
        elif cropping and len(ref_point) == 1:
            try: cv2.rectangle(display_image, ref_point[0], pyautogui.position(), (255, 0, 0), 2)
            except Exception: pass
        
        cv2.imshow(window_name, display_image)
        key = cv2.waitKey(1) & 0xFF
        if key == ord("q"): break
        elif key == ord("r"): ref_point, cropping = [], False
        elif len(ref_point) == 2:
            x1, y1 = ref_point[0]; x2, y2 = ref_point[1]
            x, y, w, h = min(x1, x2), min(y1, y2), abs(x1 - x2), abs(y1 - y2)
            if w > 0 and h > 0:
                cfg.save_setting(target_section, f'{prefix}x', x); cfg.save_setting(target_section, f'{prefix}y', y)
                cfg.save_setting(target_section, f'{prefix}w', w); cfg.save_setting(target_section, f'{prefix}h', h)
                cfg.load()
                break
    cv2.destroyAllWindows()
    for _ in range(5): cv2.waitKey(1)
    state.is_child_gui_active = False

# --- Основная логика бота ---
def run_bot():
    global model
    try: model = YOLO(cfg.model_path)
    except Exception as e:
        print(f"Ошибка загрузки модели YOLO: {e}")
        return

    sct = mss.mss()
    hud_window_name = "Bot HUD"
    is_hud_created = False

    while not state.stop_event.is_set():
        if state.is_child_gui_active:
            if is_hud_created:
                try: cv2.destroyWindow(hud_window_name)
                except cv2.error: pass
                is_hud_created = False
                for _ in range(5): cv2.waitKey(1)
            time.sleep(0.5)
            continue
            
        current_time = time.time()
        handle_safety_checks(sct, current_time)

        preview_for_hud, canny_info = None, {'count': 0, 'avg': 0, 'thr': 0}
        if state.running:
            preview_for_hud, canny_info = handle_fishing_logic(sct, current_time)

        if state.show_hud:
            is_hud_created = update_hud(hud_window_name, is_hud_created, current_time, preview_for_hud, canny_info)
        elif is_hud_created:
            try: cv2.destroyWindow(hud_window_name)
            except cv2.error: pass
            is_hud_created = False
            for _ in range(5): cv2.waitKey(1)
        
        time.sleep(0.01)
    cv2.destroyAllWindows()

def handle_safety_checks(sct, current_time):
    if current_time - state.last_safety_check_time < 1.0: return
    
    minimap_box = {'left': cfg.minimap_zone[0], 'top': cfg.minimap_zone[1], 'width': cfg.minimap_zone[2], 'height': cfg.minimap_zone[3]}
    hpbar_box = {'left': cfg.hpbar_zone[0], 'top': cfg.hpbar_zone[1], 'width': cfg.hpbar_zone[2], 'height': cfg.hpbar_zone[3]}
    new_frames = {'minimap': np.array(sct.grab(minimap_box))[:, :, :3], 'hpbar': np.array(sct.grab(hpbar_box))[:, :, :3]}
    
    if state.running:
        if compare_pixels(state.last_frames['minimap'], new_frames['minimap']) > cfg.minimap_pixel_threshold:
            state.safety_stop_reason = "Change minimap!"
            state.set_running(False)
        if compare_pixels(state.last_frames['hpbar'], new_frames['hpbar']) > cfg.hpbar_pixel_threshold:
            state.safety_stop_reason = "Change HP bar!"
            state.set_running(False)
    
    state.last_frames = new_frames
    state.last_safety_check_time = current_time

def handle_fishing_logic(sct, current_time):
    if state.fishing_timer_start == 0:
        state.reset_fishing_cycle()
        return None, {'count': 0, 'avg': 0, 'thr': 0}
    if current_time - state.fishing_timer_start > state.current_fishing_timeout:
        state.reset_fishing_cycle()
        return None, {'count': 0, 'avg': 0, 'thr': 0}
    if current_time - state.last_recast_check_time > cfg.recast_check_interval:
        state.last_recast_check_time = current_time
        min_recast, max_recast = max(0, cfg.recast_chance[0] - cfg.recast_chance[1]), min(100, cfg.recast_chance[0] + cfg.recast_chance[1])
        if random.randint(1, 100) <= random.randint(min_recast, max_recast):
            state.counters['recasts'] += 1
            state.reset_fishing_cycle()
            return None, {'count': 0, 'avg': 0, 'thr': 0}

    roi = state.target_roi if state.target_roi else cfg.scan_zone
    monitor_box = {'left': roi[0], 'top': roi[1], 'width': roi[2], 'height': roi[3]}
    frame = np.array(sct.grab(monitor_box))[:, :, :3]
    canny_info, preview_img = {'count': 0, 'avg': 0, 'thr': 0}, None

    if state.target_roi:
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 70, 150)
        edge_count = np.sum(edges > 0)
        state.edge_counts.append(edge_count)
        if len(state.edge_counts) > cfg.canny_avg_window_size: state.edge_counts.pop(0)
        
        avg_edges = np.mean(state.edge_counts) if state.edge_counts else 0
        threshold = avg_edges * cfg.canny_threshold_multiplier
        canny_info = {'count': edge_count, 'avg': avg_edges, 'thr': threshold}
        preview_img = edges
        
        # [ИЗМЕНЕНО] Проверяем поклёвку только если ещё не решили её пропустить
        if not state.is_skipping_this_cast:
            if edge_count > threshold and avg_edges > 0:
                min_skip, max_skip = max(0, cfg.skip_chance[0] - cfg.skip_chance[1]), min(100, cfg.skip_chance[0] + cfg.skip_chance[1])
                if random.randint(1, 100) > random.randint(min_skip, max_skip):
                    state.counters['fish_caught'] += 1
                    if state.last_bobber_center:
                        pyautogui.moveTo(state.last_bobber_center[0], state.last_bobber_center[1], duration=random.uniform(0.3, 0.5))
                    pyautogui.rightClick()
                    non_blocking_sleep(random.uniform(1.5, 2.5))
                    pyautogui.rightClick()
                    state.reset_fishing_cycle()
                else:
                    state.counters['skips'] += 1
                    state.is_skipping_this_cast = True # [ИЗМЕНЕНО] Устанавливаем флаг, чтобы ждать таймаута
    else:
        results = model(frame, verbose=False)
        preview_img = results[0].plot() if results else frame.copy()
        bobber_detected = False
        for box in results[0].boxes:
            if model.names[int(box.cls[0])] == 'bobber':
                bobber_detected = True
                x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                box_w, box_h = x2 - x1, y2 - y1
                state.last_bobber_center = (cfg.scan_zone[0] + (x1 + x2) // 2, cfg.scan_zone[1] + (y1 + y2) // 2)
                state.bobber_detections += 1
                state.missed_detections = 0
                if state.bobber_detections >= 5:
                    padding = get_dynamic_padding(box_w, box_h)
                    rx, ry = max(0, cfg.scan_zone[0] + x1 - padding), max(0, cfg.scan_zone[1] + y1 - padding)
                    rw, rh = box_w + 2 * padding, box_h + 2 * padding
                    state.target_roi = (rx, ry, rw, rh)
                    state.edge_counts.clear()
                    state.bobber_detections = 0
                    state.mode_start_time = time.time()
                break
        if not bobber_detected:
            state.missed_detections += 1
            if state.missed_detections > cfg.max_missed_detections:
                state.bobber_detections = 0
    return preview_img, canny_info

def update_hud(window_name, is_created, current_time, preview_img, canny_info):
    if not is_created:
        cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE)
        cv2.setWindowProperty(window_name, cv2.WND_PROP_TOPMOST, 1)
        is_created = True

    def format_time(seconds):
        m, s = divmod(seconds, 60)
        h, m = divmod(m, 60)
        return f"{int(h):02d}:{int(m):02d}:{int(s):02d}"

    uptime_str = format_time(current_time - state.bot_start_time if state.bot_start_time > 0 else 0)
    mode_time_str = format_time(current_time - state.mode_start_time if state.mode_start_time > 0 else 0)
    
    status_text = "RUNNING" if state.running else "PAUSED"
    if not state.running and state.safety_stop_reason:
        status_text = f"STOP: {state.safety_stop_reason}"

    hud_stats = {
        'status': status_text, 'uptime': uptime_str,
        'mode': "ROI" if state.target_roi else "YOLO Search", 'mode_time': mode_time_str,
        'canny_count': canny_info['count'], 'canny_avg': canny_info['avg'], 'canny_thr': canny_info['thr'],
        **state.counters
    }
    try:
        cv2.imshow(window_name, create_hud_image(hud_stats, preview_img))
        if cv2.waitKey(1) & 0xFF == ord('q'): on_exit(None, None)
    except cv2.error:
        is_created = False
        state.show_hud = False
        if state.tray_icon: state.tray_icon.update_menu()
    return is_created

# --- Функции GUI и трея ---
def open_settings_window():
    state.is_child_gui_active = True
    root = tk.Tk()
    root.title("Настройки Бота")
    root.resizable(False, False)
    root.attributes('-topmost', True)
    main_frame = ttk.Frame(root, padding="10")
    main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))

    params = {
        'fishing_timeout_base': ('Базовый таймаут (сек)', cfg.fishing_timeout[0]),
        'fishing_timeout_scatter': ('Разброс таймаута (±сек)', cfg.fishing_timeout[1]),
        'canny_threshold_multiplier': ('Множитель порога Canny', cfg.canny_threshold_multiplier),
        'canny_avg_window_size': ('Окно ср. значения Canny', cfg.canny_avg_window_size),
        'max_missed_detections': ('Макс. пропусков детекции', cfg.max_missed_detections),
        'chance_to_skip_catch_base': ('Базовый шанс пропуска (%)', cfg.skip_chance[0]),
        'chance_to_skip_catch_scatter': ('Разброс шанса пропуска (±%)', cfg.skip_chance[1]),
        'chance_to_recast_base': ('Базовый шанс переброса (%)', cfg.recast_chance[0]),
        'chance_to_recast_scatter': ('Разброс шанса переброса (±%)', cfg.recast_chance[1]),
        'recast_check_interval': ('Интервал проверки переброса (сек)', cfg.recast_check_interval),
        'minimap_pixel_threshold': ('Порог изм. миникарты (%)', cfg.minimap_pixel_threshold),
        'hpbar_pixel_threshold': ('Порог изм. HP-бара (%)', cfg.hpbar_pixel_threshold),
    }

    entries = {}
    for i, (key, (desc, val)) in enumerate(params.items()):
        ttk.Label(main_frame, text=f"{desc}:").grid(column=0, row=i, sticky=tk.W, pady=3, padx=5)
        var = tk.StringVar(value=str(val))
        ttk.Entry(main_frame, width=15, textvariable=var).grid(column=1, row=i, sticky=tk.E, padx=5)
        entries[key] = var

    def save_and_close():
        try:
            for key, var in entries.items():
                cfg.save_setting('BOT_CONFIG', key, var.get().replace(',', '.'))
            cfg.load()
            messagebox.showinfo("Успех", "Настройки сохранены.", parent=root)
            on_closing()
        except Exception as e: messagebox.showerror("Ошибка", f"Ошибка сохранения: {e}", parent=root)
    
    def on_closing():
        state.is_child_gui_active = False
        root.destroy()

    root.protocol("WM_DELETE_WINDOW", on_closing)
    button_frame = ttk.Frame(main_frame)
    button_frame.grid(column=0, row=len(params), columnspan=2, pady=15)
    ttk.Button(button_frame, text="Сохранить", command=save_and_close).pack(side=tk.LEFT, padx=10)
    ttk.Button(button_frame, text="Отмена", command=on_closing).pack(side=tk.LEFT, padx=10)
    root.mainloop()

def create_icon(color):
    image = Image.new('RGB', (64, 64), 'black')
    dc = ImageDraw.Draw(image)
    dc.ellipse([(4, 4), (60, 60)], fill=color)
    return image

def start_gui_task(task_func, *args):
    if state.is_child_gui_active: return
    threading.Thread(target=task_func, args=args, daemon=True).start()

def on_toggle_bot(icon, item): state.set_running(not state.running)
def on_toggle_hud(icon, item): 
    state.show_hud = not state.show_hud
    if state.tray_icon: state.tray_icon.update_menu()
def on_exit(icon, item):
    state.stop_event.set()
    if state.tray_icon: state.tray_icon.stop()
    keyboard.unhook_all()

if __name__ == '__main__':
    keyboard.add_hotkey('f5', lambda: state.set_running(True))
    keyboard.add_hotkey('f6', lambda: state.set_running(False))
    
    bot_thread = threading.Thread(target=run_bot, daemon=True)
    
    menu = Menu(
        item(lambda text: 'Пауза (F6)' if state.running else 'Старт (F5)', on_toggle_bot),
        item(lambda text: 'Скрыть HUD' if state.show_hud else 'Показать HUD', on_toggle_hud),
        Menu.SEPARATOR,
        item('Настройка зон', Menu(
            item('Зона сканирования', lambda: start_gui_task(zone_selector, 'COORDINATES')),
            item('Зона миникарты', lambda: start_gui_task(zone_selector, 'SAFETY_ZONES', 'minimap_')),
            item('Зона HP-бара', lambda: start_gui_task(zone_selector, 'SAFETY_ZONES', 'hpbar_')),
        )),
        item('Параметры бота', lambda: start_gui_task(open_settings_window)),
        Menu.SEPARATOR,
        item('Выход', on_exit)
    )
    
    state.tray_icon = icon("FishingBot", create_icon('red'), "Fishing Bot", menu)
    bot_thread.start()
    state.tray_icon.run()
    bot_thread.join(timeout=2)