"""Crawl a site with WebCrawlerAPI, split the markdown by headers, embed with local Ollama, write chunks.jsonl.

Usage:
    WEBCRAWLERAPI_API_KEY=... python build_chunks.py https://example.com/docs --items-limit 50 \
        --whitelist 'https://example\\.com/docs.*' --out chunks.jsonl

Output format and chunking match the shared pipeline, so load.py reads the result as is.
"""
import argparse, hashlib, json, os, re, sys, time

import requests

API = "https://api.webcrawlerapi.com/v1"
OLLAMA = "http://localhost:11434/api/embed"
MODEL, DIM = "nomic-embed-text", 768
MIN_TOKENS, MAX_TOKENS = 40, 500
HARD_CAP = 60  # refuse to crawl more pages than this


def tokens(text):
    return len(text.split()) * 4 // 3  # rough estimate, no tokenizer dependency


def api(method, path, key, **kw):
    # Errors are re-raised without the request object so the key never reaches the output.
    try:
        r = requests.request(method, API + path, headers={"Authorization": f"Bearer {key}"}, timeout=60, **kw)
    except requests.RequestException as e:
        raise SystemExit(f"request to {path} failed: {type(e).__name__}") from None
    if r.status_code >= 400:
        raise SystemExit(f"{method} {path} returned HTTP {r.status_code}: {r.text[:200]}")
    return r.json()


def crawl(url, items_limit, whitelist, key, poll=5, timeout=900):
    body = {"url": url, "items_limit": items_limit, "output_formats": ["markdown"], "respect_robots_txt": True}
    if whitelist:
        body["whitelist_regexp"] = whitelist
    job_id = api("POST", "/crawl", key, json=body)["id"]
    print(f"job {job_id} started", file=sys.stderr)
    t0 = time.time()
    while True:
        job = api("GET", f"/job/{job_id}", key)
        if job["status"] in ("done", "error"):
            return job
        if time.time() - t0 > timeout:
            raise SystemExit(f"job {job_id} still {job['status']} after {timeout}s")
        time.sleep(poll)


def clean(text):
    text = re.sub(r"^\[[^\n]*?\]\(/docs[^)]*\)[^\n]*\n", "", text)         # breadcrumb line
    text = re.sub(r"^\[Open in Claude.*\n", "", text, flags=re.M)           # "Open in ..." links
    return re.sub(r"\[([^\]]+)#\]\(#[^)]*\)", r"\1", text)                  # heading anchors


def split_sections(markdown):
    """Split on H1-H3 outside code fences. Returns [(heading_path, body)]."""
    sections, path, body, in_fence = [], [], [], False

    def flush():
        text = "\n".join(body).strip()
        if text:
            sections.append((" > ".join(path), text))
        body.clear()

    for line in markdown.splitlines():
        if line.startswith("```"):
            in_fence = not in_fence
        m = None if in_fence else re.match(r"(#{1,3}) (.+)", line)
        if m:
            flush()
            path[:] = path[: len(m.group(1)) - 1] + [m.group(2).strip()]
        body.append(line)
    flush()
    return sections


def split_paragraphs(text):
    """Break an oversized section on blank lines, never inside a code fence."""
    parts, cur, in_fence = [], [], False
    for line in text.splitlines():
        in_fence ^= line.startswith("```")
        if not line.strip() and not in_fence and cur:
            parts.append("\n".join(cur)); cur = []
        else:
            cur.append(line)
    if cur:
        parts.append("\n".join(cur))
    chunks, buf = [], ""
    for p in parts:
        if buf and tokens(buf + p) > MAX_TOKENS:
            chunks.append(buf); buf = ""
        buf = f"{buf}\n\n{p}" if buf else p
    return chunks + [buf]


def chunk_page(markdown):
    merged = []
    for heading, text in split_sections(markdown):
        if merged and tokens(text) < MIN_TOKENS:
            merged[-1] = (merged[-1][0], merged[-1][1] + "\n\n" + text)
        else:
            merged.append((heading, text))
    return [(h, part) for h, t in merged for part in (split_paragraphs(t) if tokens(t) > MAX_TOKENS else [t])]


def embed(texts, prefix):
    r = requests.post(OLLAMA, json={"model": MODEL, "input": [prefix + t for t in texts]}, timeout=300)
    r.raise_for_status()
    return r.json()["embeddings"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("url")
    ap.add_argument("--items-limit", type=int, default=50)
    ap.add_argument("--whitelist", default="", help="regexp: only matching URLs are crawled")
    ap.add_argument("--out", default="chunks.jsonl")
    args = ap.parse_args()
    if args.items_limit > HARD_CAP:
        raise SystemExit(f"items-limit above {HARD_CAP} refused")
    key = os.environ.get("WEBCRAWLERAPI_API_KEY")
    if not key:
        raise SystemExit("set WEBCRAWLERAPI_API_KEY")

    t0 = time.time()
    job = crawl(args.url, args.items_limit, args.whitelist, key)
    items = [i for i in job.get("job_items", []) if i.get("status") == "done" and i.get("markdown_content_url")]
    cost = sum(i.get("cost") or 0 for i in job.get("job_items", []))
    print(f"job status {job['status']}: {len(job.get('job_items', []))} items, {len(items)} with markdown, "
          f"summed item cost {cost}", file=sys.stderr)

    rows = []
    for item in sorted(items, key=lambda i: i["original_url"]):
        md = requests.get(item["markdown_content_url"], timeout=60)  # public URL, no auth header
        md.raise_for_status()
        url = item["original_url"]
        for i, (heading, text) in enumerate(chunk_page(clean(md.text))):
            rows.append({"id": hashlib.sha1(f"{url}#{i}".encode()).hexdigest()[:16],
                         "url": url, "heading": heading, "text": text})
    print(f"{len(rows)} chunks from {len(items)} pages, embedding...", file=sys.stderr)

    with open(args.out, "w") as out:
        for i in range(0, len(rows), 32):
            batch = rows[i:i + 32]
            vectors = embed([f"{r['heading']}\n{r['text']}" for r in batch], "search_document: ")
            for r, vec in zip(batch, vectors):
                out.write(json.dumps({**r, "embedding": vec, "model": MODEL, "dim": DIM}) + "\n")
    print(f"wrote {args.out} ({len(rows)} chunks) in {time.time() - t0:.1f}s", file=sys.stderr)


if __name__ == "__main__":
    main()