coordinate/booker_builtin.py

import os
import time
import json
import threading
import random
import redis
from typing import List, Dict, Callable

import configure
from vs_types import GroupConfig, VSPlgConfig, Task, VSQueryResult, AppointmentType
from vs_plg_factory import VSPlgFactory 
from toolkit.thread_pool import ThreadPool 
from toolkit.vs_cloud_api import VSCloudApi
from toolkit.backoff import ExponentialBackoff

class BuiltinBookerGCO:
    """
    非绑定模式 (公共内置账号池):
    - 只维护全局 target_instances 数量的实例。
    - 所有实例热机等待，发现信号后临时去云端 Pop 订单。
    """
    def __init__(self, cfg: GroupConfig, redis_conf: Dict, logger: Callable[[str], None] = None):
        self.m_cfg = cfg
        self.m_factory = VSPlgFactory()
        self.m_logger = logger
        self.m_tasks: List[Task] = []
        self.m_lock = threading.RLock()
        self.m_stop_event = threading.Event()
        self.redis_com = redis.Redis(**redis_conf)
        self.redis_sub = redis.Redis(**redis_conf)
        self.m_pending_builtin = 0
        
        self.m_tracker_key = f"vs:worker:tasks_tracker:{self.m_cfg.identifier}"
        self.group_backoff = ExponentialBackoff(base_delay=60.0, max_delay=10*60.0, factor=2.0)
        self.task_backoff = ExponentialBackoff(base_delay=5*60.0, max_delay=2*60*60.0, factor=2.0)
        self.m_last_spawn_time = 0.0
        self.heartbeat_ttl = 300

    def _log(self, message):
        if self.m_logger:
            self.m_logger(f'[BUILTIN-BOOKER] [{self.m_cfg.identifier}] {message}')

    def start(self):
        if not self.m_cfg.enable:
            return
        self._log("Starting Built-in Booker...")
        plugin_name = self.m_cfg.plugin_config.plugin_name
        class_name = "".join(part.title() for part in plugin_name.split('_'))
        plugin_path = os.path.join(self.m_cfg.plugin_config.lib_path, self.m_cfg.plugin_config.plugin_bin)
        self.m_factory.register_plugin(plugin_name, plugin_path, class_name)

        threading.Thread(target=self._booking_trigger_loop, daemon=True).start()
        threading.Thread(target=self._creator_loop, daemon=True).start()
        threading.Thread(target=self._maintain_loop, daemon=True).start()

    def stop(self):
        self._log("Stopping Booker...")
        self.m_stop_event.set()
        self._cleanup_all_tasks("booker stop")

    def _cleanup_task(self, task: Task, reason: str = ""):
        try:
            instance = getattr(task, 'instance', None)
            if instance and hasattr(instance, 'cleanup'):
                instance.cleanup()
                self._log(f"🧹 Cleaned up built-in instance. Reason: {reason}")
        except Exception as e:
            self._log(f"Cleanup failed for built-in instance. Reason: {reason}. Error: {e}")

    def _remove_task(self, task: Task, reason: str = "", cleanup: bool = True):
        removed = False
        with self.m_lock:
            if task in self.m_tasks:
                self.m_tasks.remove(task)
                removed = True
        if cleanup and removed:
            self._cleanup_task(task, reason)
        return removed

    def _cleanup_all_tasks(self, reason: str = ""):
        with self.m_lock:
            tasks = list(self.m_tasks)
            self.m_tasks.clear()
        for task in tasks:
            self._cleanup_task(task, reason)

    def _get_redis_key(self, routing_key: str) -> str:
        return f"vs:signal:{routing_key}"

    def _maintain_loop(self):
        self._log("Maintain loop started.")
        while not self.m_stop_event.is_set():
            time.sleep(1.0)
            now = time.time()
            
            with self.m_lock:
                tasks_to_check = list(self.m_tasks)
            
            if not tasks_to_check:
                continue
            
            healthy_tasks = []
            dead_tasks = []
            for t in tasks_to_check:
                if now >= t.next_remote_ping:
                    try:
                        t.instance.keep_alive()
                        if t.instance.health_check():
                            healthy_tasks.append(t)
                            next_delay = random.randint(180, 300) 
                            t.next_remote_ping = now + next_delay
                        else:
                            dead_tasks.append(t)
                            self._log(f"♻️ Instance unhealthy. Will be removed.")
                    except Exception as e:
                        dead_tasks.append(t)
                        self._log(f"Instance keep-alive failed: {e}")
                else:
                    healthy_tasks.append(t)
            
            if dead_tasks:
                with self.m_lock:
                    current_tasks = list(self.m_tasks)
                    self.m_tasks = [t for t in self.m_tasks if t in healthy_tasks]
                for t in dead_tasks:
                    if t in current_tasks:
                        self._cleanup_task(t, "unhealthy or keep-alive failed")
            else:
                with self.m_lock:
                    self.m_tasks = [t for t in self.m_tasks if t in healthy_tasks]

    def _booking_trigger_loop(self):
        self._log("Pub/Sub Trigger loop started.")
        channel_to_routing_key = {}
        for apt in self.m_cfg.appointment_types:
            channel = self._get_redis_key(apt.routing_key)
            channel_to_routing_key[channel] = apt.routing_key
            
        if not channel_to_routing_key:
            self._log("No appointment types configured. Exiting trigger loop.")
            return
            
        pubsub = None
        while not self.m_stop_event.is_set():
            try:
                if pubsub is None:
                    pubsub = self.redis_sub.pubsub(ignore_subscribe_messages=False)
                    channels_to_sub = list(channel_to_routing_key.keys())
                    self._log(f"⏳ Sending SUBSCRIBE command to Redis for: {channels_to_sub}")
                    pubsub.subscribe(*channels_to_sub)
                message = pubsub.get_message(timeout=5.0)
                if not message:
                    continue
                channel = message['channel']
                if isinstance(channel, bytes):
                    channel = channel.decode('utf-8')
                if message['type'] == 'subscribe':
                    active_subs = message['data']
                    self._log(f"📡 [Redis ACK] Successfully subscribed to: {channel} (Active connection subs: {active_subs})")
                    continue
                if message['type'] != 'message':
                    continue
                raw_data = message['data']
                if isinstance(raw_data, bytes):
                    raw_data = raw_data.decode('utf-8')
                routing_key = channel_to_routing_key.get(channel)
                if not routing_key:
                    continue
                try:
                    data = json.loads(raw_data)
                    query_result = VSQueryResult.model_validate(data['query_result'])
                    query_result.apt_type = AppointmentType.model_validate(data['apt_type'])
                except Exception as parse_err:
                    self._log(f"Data parsing error for channel {channel}: {parse_err}")
                    continue
                now = time.time()
                matching_tasks = []
                with self.m_lock:
                    for task in self.m_tasks:
                        if now < task.next_run or not task.book_allowed:
                            continue
                        if routing_key not in task.acceptable_routing_keys:
                            continue   
                        task.next_run = now + self.m_cfg.booker.booking_cooldown
                        matching_tasks.append(task)
                if matching_tasks:
                    for task in matching_tasks:
                        self._log(f"🚀 Triggering BOOK for {routing_key} | Order Ref: {task.task_ref}")
                        t = threading.Thread(
                            target=self._execute_book_job, 
                            args=(task, query_result), 
                            daemon=True
                        )
                        t.start()
            except Exception as e:
                self._log(f"Trigger loop pub/sub error: {e}")
                if pubsub:
                    try:
                        pubsub.close()
                    except:
                        pass
                    pubsub = None
                time.sleep(2)
        if pubsub:
            pubsub.close()
            self._log("Pub/Sub connection closed.")

    def _execute_book_job(self, task: Task, query_result: VSQueryResult):
        queue_name = f"auto.{query_result.apt_type.routing_key}"

        task_id = None
        task_data = None
        booking_success = False
        is_rate_limited = False
        
        try:
            task_data = VSCloudApi.Instance().get_vas_task_pop(queue_name)
            if not task_data:
                return 
            task_id = task_data['id']
            order_id = task_data.get('order_id')
            self.redis_com.zadd(self.m_tracker_key, {str(task_id): time.time() + self.heartbeat_ttl})
                        
            user_input = task_data.get('user_inputs', {})
            book_res = task.instance.book(query_result, user_input)
            
            if book_res.success:
                booking_success = True
                self._log(f"✅ BOOK SUCCESS! Order: {order_id}")
                
                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
                }
                VSCloudApi.Instance().update_vas_task(task_id, {"status": "grabbed", "grabbed_history": grab_info})
                push_content = (
                    f"🎉 【预定成功通知】\n"
                    f"━━━━━━━━━━━━━━━\n"
                    f"订单编号: {order_id}\n"
                    f"预约账号: {book_res.account}\n"
                    f"预约日期: {book_res.book_date}\n"
                    f"预约时间: {book_res.book_time}\n"
                    f"预约编号: {book_res.urn}\n"
                    f"支付链接: {book_res.payment_link if book_res.payment_link else '无需支付/暂无'}\n"
                    f"━━━━━━━━━━━━━━━\n"
                )
                VSCloudApi.Instance().push_weixin_text(push_content)
                self.redis_com.zrem(self.m_tracker_key, task_id)
                
                # === 核心：成功次数判断 ===
                task.successful_bookings += 1
                max_b = self.m_cfg.booker.max_bookings_per_account
                if max_b > 0 and task.successful_bookings >= max_b:
                    self._log(f"Account reached max bookings ({max_b}). Destroying instance.")
                    self._remove_task(task, "max bookings reached")
            else:
                self._log(f"❌ BOOK FAILED for Order: {order_id}")

        except Exception as e:
            err_str = str(e)
            self._log(f"Exception during booking: {err_str}")
            rate_limited_indicators = [
                "42901" in err_str,
                "Rate limited" in err_str
            ]
            if any(rate_limited_indicators):
                is_rate_limited = True
                self._remove_task(task, "booking rate limited")
                if task_data and task_id is not None:
                    task_meta = task_data.get('meta') or {} 
                    t_fails = task_meta.get('booking_failures', 0) + 1
                    task_meta['booking_failures'] = t_fails
                    
                    try:
                        VSCloudApi.Instance().update_vas_task(task_id, {"meta": task_meta})
                    except Exception as cloud_err:
                        self._log(f"Failed to update task meta: {cloud_err}")
                        
                    t_cd = self.task_backoff.calculate(t_fails)
                    self._log(f"⏳ Task={task_id} (Booking Attempt {t_fails}) suspended for {t_cd:.1f}s.")
                    self.redis_com.zadd(self.m_tracker_key, {str(task_id): time.time() + t_cd})
            
        finally:
           if not booking_success and task_id is not None and not is_rate_limited:
                self.redis_com.zadd(self.m_tracker_key, {str(task_id): 0})
                self._log(f"♻️ Task={task_id} normal failure. Instantly handed over to Sweeper.")
                
    def _creator_loop(self):
        self._log("Creator loop started.")
        spawn_interval = 10.0
        group_cd_key = f"vs:group:cooldown:{self.m_cfg.identifier}"
        while not self.m_stop_event.is_set():
            time.sleep(2.0)
            if self.redis_com.exists(group_cd_key):
                continue
            with self.m_lock:
                current = len(self.m_tasks)
                pending = self.m_pending_builtin
                target = self.m_cfg.booker.target_instances
            if (current + pending) < target:
                now = time.time()
                if now - self.m_last_spawn_time >= spawn_interval:
                    self.m_last_spawn_time = now 
                    self._spawn_worker()

    def _spawn_worker(self):
        with self.m_lock:
            self.m_pending_builtin += 1
            
        def _job():
            try:
                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

                if self.m_cfg.need_account:
                    acc = VSCloudApi.Instance().get_next_account(self.m_cfg.booker.account_pool_id, self.m_cfg.booker.account_cd)
                    plg_cfg.account.id = acc['id']
                    plg_cfg.account.username = acc['username']
                    plg_cfg.account.password = acc['password']

                if self.m_cfg.need_proxy:
                    proxy = VSCloudApi.Instance().get_next_proxy(self.m_cfg.proxy_pool, self.m_cfg.proxy_cd)
                    plg_cfg.proxy.id = proxy['id']
                    plg_cfg.proxy.ip = proxy['ip']
                    plg_cfg.proxy.port = proxy['port']
                    plg_cfg.proxy.proto = proxy['proto']
                    plg_cfg.proxy.username = proxy['username']
                    plg_cfg.proxy.password = proxy['password']

                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()
                
                with self.m_lock:
                    all_keys = [apt.routing_key for apt in self.m_cfg.appointment_types]
                    self.m_tasks.append(
                        Task(
                            instance=instance,
                            qw_cfg=self.m_cfg.query_wait,
                            next_run=time.time(), 
                            task_ref=None,
                            acceptable_routing_keys=all_keys,
                            source_queue="built-in",
                            book_allowed=True,
                            next_remote_ping = time.time() + random.randint(180, 300) 
                        )
                    )
                    
                    group_fail_key = f"vs:group:failures:{self.m_cfg.identifier}"
                    self.redis_com.delete(group_fail_key)
                    
                self._log(f"+++ Built-in Booker spawned: {plg_cfg.account.username}")
            except Exception as e:
                err_str = str(e)
                resource_not_found_indicators = [
                    "40401" in err_str,
                    "Account not found" in err_str,
                    "Proxy not found" in err_str,
                ]
                if any(resource_not_found_indicators):
                    return
                
                self._log(f"Spawn failed: {e}")
                
                rate_limited_indicators = [
                    "42901" in err_str,
                    "Rate limited" in err_str
                ]
                if any(rate_limited_indicators):
                    group_fail_key = f"vs:group:failures:{self.m_cfg.identifier}"
                    group_cd_key = f"vs:group:cooldown:{self.m_cfg.identifier}"
                    
                    # 更新全局（机器组）失败次数
                    g_fails = self.redis_com.incr(group_fail_key)
                    # 计算退避时间
                    g_cd = self.group_backoff.calculate(g_fails)
                    # 设置 Redis 全局冷却保护阀
                    self.redis_com.set(group_cd_key, "1", ex=int(g_cd))
                    self._log(f"📉 [Rate Limited] Group '{self.m_cfg.identifier}' failed {g_fails} times. Global Backoff: {g_cd:.1f}s.")

            finally:
                with self.m_lock:
                    self.m_pending_builtin = max(0, self.m_pending_builtin - 1)
        ThreadPool.getInstance().enqueue(_job)