# retoor <retoor@molodetz.nl>
import argparse
import asyncio
import hashlib
import importlib.util
import json
import sys
from pathlib import Path
from urllib.parse import urlsplit

MAX_JOBS = 20
HOST_JOBS = 6
BULK_THRESHOLD = 4
DEFAULT_JOBS = 10
TIMEOUT = 15.0
BUDGET = 60.0
MAX_BYTES = 8 * 1024 * 1024
FILE_MAX_BYTES = 32 * 1024 * 1024


def load_guard():
    path = Path(__file__).with_name("guard.py")
    spec = importlib.util.spec_from_file_location("fetch_guard", path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def host_of(url):
    return (urlsplit(url).hostname or "").lower()


def file_name(url, kind, index, wanted=""):
    if wanted:
        cleaned = "".join(c for c in wanted if c.isalnum() or c in ("-", "_", ".")).strip(".")
        if cleaned:
            return cleaned
    digest = hashlib.sha256(url.encode("utf-8")).hexdigest()[:12]
    if kind != "file":
        return f"{index:03d}-{digest}.html"
    suffix = Path(urlsplit(url).path).suffix.lower()
    if len(suffix) < 2 or len(suffix) > 5 or not suffix[1:].isalnum():
        suffix = ".bin"
    return f"{index:03d}-{digest}{suffix}"


def split_lanes(items, bulk_threshold=BULK_THRESHOLD):
    counts = {}
    for item in items:
        if item.get("kind") == "file" and not item.get("session"):
            host = host_of(item.get("url", ""))
            counts[host] = counts.get(host, 0) + 1
    fast, bulk = [], []
    for pos, item in enumerate(items):
        row = dict(item, _pos=pos)
        if (item.get("kind") == "file" and not item.get("session")
                and counts.get(host_of(item.get("url", "")), 0) >= bulk_threshold):
            bulk.append(row)
        else:
            fast.append(row)
    return fast, bulk


async def gather_batch(items, out_dir, jobs, timeout, fetch_bytes):
    sem = asyncio.Semaphore(max(1, min(int(jobs), MAX_JOBS)))
    host_sems = {}
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)

    def host_sem(url):
        key = host_of(url)
        found = host_sems.get(key)
        if found is None:
            found = asyncio.Semaphore(HOST_JOBS)
            host_sems[key] = found
        return found

    async def one(index, item):
        url = item.get("url", "")
        name = file_name(url, item.get("kind", "page"), index, item.get("name", ""))
        row = {"url": url, "ok": False, "path": "", "bytes": 0, "error": ""}
        if "_pos" in item:
            row["_pos"] = item["_pos"]
        async with sem, host_sem(url):
            try:
                body = await asyncio.wait_for(fetch_bytes(item, timeout), timeout + 5.0)
            except Exception as exc:
                row["error"] = str(exc)[:200]
                return row
        if not body:
            row["error"] = "empty body"
            return row
        cap = FILE_MAX_BYTES if item.get("kind") == "file" else MAX_BYTES
        if len(body) > cap:
            row["bytes"] = len(body)
            row["error"] = f"over {cap // 1024 // 1024} MB cap"
            return row
        target = out / name
        target.write_bytes(body)
        row.update(ok=True, path=str(target), bytes=len(body))
        return row

    return list(await asyncio.gather(*(one(index, item) for index, item in enumerate(items))))


async def stealth_fetcher(stealth, item, timeout):
    referer = (item.get("referer") or "").strip()
    headers = {"Referer": referer} if referer else None
    session = item.get("session") or None
    if item.get("kind") == "file":
        return await stealth.download(item["url"], headers=headers, timeout=timeout, session=session)
    response = await stealth.get(item["url"], headers=headers, timeout=timeout, session=session)
    if response.status >= 400:
        raise RuntimeError(f"HTTP {response.status}")
    return await response.read()


async def curl_fetcher(session, item, timeout):
    headers = {"Accept": "*/*"}
    referer = (item.get("referer") or "").strip()
    if referer:
        headers["Referer"] = referer
    response = await session.get(item["url"], headers=headers, timeout=timeout)
    if response.status_code >= 400:
        raise RuntimeError(f"HTTP {response.status_code}")
    return response.content


async def run_live(items, out_dir, jobs, timeout, budget):
    from stealthii import Stealthii
    from curl_cffi import AsyncSession

    fast, bulk = split_lanes(items)

    async def fast_lane():
        if not fast:
            return []
        stealth = Stealthii()
        await stealth.start()
        try:
            async def go(item, limit):
                return await stealth_fetcher(stealth, item, limit)

            return await gather_batch(fast, out_dir, jobs, timeout, go)
        finally:
            await stealth.shutdown()

    async def bulk_lane():
        if not bulk:
            return []
        async with AsyncSession(impersonate="chrome") as session:
            async def go(item, limit):
                return await curl_fetcher(session, item, limit)

            return await gather_batch(bulk, out_dir, jobs, timeout, go)

    both = await asyncio.wait_for(asyncio.gather(fast_lane(), bulk_lane()), budget)
    merged = sorted(both[0] + both[1], key=lambda row: row.get("_pos", 0))
    for row in merged:
        row.pop("_pos", None)
    return merged


def main(argv=None):
    parser = argparse.ArgumentParser(description="Fetch a JSON list of page and file URLs concurrently through stealthii, with an HTTP/2 bulk lane.")
    parser.add_argument("urls", help="JSON file holding [{\"url\", \"kind\": page|file, \"name\"?, \"referer\"?, \"session\"?}]")
    parser.add_argument("--out-dir", default="pages", help="Directory for fetched bytes")
    parser.add_argument("--jobs", type=int, default=DEFAULT_JOBS, help="Concurrent fetches overall, capped at 20; per host capped at 6")
    parser.add_argument("--timeout", type=float, default=TIMEOUT, help="Seconds per URL")
    parser.add_argument("--budget", type=float, default=BUDGET, help="Seconds for the whole batch")
    args = parser.parse_args(argv)
    try:
        with open(args.urls, encoding="utf-8") as handle:
            items = json.load(handle)
    except (OSError, ValueError) as exc:
        print(json.dumps({"manifest": [], "error": f"bad urls file: {exc}"}))
        return 1
    if not isinstance(items, list) or not all(isinstance(item, dict) and item.get("url") for item in items):
        print(json.dumps({"manifest": [], "error": "urls file must hold [{\"url\", ...}]"}))
        return 1
    guard = load_guard()
    manifest, survivors = [], []
    for item in items:
        try:
            guard.check_url(item["url"])
        except ValueError as exc:
            manifest.append({"url": item["url"], "ok": False, "path": "", "bytes": 0, "error": str(exc)[:200]})
            continue
        survivors.append(item)
    if survivors:
        try:
            fetched = asyncio.run(run_live(survivors, args.out_dir, args.jobs, args.timeout, args.budget))
        except asyncio.TimeoutError:
            fetched = [{"url": item["url"], "ok": False, "path": "", "bytes": 0, "error": "round budget exceeded"} for item in survivors]
        by_url = {row["url"]: row for row in fetched}
        manifest.extend(by_url.get(item["url"], {"url": item["url"], "ok": False, "path": "", "bytes": 0, "error": "missing from batch"}) for item in survivors)
    print(json.dumps({"manifest": manifest, "error": ""}))
    return 0


if __name__ == "__main__":
    sys.exit(main())