coordinate/gco.py

# gco.py
import os
import time
import json
import random
import threading
from datetime import datetime, timezone
from typing import List, Dict, Tuple, Any, Optional, Callable
from concurrent.futures import wait

# 导入所有依赖
from vs_types import GroupConfig, QueryWaitMode, VSPlgConfig, VSQueryResult, Task 
from vs_plg_factory import VSPlgFactory 
from toolkit.account_manager import AccountManager 
from toolkit.proxy_manager import ProxyManager 
from toolkit.thread_pool import ThreadPool 
from toolkit.vs_cloud_api import VSCloudApi
import traceback


class GCO:
    """
    @brief GCO 类
    负责管理一个组内的签证插件实例，包括实例的创建、健康检查、
    任务调度、查询和预订流程。
    """
    def __init__(self, cfg: GroupConfig, logger: Callable[[str], None] = None):
        self.m_cfg = cfg
        self.m_factory = VSPlgFactory() # 插件工厂实例
        self.m_logger = logger
        
        self.m_tasks: List[Task] = [] # 存储所有运行中的任务实例
        
        self.m_stop_event = threading.Event() # 用于停止线程
        self.m_lock = threading.RLock() # 保护共享资源 (m_tasks)
        
        self.m_monitor_thread: Optional[threading.Thread] = None
        self.m_creator_thread: Optional[threading.Thread] = None
        
        self.m_pending_builtin = 0  # 正在创建中的内置账号数量
        self.m_pending_orders = 0   # 正在创建中的订单账号数量

    def start(self):
        """
        @brief 启动协调器，包括插件注册和线程启动。
        """
        if not self.m_cfg.enable:
            self._log("Group is disabled, not starting.")
            return

        self._log("Starting coordinator...")
        self.m_stop_event.clear()

        # 注册插件
        plugin_module_path = os.path.join(self.m_cfg.plugin_config.lib_path, f"{self.m_cfg.plugin_config.plugin_bin}")
        
        # === 修复点：更智能的类名推导逻辑 ===
        # 将 snake_case (e.g., "concrete_plugin") 转换为 PascalCase (e.g., "ConcretePlugin")
        plugin_name = self.m_cfg.plugin_config.plugin_name
        class_name = "".join(part.title() for part in plugin_name.split('_'))
        
        # 调试日志：确认推导出的类名
        self._log(f"Inferring class name for plugin {plugin_name}: {class_name}")

        self.m_factory.register_plugin(plugin_name, 
                                       plugin_module_path, 
                                       class_name)

        self.m_monitor_thread = threading.Thread(target=self._monitor_loop, name=f"Monitor-{self.m_cfg.identifier}")
        self.m_creator_thread = threading.Thread(target=self._creator_loop, name=f"Creator-{self.m_cfg.identifier}")
        
        self.m_monitor_thread.start()
        self.m_creator_thread.start()
        self._log("Coordinator threads started.")

    def stop(self):
        """
        @brief 停止协调器，等待所有线程结束。
        """
        self._log("Stopping coordinator...")
        self.m_stop_event.set() # 发送停止信号
        
        if self.m_monitor_thread and self.m_monitor_thread.is_alive():
            self.m_monitor_thread.join()
        if self.m_creator_thread and self.m_creator_thread.is_alive():
            self.m_creator_thread.join()
        self._log("Coordinator stopped.")

    def group_id(self) -> str:
        """
        @brief 获取分组ID。
        """
        return self.m_cfg.identifier
    
    def _log(self, message):
        if self.m_logger:
            self.m_logger(f'[gco] [{self.m_cfg.identifier}] {message}')

    def _monitor_loop(self):
        """
        @brief 监控循环：定期检查实例健康状况，执行查询任务。
               一旦发现号源，立即触发所有实例进行批量预订。
        """
        self._log("[START] monitor loop starting...")
        rng = random.Random()
        
        # 创建一个临时线程池用于并发抢票，避免阻塞监控循环太久
        # max_workers 根据你的最大并发账号数调整，或者设为 None (默认 CPU核心数 * 5)
        
        while not self.m_stop_event.is_set():
            sleep_ms = 0.1 + rng.randint(0, 20) / 1000.0
            time.sleep(sleep_ms)
            
            now = time.time()
            
            # 1. 拷贝任务列表 (Snapshot)
            tasks_to_process = []
            with self.m_lock:
                tasks_to_process = list(self.m_tasks)
            
            # 标记本轮循环是否触发了批量抢票
            batch_booking_triggered = False

            for task in tasks_to_process:
                # 二次检查任务是否还在列表中（防止在遍历过程中被移除）
                with self.m_lock:
                    if task not in self.m_tasks:
                        continue 

                # 健康检查
                if not task.instance.health_check():
                    continue 
                
                # 检查时间窗口
                if now < task.next_run:
                    continue
                
                # === 执行查询 ===
                apt_type = None
                try:
                    apt_types = self.m_cfg.appointment_types
                    if not apt_types:
                        self._log(f"No matching appointment configuration found.")
                        continue
                    
                    weights = [float(item.weight) for item in apt_types]
                    apt_type = random.choices(apt_types, weights=weights, k=1)[0]
                    
                    VSCloudApi.Instance().slot_refresh_start(
                        apt_type.routing_key,
                        country=apt_type.country,
                        city=apt_type.city,
                        visa_type=apt_type.visa_type
                    )
                    # 这里的 task 充当了“哨兵”的角色
                    result = task.instance.query(apt_type)
                    result.apt_type = apt_type 
                    VSCloudApi.Instance().slot_refresh_success(
                        apt_type.routing_key
                    )
                    if result.success:
                        self._log(f"🔥 Slot Found by [{task.instance.get_group_id()}]! Triggering BATCH BOOKING for {len(tasks_to_process)} workers.")
                       
                        # 上报Slot Snapshot
                        # query_payload = result.to_snapshot_payload()
                        # query_payload["website"] = self.m_cfg.website
                        # query_payload["snapshot_source"] = 'worker'
                        # query_payload["snapshot_at"] = datetime.now(timezone.utc).isoformat()
                        # VSCloudApi.Instance().slot_snapshot_report(query_payload)
                        
                        # === [核心修改]：一人发现，全员出击 ===
                        # 1. 准备并发任务
                        # 我们使用刚刚快照的 tasks_to_process，或者重新获取一次全量列表
                        # 重点：所有实例 (w.instance) 都使用同一份查询结果 (result) 去抢
                        futures = []
                        for worker in tasks_to_process:
                            # 提交到线程池并发执行，极大减少时间差
                            f = ThreadPool.getInstance().enqueue(
                                self._on_query_result, 
                                worker, 
                                result
                            )
                            futures.append(f)
                        
                        # 2. 等待所有抢票任务结束 (可选，根据业务需求决定是否阻塞监控)
                        # 如果不wait，监控线程会立刻继续，可能导致重复触发
                        # 建议 wait，确保这一波抢票彻底结束
                        for f in futures:
                            try:
                                f.result() # 获取结果，捕获异常
                            except Exception as e:
                                self._log(f"Batch booking exception: {e}")

                        # 3. 标记触发状态
                        batch_booking_triggered = True
                        break 
                    
                    else:
                        # 没查到，仅当前 task 记录日志
                        self._log(f"Query done by {task.instance.get_group_id()}, No availability")

                except Exception as e:
                    traceback.print_exc()
                    self._log(f"Exception during query: {e}")
                    if apt_type:
                        VSCloudApi.Instance().slot_refresh_fail(
                            apt_type.routing_key,
                            error=str(e)
                        )

                # === 计算下次运行时间 (仅针对当前 query 的 task，除非触发了 batch) ===
                if not batch_booking_triggered:
                    interval = 30
                    mode = task.qw_cfg.mode
                    if mode == QueryWaitMode.Loop:
                        interval = 0
                    elif mode == QueryWaitMode.Fixed:
                        interval = task.qw_cfg.fixed_wait
                    elif mode == QueryWaitMode.Random:
                        interval = rng.randint(task.qw_cfg.random_min, task.qw_cfg.random_max)
                    
                    task.next_run = time.time() + interval

            # === [修正后的批量冷却逻辑] ===
            # 抢票结束后，必须严格按照每个账号配置的频率进入冷却，防止集体429
            if batch_booking_triggered:
                self._log(f"Batch booking finished. Resetting wait times based on configurations.")
                
                # 重新获取当前时间（因为抢票过程消耗了时间）
                now_ts = time.time()
                
                with self.m_lock:
                    for t in self.m_tasks:
                        # 重新读取该任务的配置
                        interval = 30 # 默认兜底
                        mode = t.qw_cfg.mode
                        
                        if mode == QueryWaitMode.Loop:
                            # 即使是 Loop 模式，在大规模抢票后建议给一个微小的缓冲（如1秒），避免死循环导致 CPU 飙升
                            # 如果你的逻辑允许立刻重试，这里可以是 0
                            interval = 1 
                        elif mode == QueryWaitMode.Fixed:
                            interval = t.qw_cfg.fixed_wait
                        elif mode == QueryWaitMode.Random:
                            interval = rng.randint(t.qw_cfg.random_min, t.qw_cfg.random_max)
                        
                        # 设置下次运行时间
                        t.next_run = now_ts + interval
                
                self._log(f"All workers cooldown reset. Resuming monitor loop.")

            # 清理不健康实例
            with self.m_lock:
                initial_size = len(self.m_tasks)
                self.m_tasks[:] = [t for t in self.m_tasks if t.instance.health_check()]
                if len(self.m_tasks) < initial_size:
                    self._log(f"Removed {initial_size - len(self.m_tasks)} unhealthy instance(s). Remaining: {len(self.m_tasks)}")

        self._log("[STOP] monitor loop exiting...")

    def _creator_loop(self):
        """
        @brief 创建者循环：双轨制并发控制 + Pending 计数防超发
        """
        self._log("[START] creator loop starting...")
        
        while not self.m_stop_event.is_set():
            time.sleep(1.0) # 保持 1秒间隔，给资源管理器缓冲时间

            # -------------------------------------------------------------
            # 1. 统计当前状态 (Active + Pending)
            # -------------------------------------------------------------
            current_builtin = 0
            current_orders = 0
            
            with self.m_lock:
                for t in self.m_tasks:
                    if t.task_ref:
                        current_orders += 1
                    else:
                        current_builtin += 1
                
                # 读取正在创建中的计数 (快照)
                pending_builtin = self.m_pending_builtin
                pending_orders = self.m_pending_orders

            # -------------------------------------------------------------
            # 2. 计算缺口
            # -------------------------------------------------------------
            
            # === [内置账号缺口] ===
            needed_builtin = 0
            # 逻辑：只要目标数量 > 0，就尝试补充 (兼容需要账号和不需要账号的模式)
            if self.m_cfg.target_instances > 0:
                total_builtin_proj = current_builtin + pending_builtin
                needed_builtin = self.m_cfg.target_instances - total_builtin_proj

            # === [订单账号缺口] ===
            needed_order = 0
            # 逻辑：必须需要账号(need_account=True) 且 限制数量大于 0
            # (order_account_enable 已移除，直接看 limit)
            if self.m_cfg.need_account and self.m_cfg.order_account_online_limit > 0:
                total_order_proj = current_orders + pending_orders
                needed_order = self.m_cfg.order_account_online_limit - total_order_proj

            # 如果两边都满了，跳过本轮
            if needed_builtin <= 0 and needed_order <= 0:
                continue

            # -------------------------------------------------------------
            # 3. 准备配置
            # -------------------------------------------------------------
            config_data = self._prepare_next_config(
                need_builtin=(needed_builtin > 0),
                need_order=(needed_order > 0),
            )

            if not config_data:
                time.sleep(10.0)
                continue

            plg_cfg, task_ref = config_data
            
            # -------------------------------------------------------------
            # 4. [绝对核心] 提交前立即增加 Pending 计数
            # -------------------------------------------------------------
            # 必须在这里加！防止下一秒循环重复创建！
            with self.m_lock:
                if task_ref:
                    self.m_pending_orders += 1
                else:
                    self.m_pending_builtin += 1
            
            p_type = "Order" if task_ref else "Built-in"
            self._log(f"+++ Spawning {p_type} (Pending: {self.m_pending_builtin}/{self.m_pending_orders})...")

            # 5. 异步提交
            try:
                ThreadPool.getInstance().enqueue(
                    self._create_and_register_plg_worker, 
                    plg_cfg, 
                    task_ref
                )
            except Exception as e:
                # 提交失败回滚计数
                self._log(f"Failed to enqueue task: {e}")
                with self.m_lock:
                    if task_ref: self.m_pending_orders -= 1
                    else: self.m_pending_builtin -= 1
            
            # 错开并发
            time.sleep(random.uniform(0.5, 1.0))

        self._log("[STOP] creator loop exiting...")

    def _prepare_next_config(self, need_builtin: bool, need_order: bool) -> Optional[Tuple[VSPlgConfig, Optional[Dict[str, Any]]]]:
        """
        @brief 准备下一个插件实例的配置
        """
        plg_cfg = VSPlgConfig()
        plg_cfg.debug = self.m_cfg.debug
        plg_cfg.free_config = self.m_cfg.free_config
        plg_cfg.session_max_life = self.m_cfg.session_max_life
        
        task_ref = None 
        config_ready = False
        pool_name = self.m_cfg.local_account_pool 

        # =================================================================
        # 1. 账号获取
        # =================================================================
        if not self.m_cfg.need_account:
            # === 游客模式 (无需账号) ===
            if need_builtin:
                plg_cfg.account.id = 0
                plg_cfg.account.username = "Guest"
                config_ready = True
                task_ref = None
        else:
            # === 标准模式 (需要账号) ===
            
            # A. 优先补充内置账号 (只要 target_instances 还有缺口)
            if need_builtin:
                # 获取并锁定账号
                account = AccountManager.Instance().next(
                    pool_name, 
                    lock_duration=self.m_cfg.account_login_interval * 60
                )
                if account:
                    plg_cfg.account.id = account["id"]
                    plg_cfg.account.username = account["username"]
                    plg_cfg.account.password = account["password"]
                    plg_cfg.account.lock_until = account.get("lock_until", 0)
                    config_ready = True
                    task_ref = None
                    self._log(f"Selected Built-in: {plg_cfg.account.username}")

            # B. 次选补充订单账号 (如果内置不需要 或 池子空了)
            # 只有当 limit > 0 时才尝试
            if not config_ready and need_order and self.m_cfg.order_account_online_limit > 0:
                try:
                    routing_key = self.m_cfg.order_account_routing
                    task_ref = VSCloudApi.Instance().get_vas_task_pop(routing_key)
                    
                    if task_ref:
                        user_inputs = task_ref.get('user_inputs', {})
                        plg_cfg.account.id = 0 # 临时账号
                        plg_cfg.account.username = user_inputs.get(self.m_cfg.input_map_username, "")
                        plg_cfg.account.password = user_inputs.get(self.m_cfg.input_map_password, "")
                        
                        if plg_cfg.account.username:
                            config_ready = True
                            self._log(f"Selected Order Acc: {plg_cfg.account.username}")
                        else:
                            VSCloudApi.Instance().return_vas_task_to_queue(task_ref['id'])
                            return None
                except Exception as e:
                    pass
                    #self._log(f"Get Order task exception, e={e}")

        if not config_ready:
            return None

        # =================================================================
        # 2. 代理配置
        # =================================================================
        if self.m_cfg.need_proxy:
            # 轮询代理
            proxy_lock_time = self.m_cfg.proxy_lock_interval
            proxy = ProxyManager.Instance().next(self.m_cfg.proxy_pool, lock_duration=proxy_lock_time)
            if not proxy:
                try:
                    if task_ref:
                        VSCloudApi.Instance().return_vas_task_to_queue(task_ref['id'])
                    return None
                except Exception as e:
                    self._log(f"Return Order task to queue exception, e={e}")

            plg_cfg.proxy.id = proxy["id"]
            plg_cfg.proxy.ip = proxy["ip"]
            plg_cfg.proxy.port = proxy["port"]
            plg_cfg.proxy.scheme = proxy["scheme"]
            plg_cfg.proxy.username = proxy.get("username", "")
            plg_cfg.proxy.password = proxy.get("password", "")
            plg_cfg.proxy.lock_until = proxy.get("lock_until", 0)

        return plg_cfg, task_ref

    def _create_and_register_plg_worker(self, plg_cfg: VSPlgConfig, task_ref: Optional[Dict[str, Any]] = None):
        """
        @brief 异步创建工作线程
        """
        instance = None
        creation_success = False 
        
        try:
            # 1. 耗时操作：实例化 & 登录
            instance = self.m_factory.create(self.m_cfg.identifier, self.m_cfg.plugin_config.plugin_name)
            instance.set_log(self.m_logger)
            instance.set_config(plg_cfg)
            
            instance.create_session() # 可能耗时很久

            # 2. 注册到任务列表
            with self.m_lock:
                # 既然允许短暂超发，这里直接添加，不做严格的拒绝并返回逻辑
                
                # 计算预定权限
                book_allowed = False
                if not self.m_cfg.account_bind_applicant:
                    book_allowed = True
                elif self.m_cfg.account_bind_applicant and task_ref:
                    book_allowed = True
                
                new_task = Task(
                    instance=instance,
                    qw_cfg=self.m_cfg.query_wait,
                    next_run=time.time(),
                    task_ref=task_ref,
                    book_allowed=book_allowed
                )
                self.m_tasks.append(new_task)
                
                creation_success = True
                
                p_type = "Order" if task_ref else "Built-in"
                self._log(f"=== Instance Registered [{p_type}]. Total Active: {len(self.m_tasks)} ===")

        except Exception as e:
            self._log(f"Creation failed for {plg_cfg.account.username}: {e}")
            
        finally:
            # -------------------------------------------------------------
            # 3. [绝对核心] 必须扣减 Pending
            # -------------------------------------------------------------
            # 无论上面是否抛出异常，或者是否注册成功，Pending 必须释放
            with self.m_lock:
                if task_ref:
                    self.m_pending_orders -= 1
                    if self.m_pending_orders < 0: self.m_pending_orders = 0
                else:
                    self.m_pending_builtin -= 1
                    if self.m_pending_builtin < 0: self.m_pending_builtin = 0
            
            # 4. 失败回滚
            if not creation_success:
                # 如果是云端订单，必须归还，否则丢单
                if task_ref:
                    self._log(f"Rolling back task {task_ref['id']}")
                    try:
                        VSCloudApi.Instance().return_vas_task_to_queue(task_ref['id'])
                    except:
                        pass

    def _on_query_result(self, t: Task, query_result: VSQueryResult):
        self._log(f"Query result received: {str(query_result)}. BLOCKING monitor loop for booking...")
        
        if not t.book_allowed:
            return

        # -------------------------------------------------------
        # 1. 准备任务数据 (Data Preparation)
        # -------------------------------------------------------
        task_data = t.task_ref
        is_cloud_task = False # 标记：是否为云端临时取出的任务（需要回滚）

        # 如果没有绑定本地任务，尝试从云端 Pop
        if not task_data:
            apt_type = query_result.apt_type
            booking_routing_key = f'auto.{apt_type.routing_key}' if apt_type.routing_key else "default"
            task_data = VSCloudApi.Instance().get_vas_task_pop(booking_routing_key)
            
            if not task_data:
                # 这种情况属于：内置账号查到了号，但云端没有待处理订单，直接放弃
                self._log(f"No pending task found for key {booking_routing_key}. Abandoning slot.")
                return
            
            is_cloud_task = True
            self._log(f"Picked up Cloud Task ID {task_data['id']} for booking...")

        # 统一提取核心参数
        task_id = task_data['id']
        order_id = task_data.get('order_id')
        user_input = task_data.get('user_inputs', {})

        # -------------------------------------------------------
        # 2. 执行预订 (Execution)
        # -------------------------------------------------------
        booking_success = False
        
        try:
            # 统一调用，无需区分来源
            book_res = t.instance.book(query_result, user_input)
            
            # -------------------------------------------------------
            # 3. 成功处理 (Success Handling)
            # -------------------------------------------------------
            if book_res.success:
                booking_success = True
                self._log(f"✅ Booking SUCCESS! Order: {order_id}")
                
                current_grab_info = {
                    "account": book_res.account,
                    "session_id": book_res.session_id,
                    "urn": book_res.urn,
                    "slot_date": book_res.book_date,
                    "slot_time": book_res.book_time,
                    "timestamp": int(time.time()),
                    "payment_link": book_res.payment_link,
                }
                
                update_data = {
                    "status": "grabbed",
                    "grabbed_history": current_grab_info
                }
                
                VSCloudApi.Instance().update_vas_task(task_id, update_data)
                self._log(f"Task {task_id} marked as GRABBED.")
            
            else:
                self._log(f"❌ Booking Failed for Order {order_id}: {book_res.message}")

        except Exception as e:
            self._log(f"Exception during booking for Order {order_id}: {e}")

        finally:
            # -------------------------------------------------------
            # 4. 回滚机制 (Rollback / Return to Queue)
            # -------------------------------------------------------
            # 只有两个条件同时满足才回滚：
            # 1. 任务来自云端队列 (is_cloud_task is True)
            # 2. 预定没有成功 (booking_success is False) - 包括预定失败或发生异常
            if is_cloud_task and not booking_success:
                self._log(f"Returning Task {task_id} to queue (status=pending).")
                try:
                    VSCloudApi.Instance().return_vas_task_to_queue(task_id)
                except Exception as ex:
                    self._log(f"Failed to return task to queue: {ex}")