# crawl.py: crawl one docs site to data/<site>/pages.jsonl (url, title, markdown)
import argparse, json, os, pathlib, sys, time
import requests

API = "https://api.webcrawlerapi.com/v1"
SITES = {
    "fastapi": {
        "url": "https://fastapi.tiangolo.com/tutorial/",
        "whitelist_regexp": r"^https://fastapi\.tiangolo\.com/(tutorial|advanced)/",
        "items_limit": 80,
    },
    "htmx": {
        "url": "https://htmx.org/docs/",
        "whitelist_regexp": r"^https://htmx\.org/(docs|reference|attributes|examples)/",
        "items_limit": 50,
    },
}

ap = argparse.ArgumentParser()
ap.add_argument("site", choices=SITES)
ap.add_argument("--out", default="data")
args = ap.parse_args()

out = pathlib.Path(args.out) / args.site
out.mkdir(parents=True, exist_ok=True)
http = requests.Session()
http.headers["Authorization"] = f"Bearer {os.environ['WEBCRAWLERAPI_API_KEY']}"

body = {**SITES[args.site], "output_formats": ["markdown"], "respect_robots_txt": True, "max_age": 0}
r = http.post(f"{API}/crawl", json=body, timeout=30)
r.raise_for_status()
job_id = r.json()["id"]
print(f"started job {job_id}")

t0 = time.time()
while True:
    r = http.get(f"{API}/job/{job_id}", timeout=30)
    r.raise_for_status()
    job = r.json()
    done = sum(1 for i in job.get("job_items", []) if i["status"] == "done")
    print(f"status={job['status']} items_done={done}")
    if job["status"] not in ("new", "in_progress"):
        break
    if time.time() - t0 > 900:
        sys.exit("deadline hit")
    time.sleep(job.get("recommended_pull_delay_ms", 5000) / 1000)
if job["status"] != "done":
    sys.exit(f"job ended with status {job['status']}")

ok = [i for i in job["job_items"] if i["status"] == "done" and i.get("markdown_content_url")]
failed = [{"url": i["original_url"], "status": i["status"], "last_error": i.get("last_error")}
          for i in job["job_items"] if i not in ok]
with open(out / "pages.jsonl", "w") as f:
    for i in ok:
        md = requests.get(i["markdown_content_url"], timeout=60).text  # no API key sent to the content host
        f.write(json.dumps({"url": i["original_url"], "title": i.get("title") or "", "markdown": md}) + "\n")
cost = round(sum(i.get("cost") or 0 for i in job["job_items"]), 4)
meta = {"site": args.site, "job_id": job_id, "finished_at": job.get("finished_at"), "pages": len(ok),
        "failed": failed, "cost_usd": cost, "params": body}
(out / "crawl_meta.json").write_text(json.dumps(meta, indent=2))
print(f"pages={len(ok)} failed={len(failed)} cost={cost}")