from __future__ import annotations

import argparse
from pathlib import Path

import sys

REPO_ROOT = Path(__file__).resolve().parents[2]
SRC_DIR = REPO_ROOT / "src"
if str(SRC_DIR) not in sys.path:
    sys.path.insert(0, str(SRC_DIR))

from data.io_utils import download_file, ensure_dir, load_yaml, verify_checksum  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Download Visual Genome raw artifacts.")
    parser.add_argument(
        "--config",
        type=Path,
        default=REPO_ROOT / "configs" / "dataset.yaml",
        help="Path to dataset config yaml.",
    )
    parser.add_argument(
        "--force-download",
        action="store_true",
        help="Force re-download even if files already exist.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    cfg = load_yaml(args.config)
    dataset_cfg = cfg["dataset"]
    download_cfg = cfg.get("download", {})

    if not download_cfg.get("enabled", True):
        print("Download disabled in config. Nothing to do.")
        return

    raw_root = REPO_ROOT / dataset_cfg["root_dir"]
    ensure_dir(raw_root)

    timeout_seconds = int(download_cfg.get("timeout_seconds", 120))
    retries = int(download_cfg.get("retries", 3))
    chunk_size = int(download_cfg.get("chunk_size_bytes", 1024 * 1024))
    force_download = args.force_download or bool(download_cfg.get("force_download", False))

    artifacts = download_cfg.get("artifacts", [])
    if not artifacts:
        raise RuntimeError("No download artifacts configured.")

    for artifact in artifacts:
        filename = artifact["filename"]
        url = artifact["url"]
        expected_sha256 = artifact.get("sha256")
        target = raw_root / filename

        if target.exists() and not force_download:
            if verify_checksum(target, expected_sha256):
                print(f"[skip] {filename} already exists and checksum passed.")
                continue
            print(f"[warn] {filename} exists but checksum mismatch. Re-downloading.")

        print(f"[download] {filename} <- {url}")
        download_file(
            url=url,
            destination=target,
            timeout_seconds=timeout_seconds,
            retries=retries,
            chunk_size=chunk_size,
        )

        if not verify_checksum(target, expected_sha256):
            raise RuntimeError(f"Checksum validation failed for {target}")
        print(f"[ok] {filename}")

    print("All configured Visual Genome artifacts are downloaded.")


if __name__ == "__main__":
    main()