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 _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)
        time.sleep(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.load_start(timeout=2)
            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")
        time.sleep(3)
        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)
            
            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.page.wait.load_start(timeout=5)
                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})...")
                time.sleep(2)
                self.page.refresh()
                self.page.wait.load_start(timeout=5)
                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:
                cf_bypasser.bypass(max_retry=3)
                time.sleep(3)
                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')
                        ans_ele.input(ans_text)
                        self._log(f"-> Find input {ele_id}, successfully filled in the answer for question {q_num}.")
          

                self.page.ele('#continue').click()
                self._log("Security answers submitted. Waiting for redirection...")
                time.sleep(3)
                continue

            # 状态 1: 登录页面
            elif self.page.ele('#signInName', timeout=1):
                self._log("[State] Login page detected. Submitting credentials...")
                
                username_input = self.page.ele('#signInName')
                username_input.clear()
                username_input.input(username)
                
                password_input = self.page.ele('#password')
                password_input.clear()
                password_input.input(password)
                
                self.page.ele('#continue').click()
                has_submitted_login = True
                self._log("Login form submitted. Waiting for the next step to load...")
                time.sleep(3)
                continue

            # 状态 3: 预约主页(控制台) -> 选择首签或改签
            elif self.page.ele('#atlas-sidebar', timeout=1):
                self._log("[State] At the main booking dashboard. Looking for navigation button...")
                
                reschedule_btn = self.page.ele('#reschedule_appointment', timeout=0.5)
                if reschedule_btn:
                    self._log("-> Detected [Reschedule Appointment]. Currently in rescheduling mode, clicking to proceed...")
                    reschedule_btn.click()
                else:
                    schedule_btn = self.page.ele('xpath://ul[@id="atlas-sidebar"]//a[text()="安排预约" or text()="New Appointment" or text()="Schedule Appointment"]', timeout=0.5)
                    if schedule_btn:
                        self._log("-> Detected [Schedule Appointment]. Currently in first-time booking mode, clicking to proceed...")
                        schedule_btn.click()
                    else:
                        self._log("-> [WARN] Sidebar found, but no 'Schedule' or 'Reschedule' button detected. The page may still be loading...")
                        
                time.sleep(3)
                continue

            else:
                self._log("[State] In unknown or transitional state. No matching UI elements found. 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
        # 1. 刷新页面以获取最新数据
        self.page.refresh()
        time.sleep(3)
        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)
        
        # 3. 选择领事馆
        self.page.ele(f"xpath://select[@id='post_select']/option[@value='{location_id}']", timeout=60)
        post_select.select.by_value(location_id) 
        
        # 4. 等待日历加载
        self.page.ele('xpath://p[@id="datepicker-message"]', timeout=60)        
        # 5. 抓取所有可用日期
        available_dates = []
        day_cells = self.page.eles("css:td[data-handler='selectDay'].greenday")
        
        for cell in day_cells:
            day = cell.ele("css:a.ui-state-default").text
            month = int(cell.attr("data-month")) + 1
            year = int(cell.attr("data-year"))
            available_dates.append(date(year, month, int(day)).isoformat())

        if available_dates:
            res.success = True
            res.availability_status = AvailabilityStatus.Available
            earliest_date = available_dates[0]
            earliest_dt = datetime.strptime(earliest_date, "%Y-%m-%d")
            res.earliest_date = earliest_dt
            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_start_date', '')
        exp_end = user_inputs.get('expected_end_date', '')

        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_date = random.choice(valid_dates_list)

        book_date_obj = datetime.strptime(selected_date, "%Y-%m-%d").date()
        day_to_click = str(book_date_obj.day)
        month_to_click = str(book_date_obj.month - 1)
        year_to_click = str(book_date_obj.year)
        
        target_cell_xpath = f"xpath://td[@data-year='{year_to_click}' and @data-month='{month_to_click}']//a[text()='{day_to_click}']"
        
        # 1. 点击目标日期
        self._log(f"Clicking date {selected_date}...")
        self.page.ele(target_cell_xpath, timeout=30).click(by_js=True)

        # 2. 等待并选择时间
        self._log("Selecting earliest available time...")
        first_time_radio = self.page.ele("css:#time_select input[name='schedule-entries']", timeout=30)
        booked_time = first_time_radio.parent().text.strip()
        first_time_radio.click()

        # 3. 提交
        self._log("Submitting booking...")
        submit_button = self.page.ele("#submitbtn", timeout=30)
        submit_button.click()
        
        self._log("Booking submitted successfully!")
        
        # 构造返回结果
        res = VSBookResult()
        res.success = True
        res.book_date = selected_date
        res.book_time = booked_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()