# Trouve-KA — enfilement d'URLs depuis un sitemap
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai

"""Enfile dans le frontier les URLs d'un sitemap (urlset ou sitemapindex).

Usage : .venv/bin/python scripts/bootstrap-seeds/enqueue-sitemap.py <sitemap_url> [...]
        [--priority 0.9] [--limit 5000] [--no-follow-index]

- Suit les index de sitemaps (1 niveau), dans l'ordre du fichier.
- Trie les URLs par profondeur de chemin (pages de haut niveau d'abord)
  avant de tronquer à --limit, pour respecter le cap par domaine avec
  les pages les plus importantes.
- N'enfile que les URLs appartenant au domaine du sitemap (garde-fou).
- Le cap dur de 5000 URLs connues/domaine reste appliqué par la base
  (enqueue_urls_bulk).
"""

import argparse
import asyncio
import gzip
import re
import sys
import urllib.request

from trouveka.config import get_settings
from trouveka.database import Database
from trouveka.logging import get_logger
from trouveka.shared import canonicalize_url, extract_domain

log = get_logger("crawler.enqueue_sitemap")

LOC_RE = re.compile(r"<loc>\s*([^<\s]+)\s*</loc>")
CHUNK = 1000


def fetch(url: str, user_agent: str, timeout: float = 30.0) -> str:
    req = urllib.request.Request(url, headers={"User-Agent": user_agent})
    with urllib.request.urlopen(req, timeout=timeout) as resp:
        body = resp.read(20_000_000)
        if resp.headers.get("Content-Encoding") == "gzip" or url.endswith(".gz"):
            try:
                body = gzip.decompress(body)
            except OSError:
                pass
        return body.decode("utf-8", errors="replace")


def collect_locs(sitemap_url: str, user_agent: str, *, follow_index: bool) -> list[str]:
    body = fetch(sitemap_url, user_agent)
    locs = LOC_RE.findall(body)
    if "<sitemapindex" in body:
        if not follow_index:
            log.info("index ignoré (--no-follow-index)", extra={"ctx": {"url": sitemap_url}})
            return []
        urls: list[str] = []
        for child in locs:
            try:
                child_body = fetch(child, user_agent)
            except Exception as exc:  # noqa: BLE001 — un enfant cassé ne bloque pas le reste
                log.info("sitemap enfant illisible", extra={"ctx": {"url": child, "err": str(exc)}})
                continue
            if "<sitemapindex" in child_body:
                continue  # un seul niveau d'indirection
            urls.extend(LOC_RE.findall(child_body))
        return urls
    return locs


async def enqueue(sitemap_url: str, *, priority: float, limit: int, follow_index: bool) -> int:
    settings = get_settings()
    sitemap_domain = extract_domain(sitemap_url)
    if not sitemap_domain:
        log.info("URL de sitemap invalide", extra={"ctx": {"url": sitemap_url}})
        return 0

    raw = collect_locs(sitemap_url, settings.crawler_user_agent, follow_index=follow_index)
    items: list[tuple[str, str, float]] = []
    seen: set[str] = set()
    for loc in raw:
        url = canonicalize_url(loc)
        if not url or url in seen:
            continue
        domain = extract_domain(url)
        if domain != sitemap_domain:
            continue  # garde-fou : on reste sur le domaine du sitemap
        seen.add(url)
        items.append((url, domain, priority))

    # Pages de haut niveau d'abord (moins de segments de chemin), ordre du
    # sitemap préservé à profondeur égale, puis troncature au budget.
    items.sort(key=lambda t: len([s for s in t[0].split("/", 3)[-1].split("/") if s]))
    items = items[:limit]

    db = Database(settings.database_url, pool_min=1, pool_max=3)
    await db.connect()
    added = 0
    try:
        for i in range(0, len(items), CHUNK):
            _, n = await db.enqueue_urls_bulk(items[i : i + CHUNK], depth=1, source_url_id=None)
            added += n
    finally:
        await db.close()
    log.info(
        "sitemap enfilé",
        extra={"ctx": {"sitemap": sitemap_url, "locs": len(raw), "candidates": len(items), "added": added}},
    )
    return added


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("sitemaps", nargs="+", help="URL(s) de sitemap.xml ou d'index")
    ap.add_argument("--priority", type=float, default=0.9)
    ap.add_argument("--limit", type=int, default=5000, help="max d'URLs candidates par sitemap")
    ap.add_argument("--no-follow-index", action="store_true")
    args = ap.parse_args()

    total = 0
    for sm in args.sitemaps:
        total += asyncio.run(
            enqueue(sm, priority=args.priority, limit=args.limit, follow_index=not args.no_follow_index)
        )
    print(f"total ajouté au frontier : {total}", file=sys.stderr)


if __name__ == "__main__":
    main()