import subprocess
import signal

class WorkerManager:
    def __init__(self, shutdown_timeout=5):
        self.process = None
        self.running = False
        self.shutdown_timeout = shutdown_timeout
    
    def start_worker(self):
        """Запуск воркера"""
        cmd = [
            "taskiq", "worker", 
            "app.taskiq.broker:broker",  
            "--workers", "1",         
            "--max-prefetch", "1",    
        ]
        
        print("Запуск воркера...")
        self.process = subprocess.Popen(cmd)
        self.running = True
        
        # Установить обработчики сигналов
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)
        
        try:
            return_code = self.process.wait()
            print(f"Воркер завершился с кодом: {return_code}")
        except KeyboardInterrupt:
            print("Получен SIGINT, останавливаем...")
            self.stop_worker()
        except Exception as e:
            print(f"Ошибка при ожидании воркера: {e}")
            self.stop_worker()
    
    def _signal_handler(self, sig, frame):
        """Обработчик сигнала остановки"""
        print(f"\nПолучен сигнал {sig}, останавливаю воркер...")
        self.stop_worker()
    
    def stop_worker(self):
        """Корректная остановка воркера"""
        if self.process and self.running:
            print("Отправляем SIGTERM воркеру...")
            self.process.terminate()
            
            try:
                # Ждать завершения в течение self.shutdown_timeout секунд
                return_code = self.process.wait(timeout=self.shutdown_timeout)
                print(f"Воркер завершился с кодом: {return_code}")
            except subprocess.TimeoutExpired:
                print(f"Воркер не завершился за {self.shutdown_timeout} секунд, отправляем SIGKILL...")
                self.process.kill()
                try:
                    self.process.wait(timeout=2)  # 2 секунды на SIGKILL
                    print("Процесс принудительно завершен")
                except subprocess.TimeoutExpired:
                    print("Не удалось завершить процесс принудительно")
            
            self.running = False
            print("Воркер остановлен")

def main():
    # Использовать короткий таймаут
    manager = WorkerManager(shutdown_timeout=3)
    manager.start_worker()

if __name__ == "__main__":
    main()