"""Fetch a small real DeepFly3D sample from Harvard Dataverse.

The DeepFly3D experiment zips are ~750 MB each, but store the 7-camera JPEG frames
*uncompressed* — so a partial HTTP range download of the first few MB is enough to
extract real frames. This builds a 7-camera montage for inspection.

Note: these are markerless real images at native resolution (960×480 per camera);
the `fast` engine (color-marker triangulation on a synthetic rig) cannot reconstruct
3D from them — that needs the real `deepfly3d` engine. This loader brings the real
data in for viewing / for the deep pipeline.

Usage:
  python -m scripts.fetch_deepfly3d [datafile_id] [megabytes] [out_dir]
Dataverse DeepFly3D (aDN>CsChrimson): doi:10.7910/DVN/S4L4KX
"""
from __future__ import annotations

import io
import re
import struct
import sys
import urllib.request
from collections import defaultdict
from pathlib import Path

import numpy as np
import tifffile
from PIL import Image

from core.io import APP_TMP_DIR, register_protected

DATAVERSE = "https://dataverse.harvard.edu/api/access/datafile/"


def _range_download(datafile_id: int, mb: int) -> bytes:
    url = f"{DATAVERSE}{datafile_id}"
    req = urllib.request.Request(url, headers={
        "Range": f"bytes=0-{mb * 1_000_000}",
        "User-Agent": "Mozilla/5.0 (imaging-plaza runnable-example)",
    })
    with urllib.request.urlopen(req) as r:
        return r.read()


def _extract_jpgs(data: bytes) -> dict:
    imgs, i = {}, 0
    while True:
        j = data.find(b"PK\x03\x04", i)
        if j < 0 or j + 30 > len(data):
            break
        _, _, _, comp, _, _, _, csize, _, fnl, efl = struct.unpack("<IHHHHHIIIHH", data[j:j + 30])
        name = data[j + 30:j + 30 + fnl].decode("latin1", "ignore")
        start = j + 30 + fnl + efl
        if comp == 0 and 0 < csize and start + csize <= len(data):
            if name.endswith(".jpg"):
                imgs[name] = data[start:start + csize]
            i = start + csize
        else:
            i = j + 4
    return imgs


def main() -> int:
    datafile_id = int(sys.argv[1]) if len(sys.argv) > 1 else 3443269
    mb = int(sys.argv[2]) if len(sys.argv) > 2 else 30
    out_dir = Path(sys.argv[3]) if len(sys.argv) > 3 else APP_TMP_DIR

    print(f"range-downloading first {mb} MB of Dataverse datafile {datafile_id}…")
    imgs = _extract_jpgs(_range_download(datafile_id, mb))
    bycam = defaultdict(list)
    for n in imgs:
        m = re.search(r"camera_(\d)_img_(\d+)", n)
        if m:
            bycam[int(m.group(1))].append((int(m.group(2)), n))
    print(f"extracted {len(imgs)} real frames; cameras present: {sorted(bycam)}")

    tiles, native = [], None
    for cam in range(7):
        if bycam.get(cam):
            im = Image.open(io.BytesIO(imgs[sorted(bycam[cam])[0][1]]))
            native = im.size
            tiles.append(np.array(im.convert("L").resize((200, 200))))
        else:
            tiles.append(np.zeros((200, 200), np.uint8))
    tiles.append(np.zeros((200, 200), np.uint8))
    grid = np.block([[tiles[0], tiles[1], tiles[2], tiles[3]],
                     [tiles[4], tiles[5], tiles[6], tiles[7]]])

    out_dir.mkdir(parents=True, exist_ok=True)
    montage = out_dir / "deepfly3d_real.tif"
    tifffile.imwrite(montage, grid)
    register_protected(montage)
    png = out_dir / "deepfly3d_real.png"
    Image.fromarray(grid).save(png)
    register_protected(png)
    print(f"wrote {montage}  (7 real cameras, native {native})")
    return 0


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