"""Download pinned data artifacts without executing upstream code."""
import argparse
import hashlib
import json
import os
import subprocess
import fcntl
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
RAW = ROOT / "data/vulnerability/raw"
PROVENANCE = ROOT / "data/vulnerability/provenance"
REPOS = {
    "benchproctor": "TheAuditorTool/BenchProctor",
    "owasp-python": "OWASP-Benchmark/BenchmarkPython",
    "owasp-java": "OWASP-Benchmark/BenchmarkJava",
    "codexglue": "microsoft/CodeXGLUE",
    "primevul": "DLVulDet/PrimeVul",
    "diversevul": "wagner-group/diversevul",
    "bigvul": "ZeoVan/MSR_20_Code_vulnerability_CSV_Dataset",
    "devign": "epicosy/devign",
    "vulvalidate": "lzhan011/VulValidate",
    "sven": "eth-sri/sven",
}


def sha256(path):
    h = hashlib.sha256()
    with Path(path).open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def git_blob_sha1(path):
    path = Path(path)
    h = hashlib.sha1(f"blob {path.stat().st_size}\0".encode())
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def download(url, dest, max_mb=2048):
    dest = Path(dest)
    dest.parent.mkdir(parents=True, exist_ok=True)
    with dest.with_suffix(dest.suffix + ".lock").open("a") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX)
        return _download(url, dest, max_mb)


def _download(url, dest, max_mb=2048):
    dest = Path(dest)
    dest.parent.mkdir(parents=True, exist_ok=True)
    if dest.exists() and dest.stat().st_size:
        return
    temporary = dest.with_suffix(dest.suffix + ".part")
    # subprocess argument arrays: URLs and filenames are never interpreted by a shell.
    command = [
        "curl", "-sSL", "--fail", "--connect-timeout", "30", "--max-time", "1800",
        "--http1.1",
        "--continue-at", "-",
        "--speed-time", "60", "--speed-limit", "1024",
        "--max-filesize", str(max_mb * 1024 * 1024), "--output", str(temporary), url,
    ]
    if dest.suffix in (".json", ".jsonl", ".md") and not temporary.exists():
        command.insert(1, "--compressed")
    if os.environ.get("LAYA_DOWNLOAD_DIRECT") == "1":
        command[1:1] = ["--noproxy", "*"]
    result = subprocess.run(command, capture_output=True, text=True)
    if result.returncode:
        raise RuntimeError(result.stderr.strip())
    temporary.replace(dest)


def api(url, dest):
    download(url, dest, max_mb=100)
    return json.loads(Path(dest).read_text())


def pin(source, repo):
    folder = PROVENANCE / source
    info = api(f"https://api.github.com/repos/{repo}", folder / "repository.json")
    branch = info["default_branch"]
    commit = api(f"https://api.github.com/repos/{repo}/commits/{branch}", folder / "commit.json")
    revision = commit["sha"]
    tree = api(f"https://api.github.com/repos/{repo}/git/trees/{revision}?recursive=1", folder / "tree.json")
    if tree.get("truncated"):
        raise RuntimeError(f"{repo}: truncated tree, refusing an incomplete inventory")
    paths = {x["path"]: x for x in tree["tree"] if x["type"] == "blob"}
    record = {
        "source": source, "repository": f"https://github.com/{repo}", "revision": revision,
        "retrieved_at": datetime.now(timezone.utc).isoformat(),
        "repository_license": (info.get("license") or {}).get("spdx_id", "UNKNOWN"),
        "underlying_code_license": "See upstream source projects; repository license is not a blanket relicensing of third-party code.",
        "artifacts": [], "errors": [],
    }
    for path in paths:
        name = Path(path).name.lower()
        if "/" not in path and (name.startswith(("license", "copying", "notice")) or name in ("readme.md", "data_licence", "changelog.md")):
            try:
                download(f"https://raw.githubusercontent.com/{repo}/{revision}/{path}", folder / path)
            except RuntimeError as exc:
                record["errors"].append({"path": path, "error": str(exc)})
    return record, paths


def fetch_repo(source, bench_tier="normal"):
    repo = REPOS[source]
    try:
        record, paths = pin(source, repo)
    except Exception as exc:
        return {"source": source, "repository": f"https://github.com/{repo}", "errors": [str(exc)], "artifacts": []}
    revision = record["revision"]
    jobs = []
    if source == "benchproctor":
        # Download all languages and all complete parts of the published normal corpus.
        for path, entry in paths.items():
            if path.startswith(f"Benchmarks/{bench_tier}/") and path.endswith(".zip"):
                jobs.append((f"https://raw.githubusercontent.com/{repo}/{revision}/{path}", RAW / source / Path(path).name, path))
    elif source.startswith("owasp-"):
        jobs.append((f"https://codeload.github.com/{repo}/zip/{revision}", RAW / source / "repository.zip", "repository snapshot"))
    elif source == "codexglue":
        for path in paths:
            if path.startswith("Code-Code/Defect-detection/") and path.endswith((".jsonl", ".json", ".py", ".md")):
                jobs.append((f"https://raw.githubusercontent.com/{repo}/{revision}/{path}", RAW / source / path.removeprefix("Code-Code/Defect-detection/"), path))
    elif source == "devign":
        path = "data/raw/dataset.json"
        jobs.append((f"https://raw.githubusercontent.com/{repo}/{revision}/{path}", RAW / source / "dataset.json", path))
    elif source == "sven":
        for path in paths:
            if path.startswith("data_train_val/") and path.endswith(".jsonl"):
                name = path.removeprefix("data_train_val/").replace("/", "--")
                jobs.append((f"https://raw.githubusercontent.com/{repo}/{revision}/{path}", RAW / source / name, path))
    elif source == "vulvalidate":
        for path in paths:
            if path.startswith("combined_dataset/") and (path.endswith((".md", ".json")) or path.endswith(".jsonl.gz") and "paired" not in path):
                jobs.append((f"https://raw.githubusercontent.com/{repo}/{revision}/{path}", RAW / source / Path(path).name, path))
    # Real-world datasets often host data outside GitHub. Preserve exact official links.
    def get_artifact(job):
        url, dest, upstream_path = job
        try:
            download(url, dest)
            if upstream_path in paths and git_blob_sha1(dest) != paths[upstream_path]["sha"]:
                raise ValueError(f"Git blob integrity mismatch: {upstream_path}")
            artifact = {"url": url, "path": str(dest.relative_to(ROOT)), "upstream_path": upstream_path,
                        "bytes": dest.stat().st_size, "sha256": sha256(dest)}
            # Git LFS pointers are not datasets; record that distinction.
            if dest.stat().st_size < 1024 and dest.read_bytes().startswith(b"version https://git-lfs.github.com"):
                artifact["status"] = "lfs_pointer"
            else:
                artifact["status"] = "downloaded"
            record["artifacts"].append(artifact)
            print(source, dest.name, artifact["bytes"], artifact["status"], flush=True)
        except Exception as exc:
            record["errors"].append({"url": url, "error": str(exc)})
    with ThreadPoolExecutor(max_workers=2) as pool:
        list(pool.map(get_artifact, jobs))
    manifest_name = "manifest.json" if bench_tier == "normal" else f"manifest-{bench_tier}.json"
    (PROVENANCE / source / manifest_name).parent.mkdir(parents=True, exist_ok=True)
    (PROVENANCE / source / manifest_name).write_text(json.dumps(record, indent=2, ensure_ascii=False))
    return record


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--sources", nargs="+", default=list(REPOS), choices=list(REPOS))
    parser.add_argument("--workers", type=int, default=3)
    parser.add_argument("--external-only", action="store_true")
    parser.add_argument("--bench-tier", choices=["normal", "quicktest", "enterprise"], default="normal")
    args = parser.parse_args()
    PROVENANCE.mkdir(parents=True, exist_ok=True)
    if args.external_only:
        fetch_external()
        return
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = [pool.submit(fetch_repo, name, args.bench_tier) for name in args.sources]
        records = [future.result() for future in as_completed(futures)]
    by_source = {r["source"]: r for r in records}
    for path in PROVENANCE.glob("*/manifest.json"):
        r = json.loads(path.read_text())
        by_source.setdefault(r["source"], r)
    (PROVENANCE / "sources.json").write_text(json.dumps(sorted(by_source.values(), key=lambda r: r["source"]), indent=2, ensure_ascii=False))
    print("Source inventory saved; external data links may require a separate download.", flush=True)


def fetch_external():
    """Official DiverseVul download and a pinned, explicitly marked PrimeVul mirror."""
    import re
    from urllib.parse import urlencode
    def diverse():
        folder = PROVENANCE / "diversevul"
        original = "https://drive.google.com/uc?export=download&id=12IWKhmLhq7qn5B_iXgn5YerOQtkH-6RG"
        response = folder / "download-response.html"
        download(original, response, max_mb=1024)
        text = response.read_text()
        values = dict(re.findall(r'name="([^"]+)" value="([^"]*)"', text))
        if "uuid" not in values:
            raise RuntimeError("Official DiverseVul download confirmation format changed")
        return [("https://drive.usercontent.google.com/download?" + urlencode(values), RAW / "diversevul" / "dataset.jsonl")]
    def prime():
        folder = PROVENANCE / "primevul"
        info = api("https://huggingface.co/api/datasets/colin/PrimeVul", folder / "mirror-info.json")
        revision = info["sha"]
        return [(f"https://huggingface.co/datasets/colin/PrimeVul/resolve/{revision}/primevul_{split}.jsonl",
                 RAW / "primevul" / f"{split}.jsonl") for split in ("train", "valid", "test")]
    for source, get_jobs in (("diversevul", diverse), ("primevul", prime)):
        path = PROVENANCE / source / "external-manifest.json"
        record = {"source": source, "retrieved_at": datetime.now(timezone.utc).isoformat(), "artifacts": [], "errors": [],
                  "distribution": "official Google Drive" if source == "diversevul" else "third-party Hugging Face mirror colin/PrimeVul",
                  "original_repository": "https://github.com/" + REPOS[source]}
        try:
            for url, dest in get_jobs():
                download(url, dest)
                with dest.open("rb") as f:
                    prefix = f.read(100).lstrip()
                if not prefix.startswith((b"{", b"[")):
                    raise ValueError(f"Not JSON data: {dest}")
                record["artifacts"].append({"url": url, "path": str(dest.relative_to(ROOT)), "bytes": dest.stat().st_size, "sha256": sha256(dest)})
                print(source, dest.name, dest.stat().st_size, flush=True)
        except Exception as exc:
            record["errors"].append(str(exc))
            print(source, "external error", str(exc), flush=True)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(record, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()