coordinate/plugins/usa_plugin.py

import time
import json
import random
import re
import os
import uuid
import shutil
import socket
from datetime import date, datetime
from typing import List, Dict, Optional, Any, Callable
from urllib.parse import urljoin, urlparse, urlencode, parse_qs
from concurrent.futures import ThreadPoolExecutor
from DrissionPage import ChromiumPage, ChromiumOptions

import configure
from vs_plg import IVSPlg
from utils.cloudflare_bypass_for_scraping import CloudflareBypasser
from toolkit.mihomo_tunnel import MihomoTunnel
from utils.mouse import HumanMouse
from utils.keyboard import HumanKeyboard
from utils.fingerprint_utils import FingerprintGenerator
from vs_types import VSPlgConfig, AppointmentType, VSQueryResult, VSBookResult, AvailabilityStatus, TimeSlot, DateAvailability, NotFoundError, PermissionDeniedError, RateLimiteddError, SessionExpiredOrInvalidError, BizLogicError 


class BrowserResponse:
    """模拟 requests.Response"""
    def __init__(self, result_dict):
        result_dict = result_dict or {}
        self.status_code = result_dict.get('status', 0)
        self.text = result_dict.get('body', '')
        self.headers = result_dict.get('headers', {})
        self.url = result_dict.get('url', '')
        self._json = None

    def json(self):
        if self._json is None:
            if not self.text:
                return {}
            try:
                self._json = json.loads(self.text)
            except:
                self._json = {}
        return self._json


class UsaPlugin(IVSPlg):
    
    LOCATIONS = {
        "SHANGHAI": {"name": "SHANGHAI", "id": "096bf614-b0db-ec11-a7b4-001dd80234f6"},
        "WUHAN": {"name": "WUHAN", "id": "7b6af614-b0db-ec11-a7b4-001dd80234f6"},
        "SHENYANG": {"name": "SHENYANG", "id": "0f6bf614-b0db-ec11-a7b4-001dd80234f6"},
    }

    def __init__(self, group_id: str):
        self.group_id = group_id
        self.config: Optional[VSPlgConfig] = None
        self.free_config: Dict[str, Any] = {}
        self.is_healthy = True
        self.logger = None
        
        self.mouse = None
        self.keyboard = None
        self.page: Optional[ChromiumPage] = None
        
        self.instance_id = uuid.uuid4().hex[:8]
        self.root_workspace = os.path.abspath(os.path.join("data/temp_browser_data", f"{self.group_id}.{self.instance_id}"))
        self.user_data_path = os.path.join(self.root_workspace, "user_data")
        
        if not os.path.exists(self.root_workspace):
            os.makedirs(self.root_workspace)     

        self.tunnel = None
        self.session_create_time: float = 0

    def _log(self, message):
        if self.logger:
            self.logger(f'[UsaPlugin] [{self.group_id}] [{self.instance_id}] {message}')
        else:
            print(f'[UsaPlugin] [{self.group_id}] [{self.instance_id}] {message}')

    def _random_sleep(self, min_sec=30, max_sec=60):
        """
        核心防限速控制：模拟人类阅读和操作的长时间停顿，确保网络请求间隔在30-60秒
        """
        sleep_time = random.uniform(min_sec, max_sec)
        self._log(f"Anti-Rate-Limit: Sleeping for {sleep_time:.2f} seconds...")
        time.sleep(sleep_time)

    def _get_free_port(self):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(('', 0))
            return s.getsockname()[1]

    def set_config(self, config: VSPlgConfig):
        """设置 API 的配置信息"""
        self.config = config
        self.free_config = config.free_config or {}

    def set_log(self, logger: Callable[[str], None]) -> None:
        """设置日志输出工具"""
        self.logger = logger
        
    def keep_alive(self):
        pass

    def health_check(self) -> bool:
        if not self.is_healthy:
            return False
        if self.page is None:
            return False
        try:
            if not self.page.run_js("return 1;"):
                return False
        except:
            return False
            
        if self.config.session_max_life > 0:
            current_time = time.time()
            elapsed_time = current_time - self.session_create_time
            if elapsed_time > self.config.session_max_life:
                self._log(f"Session expired.")
                return False
        return True
    
    def _save_screenshot(self, name_prefix):
        try:
            timestamp = int(time.time())
            filename = f"{self.instance_id}_{name_prefix}_{timestamp}.jpg"
            save_path = os.path.join("data", filename)
            os.makedirs("data", exist_ok=True)
            self.page.get_screenshot(path=save_path, full_page=False)
            self._log(f"Screenshot saved to {save_path}")
        except Exception as e:
            self._log(f"Failed to save screenshot: {e}")

    def create_session(self) -> None:
        """创建一个新的会话 (包含初始化浏览器、过CF验证和执行登录)"""
        
        self._log(f"Initializing Session (ID: {self.instance_id})...")

        def get_free_port():
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.bind(('', 0))
                return s.getsockname()[1]
        
        co = ChromiumOptions()
        debug_port = get_free_port()
        self._log(f"Assigned Debug Port: {debug_port}")
        self._log(f"Account id={self.config.account.id}, proxy id={self.config.proxy.id}")
        
        co.set_local_port(debug_port)       
        co.set_user_data_path(self.user_data_path)
        
        chrome_path = configure.CHROME_PATH
        if not chrome_path:
            chrome_path = os.getenv("CHROME_BIN")
        if chrome_path and os.path.exists(chrome_path):
            co.set_paths(browser_path=chrome_path)
        
        if self.config.proxy and self.config.proxy.ip:
            p = self.config.proxy
            if p.username and p.password:
                self._log(f"Starting Proxy Tunnel for {p.ip}...")
                exit_node = {
                    "name": "ExitNode",
                    "type": p.proto,
                    "server": p.ip,
                    "port": p.port,
                    "username": p.username,
                    "password": p.password
                }
                relay_node = None
                if configure.MIHOMO_RELAY_NODES:
                    relay_node = random.choice(configure.MIHOMO_RELAY_NODES)
                mihomo_path = configure.MIHOMO_BIN_PATH
                if not mihomo_path:
                    mihomo_path = os.getenv("MIHOMO_BIN")
                if not mihomo_path:
                    raise BizLogicError(message='Mihomo path is null, You need set mihomo bin path in configure or os env')
                self.tunnel = MihomoTunnel(mihomo_path, exit_node=exit_node, relay_node=relay_node)
                local_proxy = self.tunnel.start()
                self._log(f"Tunnel started at {local_proxy}")
                co.set_argument(f'--proxy-server={local_proxy}')
            else:
                proxy_str = f"{p.proto}://{p.ip}:{p.port}"
                co.set_argument(f'--proxy-server={proxy_str}')
        else:
            self._log("[WARN] No proxy configured!")

        specific_fp = FingerprintGenerator().generate(self.config.account.username)
        fp_seed = specific_fp.get("seed")
        fp_platform = specific_fp.get("platform")
        fp_brand = specific_fp.get("brand")
        self._log(f'browser fingerprint seed={fp_seed}')
        
        co.headless(False)
        co.set_argument('--no-sandbox')
        co.set_argument('--disable-dev-shm-usage')
        co.set_argument('--window-size=1920,1080')
        co.set_argument('--disable-blink-features=AutomationControlled')
        co.set_argument(f"--fingerprint={fp_seed}")
        co.set_argument(f"--fingerprint-platform={fp_platform}")
        co.set_argument(f"--fingerprint-brand={fp_brand}")
        
        self.page = ChromiumPage(co)
        # 获取基础 URL
        usa_url = self.free_config.get('usa_url', '') 
        self._log(f"Navigating: {usa_url}")
        self.page.get(usa_url)
        self._random_sleep(3, 5)
        
        if 'Attention Required! | Cloudflare' in self.page.title and 'Sorry, you have been blocked' in self.page.html:
            self._log(f'Block by cloudflare, try refresh...')
            self.page.refresh()
            self.page.wait.doc_loaded()
        
        cf_bypasser = CloudflareBypasser(self.page, log=self.config.debug)
        if not cf_bypasser.bypass(max_retry=6):
            raise BizLogicError("Cloudflare bypass timeout")
        self._random_sleep(3, 5) 
        cf_bypasser.handle_waiting_room()
        
        self._log("Init humanize tools...")
        self.mouse = HumanMouse(self.page, debug=False)
        self.keyboard = HumanKeyboard(self.page)
        
        viewport_width = self.page.rect.viewport_size[0]
        viewport_height = self.page.rect.viewport_size[1]
        init_x = random.randint(10, viewport_width - 10)
        init_y = random.randint(10, viewport_height - 10)
        self.mouse.move(init_x, init_y) 
        
        username = self.config.account.username
        password = self.config.account.password
        security = self.free_config.get('security', {})

        max_steps = 15 # 由于状态多，步数可以稍微调大一点
        stuck_counter = 0
        last_url = ""
        session_created = False
        has_submitted_login = False
        
        for step in range(max_steps):
            self.page.wait.doc_loaded()
            time.sleep(1) # 这个用于等待页面DOM渲染，保留短时，因为不是发请求
            
            current_url = self.page.url
            current_title = self.page.title.lower()
            current_html_content = self.page.html
            self._log(f"--- [Router Step {step+1}] Current URL: {current_url} ---")
            
            # --- [异常处理层] ---
            if current_url == last_url:
                stuck_counter += 1
            else:
                last_url = current_url
                stuck_counter = 0
            if stuck_counter >= 3:
                self._log("[WARN] Page stucked, try to refresh...")
                self.page.refresh()
                self._random_sleep(30, 40)
                stuck_counter = 0
                continue
            
            # 网络出现故障,直接重试
            server_error_indicators = ["502 Bad Gateway", "503 Service Temporarily Unavailable"]
            if any(err in current_html_content for err in server_error_indicators):
                self._log(f"[WARN] Server network error, try to refresh (Step: {step})...")
                self.page.refresh()
                self._random_sleep(30, 40)
                continue
            
            cloudflare_blocked_indicators = [
                "Sorry, you have been blocked" in current_html_content,
                "You are being rate limited" in current_html_content,
                "Cloudflare Ray ID" in current_html_content
            ]
            if any(cloudflare_blocked_indicators):
                raise BizLogicError(message="Blocked by Cloudflare WAF. Need to change IP or browser fingerprint.")
            
            # 遇到五秒盾先绕盾
            if "just a moment" in current_title:
                if not cf_bypasser.bypass(max_retry=3):
                    continue
                self._random_sleep(3, 5)
                cf_bypasser.handle_waiting_room()
                continue
            
            if self.page.ele('#post_select', timeout=1):
                self._log("🎉 Successfully reached the Slot Search page (Target Page). Session created successfully!")                
                self.session_create_time = time.time()
                session_created = True
                break

            # 状态 2: 密保问题页面
            elif self.page.ele('xpath://input[starts-with(@id, "kba") and contains(@id, "_response")]', timeout=1):
                self._log("[State] Security question verification detected. Filling in answers...")
                answer_eles = self.page.eles('xpath://input[starts-with(@id, "kba") and contains(@id, "_response")]')
                for ans_ele in answer_eles:
                    ele_id = ans_ele.attr('id')
                    match = re.search(r'kba(\d+)_response', ele_id)
                    if match:
                        q_num = match.group(1)
                        config_key = f"{q_num}_quest"
                        q_data = security.get(config_key)
                        ans_text = q_data.get('a')
                        self.mouse.human_click_ele(ans_ele)
                        self.keyboard.type_text(ans_text, humanize=True)
                        self._log(f"-> Find input {ele_id}, successfully filled in the answer for question {q_num}.")
                        
                continue_btn = self.page.ele('#continue')
                self.mouse.human_click_ele(continue_btn)
                self._log("Security answers submitted. Waiting for redirection...")
                self._random_sleep(30, 40) 
                continue

            # 状态 1: 登录页面
            elif self.page.ele('#signInName', timeout=1):
                self._log("[State] Login page detected. Submitting credentials...")
                username_selector = '#signInName'
                username_input = self.page.ele(username_selector)
                self.mouse.human_click_ele(username_input)
                username_input.clear()
                self.keyboard.type_text(username, humanize=True)
                self._random_sleep(3, 5) 
                
                password_selector = '#password'
                password_input = self.page.ele(password_selector)
                self.mouse.human_click_ele(password_input)
                password_input.clear()
                self.keyboard.type_text(password, humanize=True)
                self._random_sleep(3, 5) 
                
                continue_btn_selector = '#continue'
                continue_btn = self.page.ele(continue_btn_selector)
                self.mouse.human_click_ele(continue_btn)                
                has_submitted_login = True
                self._log("Login form submitted. Waiting for the next step to load...")
                self._random_sleep(30, 40)
                continue
            
            elif self.page.ele('#continue_application', timeout=1):
                schedule_btn = self.page.ele('#continue_application', timeout=0.5)
                self._log("Currently in first-time booking mode, clicking to proceed...")
                self.mouse.human_click_ele(schedule_btn)
                self._random_sleep(30, 40)
            

            # 状态 3: 预约主页(控制台) -> 选择首签或改签
            elif self.page.ele('#atlas-sidebar', timeout=1):
                self._log("[State] At the main booking dashboard. Looking for navigation button...")
                reschedule_btn_selector = '#reschedule_appointment'
                schedule_btn_selector = '#schedule_appointment'
                reschedule_btn = self.page.ele(reschedule_btn_selector, timeout=0.5)
                if reschedule_btn:
                    self._log("Currently in rescheduling mode, clicking to proceed...")
                    self.mouse.human_click_ele(reschedule_btn)
                    self._random_sleep(30, 40)
                else:
                    schedule_btn = self.page.ele(schedule_btn_selector, timeout=0.5)
                    if schedule_btn:
                        self._log("Currently in first-time booking mode, clicking to proceed...")
                        self.mouse.human_click_ele(schedule_btn)
                        self._random_sleep(30, 40)
                    else:
                        self._log("Not found schedule or reschedule button. The page may still be loading...")
                        time.sleep(2)
                continue

            else:
                self._log("[State] In unknown or transitional state. Waiting for next polling cycle...")
                time.sleep(2)
                
        if not session_created:
            raise BizLogicError(f"Failed to reach appointment-booking after {max_steps} navigation steps. Stuck at: {self.page.url}")

    def query(self, apt_type: AppointmentType) -> VSQueryResult:
        """查询可用的签证预约信息"""
        self._log("Querying available slots...")
        res = VSQueryResult()
        res.success = False
        
        self.page.refresh()
        
        if "just a moment" in self.page.title:
            cf_bypasser = CloudflareBypasser(self.page, log=self.config.debug)
            if not cf_bypasser.bypass(max_retry=5):
                raise BizLogicError("Cloudflare bypass timeout")
            self._random_sleep(3, 5)
            cf_bypasser.handle_waiting_room()
        
        current_url = self.page.url.lower()
        if 'auth' in current_url or 'login' in current_url:
            self.is_healthy = False
            raise SessionExpiredOrInvalidError()
        
        applicant = self.free_config.get('applicant')
        location_name = self.free_config.get('location')
        location_id = self.LOCATIONS.get(location_name.upper(), {}).get('id')

        # 2. 等待页面元素
        self.page.ele(f"xpath://label[text()='{applicant}']", timeout=60)
        post_select = self.page.ele('#post_select', timeout=60)
        
        # ==================== 新增：网络监听逻辑 ====================
        
        # 开启监听目标 API
        target_api = 'get-family-consular-schedule-days'
        self.page.listen.start(target_api)
        
        # 3. 选择领事馆 (此操作会触发上述 API 的 AJAX 请求)
        self.page.ele(f"xpath://select[@id='post_select']/option[@value='{location_id}']", timeout=60)
        post_select.select.by_value(location_id) 
        
        # 4. 等待拦截 API 响应 (设置超时时间)
        self._log("Waiting for schedule dates API response...")
        packet = self.page.listen.wait(timeout=30)
        self.page.listen.stop()
        
        if not packet:
            raise BizLogicError("Timeout waiting for schedule API response")
            
        status_code = packet.response.status
        raw_resp = packet.response.raw_body
        self._log(f"API Response Status: {status_code}")
        
        # 处理 HTTP 返回码不是200的情况
        if status_code != 200:
            if status_code == 403:
                raise PermissionDeniedError(f"HTTP 403: {raw_resp[:512]}")
            if status_code == 429:
                self.is_healthy = False
                raise RateLimiteddError(f"HTTP 429: {raw_resp[:512]}")
            raise BizLogicError(f"HTTP {status_code} error. resp={raw_resp[0:512]}")

        # 5. 解析返回的数据
        available_dates = []
        match = re.search(r'(\{.*\})', raw_resp, re.DOTALL)
        if match:
            json_str = match.group(1)
            data = json.loads(json_str)
        else:
            data = json.loads(raw_resp)

        schedule_days = data.get("ScheduleDays", [])
        if schedule_days:
            for day_obj in schedule_days:
                date_str = day_obj.get("Date")
                if date_str:
                    available_dates.append(date_str)

        # 6. 处理最终结果 (保持与你原有返回结构一致)
        if available_dates:
            # 确保日期是有序的
            available_dates.sort()
            
            res.success = True
            res.availability_status = AvailabilityStatus.Available
            earliest_date = available_dates[0]
            res.earliest_date = datetime.strptime(earliest_date, "%Y-%m-%d")
            res.availability = [
                DateAvailability(date=datetime.strptime(d, "%Y-%m-%d"), times=[])
                for d in available_dates
            ]
            self._log(f"Slot Found! earliest_date={earliest_date}, size={len(available_dates)}")
        else:
            res.success = False
            res.availability_status = AvailabilityStatus.NoneAvailable
            self._log("No slots available.")
            
        return res

    def book(self, slot_info: VSQueryResult, user_inputs) -> VSBookResult:
        """进行预约操作"""
        res = VSBookResult()
        res.success = False
        
        exp_start = user_inputs.get('expected_date_start', '')
        exp_end = user_inputs.get('expected_date_end', '')

        available_dates_str =[
            da.date.strftime("%Y-%m-%d")
            for da in slot_info.availability if da.date
        ]
        
        valid_dates_list = self._filter_dates(available_dates_str, exp_start, exp_end)
        if not valid_dates_list:
            raise NotFoundError(message="No dates match user constraints")
        
        selected_slot_date = random.choice(valid_dates_list)
        book_date_obj = datetime.strptime(selected_slot_date, "%Y-%m-%d").date()
        
        # jQuery UI Datepicker 月份是 0-11
        target_day = str(book_date_obj.day)
        target_month = str(book_date_obj.month - 1) 
        target_year = str(book_date_obj.year)
        
        self._log(f"Target booking date: {selected_slot_date}. Navigating calendar...")
        
        # 1. 适配不在当前日历的情况：通过下拉框选择年份和月份
        year_select = self.page.ele('.ui-datepicker-year', timeout=10)  # 查找年份下拉框
        if year_select and year_select.value != target_year:
            self._log(f"Changing year to {target_year}")
            year_select.select.by_value(target_year)
            self._random_sleep(0.5, 1)

        month_select = self.page.ele('.ui-datepicker-month', timeout=10) # 查找月份下拉框
        if month_select and month_select.value != target_month:
            self._log(f"Changing month to {target_month}")
            month_select.select.by_value(target_month)
            self._random_sleep(0.5, 1)

        # 2. 点击目标日期
        self._log(f"Clicking date {selected_slot_date}...")
        target_cell_selector = f"xpath://td[@data-year='{target_year}' and @data-month='{target_month}']//a[text()='{target_day}']"
        
        target_date_cell = self.page.ele(target_cell_selector, timeout=10)
        if not target_date_cell:
             raise BizLogicError(f"Target date element not found for {selected_slot_date} after navigating.")
             
        self.mouse.human_click_ele(target_date_cell)
        
        self._log("Waiting for available times to load...")
        self._random_sleep(3, 5)

        # 3. 等待并选择时间
        self._log("Selecting earliest available time...")
        slot_time_selector = "css:#time_select input[name='schedule-entries']"
        first_time_radio = self.page.ele(slot_time_selector, timeout=15)
        if not first_time_radio:
            raise BizLogicError("Failed to load time slots after clicking date. Might be rate limited or slots gone.")
            
        selected_slot_time = first_time_radio.parent().text.strip()
        self.mouse.human_click_ele(first_time_radio)
        
        # 4. 点击提交
        self._random_sleep(1, 2)
        self._log("Submitting booking...")
        submit_appointment_selector = "#submitbtn"
        submit_button = self.page.ele(submit_appointment_selector, timeout=10)
        self.mouse.human_click_ele(submit_button)
        self._log("Booking submitted successfully!")

        # 构造返回结果
        res = VSBookResult()
        res.success = True
        res.book_date = selected_slot_date
        res.book_time = selected_slot_time
        res.account = self.config.account.username
        return res
            
    def _filter_dates(self, dates: List[str], start_str: str, end_str: str) -> List[str]:
        if not start_str or not end_str:
            return dates
        valid_dates = []
        s_date = datetime.strptime(start_str[:10], "%Y-%m-%d")
        e_date = datetime.strptime(end_str[:10], "%Y-%m-%d")
        for date_str in dates:
            curr_date = datetime.strptime(date_str, "%Y-%m-%d")
            if s_date <= curr_date <= e_date:
                valid_dates.append(date_str)
        random.shuffle(valid_dates)
        return valid_dates
    
    # --- 资源清理核心方法 ---
    def cleanup(self):
        """
        销毁浏览器并彻底删除临时文件
        """
        if self.page:
            try:
                self.page.quit(force=True)
            except Exception:
                pass
            self.page = None
        
        if os.path.exists(self.root_workspace):
            for _ in range(3):
                try:
                    time.sleep(0.2)
                    shutil.rmtree(self.root_workspace, ignore_errors=True)
                    break
                except Exception as e:
                    self._log(f"Cleanup retry: {e}")
                    time.sleep(0.5)
            
            if os.path.exists(self.root_workspace):
                 self._log(f"[WARN] Failed to fully remove workspace: {self.root_workspace}")
        if self.tunnel:
            try: self.tunnel.stop()
            except: pass
            self.tunnel = None
        
    def __del__(self):
        """
        析构函数：当对象被垃圾回收时自动调用
        """
        self.cleanup()