bangwen/collectors/runner.py

import argparse
from concurrent.futures import ThreadPoolExecutor, as_completed
import fnmatch
import importlib.util
import os
from pathlib import Path
import re
import signal
import sys
import time

try:
    import tomllib
except ModuleNotFoundError:
    import tomli as tomllib

from common.declarative import load_declarative_collectors
from common.protocol import emit_payload

ROOT_DIR = Path(__file__).resolve().parent
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

running = True


def handle_signal(sig, frame):
    global running
    sys.stderr.write("\nReceived shutdown signal. Exiting runner...\n")
    running = False


signal.signal(signal.SIGINT, handle_signal)
signal.signal(signal.SIGTERM, handle_signal)


def load_config() -> dict:
    config_path = os.getenv("CONFIG_PATH", str(ROOT_DIR / "config.toml"))
    cfg = {
        "runner": {
            "schedule_interval_hours": 6,
            "concurrency": 3,
            "domain_delay_seconds": 3.0,
            "core_ingest_url": os.getenv("CORE_INGEST_URL", "http://core:8080/api/v1/ingest"),
        },
        "filter": {
            "provinces": [],
            "include": [],
            "tags": [],
            "exclude": [],
        },
    }

    if Path(config_path).is_file():
        try:
            with open(config_path, "rb") as f:
                loaded = tomllib.load(f)
                if "runner" in loaded:
                    cfg["runner"].update(loaded["runner"])
                if "filter" in loaded:
                    cfg["filter"].update(loaded["filter"])
        except Exception as e:
            sys.stderr.write(f"Warning: Failed to load {config_path}: {e}\n")

    if "domain_concurrency" in cfg["runner"]:
        cfg["runner"]["concurrency"] = cfg["runner"]["domain_concurrency"]

    if os.getenv("CORE_INGEST_URL"):
        cfg["runner"]["core_ingest_url"] = os.getenv("CORE_INGEST_URL")
    if os.getenv("SCHEDULE_INTERVAL_HOURS"):
        cfg["runner"]["schedule_interval_hours"] = int(os.getenv("SCHEDULE_INTERVAL_HOURS"))
    if os.getenv("DOMAIN_CONCURRENCY"):
        cfg["runner"]["concurrency"] = int(os.getenv("DOMAIN_CONCURRENCY"))
    if os.getenv("DOMAIN_DELAY_SECONDS"):
        cfg["runner"]["domain_delay_seconds"] = float(os.getenv("DOMAIN_DELAY_SECONDS"))
    if os.getenv("COLLECTORS_PROVINCES"):
        cfg["filter"]["provinces"] = [
            p.strip() for p in os.getenv("COLLECTORS_PROVINCES").split(",") if p.strip()
        ]
    if os.getenv("COLLECTORS_INCLUDE"):
        cfg["filter"]["include"] = [
            i.strip() for i in os.getenv("COLLECTORS_INCLUDE").split(",") if i.strip()
        ]

    return cfg


def discover_collectors() -> list[dict]:
    collectors = []

    sources_dir = ROOT_DIR / "sources"
    adapters_dir = ROOT_DIR / "adapters"

    if sources_dir.is_dir():
        for toml_file in sorted(sources_dir.rglob("*.toml")):
            try:
                for entry in load_declarative_collectors(toml_file, base_dir=ROOT_DIR):
                    collectors.append(entry)
            except Exception as e:
                sys.stderr.write(f"Error loading source {toml_file}: {e}\n")

    if adapters_dir.is_dir():
        for py_file in sorted(adapters_dir.rglob("*.py")):
            if py_file.name == "__init__.py":
                continue

            rel_parts = py_file.relative_to(adapters_dir).with_suffix("").parts
            module_name = "adapters." + ".".join(rel_parts)
            try:
                spec = importlib.util.spec_from_file_location(module_name, py_file)
                if not spec or not spec.loader:
                    continue
                mod = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(mod)

                crawl_fn = getattr(mod, "crawl", None)
                if not callable(crawl_fn):
                    continue

                meta = getattr(mod, "META", None)
                if not meta or not isinstance(meta, dict):
                    meta = {
                        "source_id": f"adapter.{'.'.join(rel_parts)}",
                        "name": py_file.stem,
                        "province": rel_parts[0] if len(rel_parts) > 1 else "",
                        "city": None,
                        "tags": [],
                    }

                domain = meta.get("domain")
                if not domain:
                    start_url = getattr(mod, "START_URL", None)
                    if start_url:
                        from urllib.parse import urlsplit
                        domain = urlsplit(start_url).netloc.lower()
                    else:
                        domain = meta.get("source_id", "adapter")
                meta["domain"] = domain

                collectors.append({
                    "meta": meta,
                    "crawl_fn": crawl_fn,
                    "file": str(py_file.relative_to(ROOT_DIR)),
                    "type": "py",
                    "domain": domain,
                })
            except Exception as e:
                sys.stderr.write(f"Error loading adapter {py_file}: {e}\n")

    return collectors


PROVINCE_ALIAS = {
    "yn": ["yn", "云南", "云南省"],
    "云南": ["yn", "云南", "云南省"],
    "云南省": ["yn", "云南", "云南省"],
    "sc": ["sc", "四川", "四川省"],
    "四川": ["sc", "四川", "四川省"],
    "四川省": ["sc", "四川", "四川省"],
}


def is_enabled(meta: dict, filter_cfg: dict) -> bool:
    source_id = meta.get("source_id", "")
    province = meta.get("province", "")
    tags = set(meta.get("tags", []))

    for pat in filter_cfg.get("exclude", []):
        if fnmatch.fnmatch(source_id, pat):
            return False

    provinces = filter_cfg.get("provinces", [])
    include = filter_cfg.get("include", [])
    req_tags = filter_cfg.get("tags", [])

    if not provinces and not include and not req_tags:
        return True

    if provinces:
        for p in provinces:
            aliases = PROVINCE_ALIAS.get(p.lower(), [p])
            if province in aliases or source_id.startswith(f"{p}."):
                return True
    if any(fnmatch.fnmatch(source_id, pat) for pat in include):
        return True
    if any(tag in tags for tag in req_tags):
        return True

    return False


def run_one(collector: dict, ingest_url: str):
    meta = collector["meta"]
    name = meta.get("name", meta.get("source_id"))
    crawl_fn = collector["crawl_fn"]

    try:
        sys.stderr.write(f"[{meta['source_id']}] Starting: {name}...\n")
        payload = crawl_fn(max_items=10)
        emit_payload(payload, push_url=ingest_url)
        sys.stderr.write(
            f"[{meta['source_id']}] Completed: {len(payload.items)} items pushed.\n"
        )
        return True
    except Exception as exc:
        raw_err = str(exc)
        sanitized_err = re.sub(r"([?&][\w-]+)=([^&\s]+)", r"\1=[redacted]", raw_err)
        if len(sanitized_err) > 150:
            sanitized_err = sanitized_err[:147] + "..."
        sys.stderr.write(f"[{meta['source_id']}] Failed: {sanitized_err}\n")
        try:
            from common.protocol import CollectorMeta, IngestPayload
            err_payload = IngestPayload(
                source_id=meta["source_id"],
                batch_id=f"err-{int(time.time())}",
                collector=CollectorMeta(
                    name=name,
                    runtime="python/3.12",
                    error_message=sanitized_err,
                ),
                items=[],
            )
            emit_payload(err_payload, push_url=ingest_url)
        except Exception as e:
            sys.stderr.write(f"[{meta['source_id']}] Failed to emit heartbeat: {e}\n")
        return False


def run_domain_queue(
    domain: str, collectors: list[dict], ingest_url: str, delay_sec: float
):
    total = len(collectors)
    sys.stderr.write(
        f"[{domain}] Starting queue ({total} channel{'s' if total > 1 else ''})...\n"
    )
    for i, col in enumerate(collectors):
        if not running:
            break
        if i > 0 and delay_sec > 0:
            sys.stderr.write(
                f"[{domain}] Polite pause ({delay_sec}s) before next channel...\n"
            )
            time.sleep(delay_sec)
        run_one(col, ingest_url)


def run_cycle(cfg: dict, active_collectors: list[dict]):
    ingest_url = cfg["runner"]["core_ingest_url"]
    concurrency = max(1, int(cfg["runner"].get("concurrency", 3)))
    delay_sec = max(0.0, float(cfg["runner"].get("domain_delay_seconds", 3.0)))

    domain_groups: dict[str, list[dict]] = {}
    for col in active_collectors:
        domain = col.get("domain") or col.get("meta", {}).get("domain") or "default"
        domain_groups.setdefault(domain, []).append(col)

    sys.stderr.write(
        f"\n=== Starting crawl cycle ({len(active_collectors)} active collectors across {len(domain_groups)} domains, domain concurrency: {concurrency}, intra-domain delay: {delay_sec}s) ===\n"
    )

    with ThreadPoolExecutor(max_workers=concurrency) as pool:
        futures = {
            pool.submit(run_domain_queue, domain, cols, ingest_url, delay_sec): domain
            for domain, cols in domain_groups.items()
        }
        for future in as_completed(futures):
            if not running:
                pool.shutdown(wait=False, cancel_futures=True)
                break

    sys.stderr.write("=== Crawl cycle finished ===\n\n")


def wait_for_gateway(url: str, timeout: int = 30):
    import httpx
    base_url = url.split("/api/v1")[0]
    health_url = f"{base_url}/health"
    start = time.time()
    while time.time() - start < timeout and running:
        try:
            r = httpx.get(health_url, timeout=2.0)
            if r.status_code == 200:
                sys.stderr.write(f"Core gateway at {health_url} is ready.\n")
                return True
        except Exception:
            pass
        time.sleep(1)
    sys.stderr.write(f"Gateway health check reached timeout at {health_url}, proceeding.\n")
    return False


def main():
    parser = argparse.ArgumentParser(description="Bangwen Collector Runner")
    parser.add_argument("--list", action="store_true", help="List all discovered and enabled collectors")
    parser.add_argument("--once", action="store_true", help="Run one crawl cycle and exit")
    args = parser.parse_args()

    cfg = load_config()
    all_collectors = discover_collectors()
    active_collectors = [c for c in all_collectors if is_enabled(c["meta"], cfg["filter"])]

    if args.list:
        print(f"Total collectors discovered: {len(all_collectors)}")
        print(f"Enabled collectors ({len(active_collectors)}):")
        for c in all_collectors:
            enabled = is_enabled(c["meta"], cfg["filter"])
            mark = "[x]" if enabled else "[ ]"
            meta = c["meta"]
            mode = c.get("type", "py")
            domain = c.get("domain", "")
            print(f"  {mark} {meta['source_id']:<20} ({meta.get('province')}) [{mode}] [{domain}] - {meta.get('name')} ({c['file']})")
        return

    interval_sec = cfg["runner"]["schedule_interval_hours"] * 3600
    active_ids = [c["meta"]["source_id"] for c in active_collectors]
    sys.stderr.write(
        f"Bangwen Collector Runner started.\n"
        f"Configured filter: provinces={cfg['filter']['provinces']}, include={cfg['filter']['include']}, tags={cfg['filter']['tags']}\n"
        f"Active collectors ({len(active_collectors)}): {', '.join(active_ids)}\n"
        f"Interval: {interval_sec}s, Concurrency: {cfg['runner']['concurrency']}\n\n"
    )

    wait_for_gateway(cfg["runner"]["core_ingest_url"])
    run_cycle(cfg, active_collectors)

    if args.once:
        return

    last_run = time.time()
    while running:
        time.sleep(1)
        if time.time() - last_run >= interval_sec:
            run_cycle(cfg, active_collectors)
            last_run = time.time()


if __name__ == "__main__":
    main()