boss_dp/worker/main.py

# -*- coding: utf-8 -*-
"""
Worker startup entrypoint.
Usage:
  python -m worker.main [--server ws://IP:8000/ws] [--worker-id pc-a] [--worker-name name]
"""
from __future__ import annotations

import argparse
import asyncio
import logging
import os

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "server.settings")

import django  # noqa: E402

django.setup()  # noqa: E402

from tunnel.client import TunnelClient
from worker import config
from worker.recruit_filter_sync import bootstrap_recruit_filter_snapshot
from worker.tasks.registry import register_all_handlers
from worker.ws_client import WorkerWSClient

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(name)-28s  %(levelname)-5s  %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("worker.main")


def parse_args():
    parser = argparse.ArgumentParser(description="Browser Control Worker Agent")
    parser.add_argument("--worker", action="store_true", help=argparse.SUPPRESS)
    parser.add_argument("--server", default=config.SERVER_WS_URL, help=f"WebSocket server URL (default: {config.SERVER_WS_URL})")
    parser.add_argument("--worker-id", default=config.WORKER_ID, help=f"Worker ID (default: {config.WORKER_ID})")
    parser.add_argument("--worker-name", default=config.WORKER_NAME, help=f"Worker name (default: {config.WORKER_NAME})")
    parser.add_argument("--bit-api", default=config.BIT_API_BASE, help=f"BitBrowser local API URL (default: {config.BIT_API_BASE})")
    parser.add_argument("--no-tunnel", action="store_true", help="Disable tunnel client")
    return parser.parse_args()


def _local_port_from_bit_api(bit_api_base: str) -> int:
    try:
        from urllib.parse import urlparse

        p = urlparse(bit_api_base if "://" in bit_api_base else "http://" + bit_api_base)
        return p.port or 54345
    except Exception:
        return config.TUNNEL_LOCAL_PORT


def _extract_host_from_ws_url(ws_url: str) -> str:
    try:
        from urllib.parse import urlparse

        p = urlparse(ws_url)
        return p.hostname or "127.0.0.1"
    except Exception:
        return "127.0.0.1"


async def run(args):
    register_all_handlers()
    logger.info("Task handlers registered")

    # Startup bootstrap: fetch site filters (main1 logic) and persist for frontend dropdown.
    try:
        bootstrap_result = bootstrap_recruit_filter_snapshot(
            worker_id=args.worker_id,
            bit_api_base=args.bit_api,
            logger=logger,
        )
        if bootstrap_result:
            logger.info(
                "Recruit filters synced on startup: account=%s groups=%s options=%s",
                bootstrap_result.get("account_name", ""),
                bootstrap_result.get("groups", 0),
                bootstrap_result.get("flat_options", 0),
            )
    except Exception as e:
        logger.warning("Recruit filter startup sync failed: %s", e)

    client = WorkerWSClient(
        server_url=args.server,
        worker_id=args.worker_id,
        worker_name=args.worker_name,
        bit_api_base=args.bit_api,
    )

    tunnel_enabled = config.TUNNEL_ENABLED and not args.no_tunnel
    tunnel_client = None
    if tunnel_enabled:
        tunnel_host = _extract_host_from_ws_url(args.server)
        tunnel_client = TunnelClient(
            server_host=tunnel_host,
            control_port=config.TUNNEL_CONTROL_PORT,
            stream_port=config.TUNNEL_STREAM_PORT,
            worker_id=args.worker_id,
            local_port=_local_port_from_bit_api(args.bit_api),
        )
        logger.info(
            "Tunnel enabled: expose local %s -> %s (worker_id=%s)",
            _local_port_from_bit_api(args.bit_api),
            tunnel_host,
            args.worker_id,
        )

    logger.info(
        "Worker startup: id=%s, name=%s, server=%s",
        args.worker_id,
        args.worker_name,
        args.server,
    )

    async def run_worker():
        await client.run()

    async def run_tunnel():
        if tunnel_client:
            await tunnel_client.run()

    worker_task = asyncio.create_task(run_worker())
    tunnel_task = asyncio.create_task(run_tunnel()) if tunnel_client else None

    try:
        if tunnel_task is not None:
            await asyncio.gather(worker_task, tunnel_task)
        else:
            await worker_task
    except KeyboardInterrupt:
        logger.info("Keyboard interrupt received, stopping...")
        worker_task.cancel()
        if tunnel_task is not None:
            tunnel_task.cancel()
        try:
            await worker_task
        except asyncio.CancelledError:
            pass
        if tunnel_task is not None:
            try:
                await tunnel_task
            except asyncio.CancelledError:
                pass
            await tunnel_client.stop()
        await client.stop()


def main():
    args = parse_args()
    try:
        asyncio.run(run(args))
    except KeyboardInterrupt:
        logger.info("Worker exited")


if __name__ == "__main__":
    main()