#!/usr/bin/env python3
"""Resumable chunked model bundle downloader using only Python 3.9+ stdlib."""
import argparse, hashlib, json, os, re, stat, sys, time, urllib.error, urllib.parse, urllib.request
from pathlib import Path

PROD_ORIGIN = "https://relay-sports-trained-v01.coursework-site.workers.dev"
DEV_ORIGIN = "http://127.0.0.1:8792"
EXPECTED_BUNDLE = ("Relay-Sports-8B-Trained-Preview-0.1.zip", 5113783892, "c8eae0218480291a48ab23b1c208b3b0b03756d799f0520fd3c07042fa9d6467")
UA = "Mozilla/5.0 RelayDownloader/0.1"
BUF_SIZE, MAX_MANIFEST, MAX_CHUNKS, MAX_CHUNK_BYTES = 256 * 1024, 128 * 1024, 1000, 25 * 1024 * 1024
HEX64_RE = re.compile(r"^[0-9a-f]{64}$")
SAFE_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,120}\.zip$")
CHUNK_PATH_RE = re.compile(r"^/chunks/[0-9]{4}\.bin$")


def die(msg: str) -> None:
    print(f"Error: {msg}", file=sys.stderr)
    sys.exit(1)


class SameOriginRedirectHandler(urllib.request.HTTPRedirectHandler):
    def __init__(self, allowed_origin: str):
        self.scheme, self.netloc = urllib.parse.urlsplit(allowed_origin)[:2]

    def redirect_request(self, req, fp, code, msg, headers, newurl):
        u = urllib.parse.urlsplit(newurl)
        if (u.scheme, u.netloc) != (self.scheme, self.netloc):
            raise urllib.error.URLError(f"Cross-origin redirect rejected: {newurl}")
        return super().redirect_request(req, fp, code, msg, headers, newurl)


def open_url(origin: str, path: str):
    opener = urllib.request.build_opener(SameOriginRedirectHandler(origin))
    req = urllib.request.Request(f"{origin}{path}", headers={"User-Agent": UA})
    return opener.open(req, timeout=60)


def fetch_manifest(origin: str) -> dict:
    try:
        with open_url(origin, "/download-manifest.json") as resp:
            raw = resp.read(MAX_MANIFEST + 1)
    except Exception as e:
        die(f"Failed to download manifest: {e}")
    if len(raw) > MAX_MANIFEST:
        die("Manifest exceeds 128 KiB limit")
    try:
        m = json.loads(raw.decode("utf-8"))
    except Exception as e:
        die(f"Invalid manifest JSON: {e}")
    if not isinstance(m, dict):
        die("Manifest root must be an object")
    name, size, sha, csize, chunks = (
        m.get("name"), m.get("size"), m.get("sha256"), m.get("chunkSize"), m.get("chunks")
    )
    if (name, size, sha) != EXPECTED_BUNDLE:
        die("Manifest does not match this pinned Relay release")
    if not isinstance(name, str) or Path(name).name != name or not SAFE_NAME_RE.match(name):
        die(f"Unsafe or invalid bundle filename: {name!r}")
    if not isinstance(size, int) or size <= 0 or not isinstance(csize, int) or not (0 < csize <= MAX_CHUNK_BYTES):
        die("Invalid bundle size or chunkSize")
    if not isinstance(sha, str) or not HEX64_RE.match(sha):
        die("Invalid bundle sha256")
    if not isinstance(chunks, list) or not (1 <= len(chunks) <= MAX_CHUNKS):
        die("Invalid chunks list count")
    total = 0
    for i, c in enumerate(chunks):
        if not isinstance(c, dict):
            die(f"Chunk {i} is not an object")
        cp, cs, ch = c.get("path"), c.get("size"), c.get("sha256")
        if not isinstance(cp, str) or not CHUNK_PATH_RE.match(cp):
            die(f"Chunk {i} has invalid path: {cp!r}")
        if not isinstance(cs, int) or not (0 < cs <= MAX_CHUNK_BYTES):
            die(f"Chunk {i} has invalid size: {cs!r}")
        if not isinstance(ch, str) or not HEX64_RE.match(ch):
            die(f"Chunk {i} has invalid sha256")
        total += cs
    if total != size:
        die(f"Chunk size sum ({total}) does not match manifest size ({size})")
    return m


def verify_file_sha256(path: Path, expected_size: int, expected_sha: str) -> bool:
    h = hashlib.sha256()
    total = 0
    with path.open("rb") as f:
        while chunk := f.read(BUF_SIZE):
            h.update(chunk)
            total += len(chunk)
    return total == expected_size and h.hexdigest() == expected_sha


def scan_existing_part(f, chunks: list) -> int:
    valid_offset, verified = 0, 0
    f.seek(0)
    for c in chunks:
        rem, h = c["size"], hashlib.sha256()
        while rem > 0:
            buf = f.read(min(BUF_SIZE, rem))
            if not buf:
                break
            h.update(buf)
            rem -= len(buf)
        if rem == 0 and h.hexdigest() == c["sha256"]:
            valid_offset += c["size"]
            verified += 1
        else:
            break
    f.seek(0, os.SEEK_END)
    if f.tell() > valid_offset:
        print(f"Truncating incomplete/corrupt tail at byte {valid_offset}")
        f.truncate(valid_offset)
    return verified


def download_chunk(origin: str, c: dict, f, offset: int) -> None:
    for attempt in range(1, 4):
        f.seek(offset)
        f.truncate(offset)
        h, got = hashlib.sha256(), 0
        try:
            with open_url(origin, c["path"]) as resp:
                while got < c["size"]:
                    buf = resp.read(min(BUF_SIZE, c["size"] - got))
                    if not buf:
                        break
                    f.write(buf)
                    h.update(buf)
                    got += len(buf)
                extra = resp.read(1)
            if got == c["size"] and not extra and h.hexdigest() == c["sha256"]:
                f.flush()
                return
            err = f"size/hash mismatch (got {got}/{c['size']} bytes)"
        except Exception as e:
            err = str(e)
        f.seek(offset)
        f.truncate(offset)
        if attempt < 3:
            print(f"  Retry {attempt}/3 for {c['path']} after error: {err}", file=sys.stderr)
            time.sleep(attempt)
        else:
            die(f"Failed chunk {c['path']} after 3 attempts: {err}")


def main() -> None:
    ap = argparse.ArgumentParser(description="Resumable public model bundle downloader.")
    ap.add_argument("--origin", default=PROD_ORIGIN, help="Fixed public origin (or local dev origin)")
    args = ap.parse_args()
    if args.origin not in (PROD_ORIGIN, DEV_ORIGIN):
        die(f"Origin not permitted. Allowed: {PROD_ORIGIN} or {DEV_ORIGIN}")

    manifest = fetch_manifest(args.origin)
    cwd = Path.cwd()
    final_path = cwd / manifest["name"]
    part_path = cwd / f"{manifest['name']}.part"
    lock_path = cwd / f"{manifest['name']}.lock"

    for p in (final_path, part_path, lock_path):
        if p.is_symlink():
            die(f"Refusing symlink target: {p}")
    if final_path.exists():
        die(f"Destination file already exists (refusing to overwrite): {final_path}")

    try:
        lock_fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
        os.write(lock_fd, f"{os.getpid()}\n".encode())
        os.close(lock_fd)
    except FileExistsError:
        die(f"Lock file exists ({lock_path.name}); another download may be running. If a previous run crashed, confirm no downloader is running before removing that lock file and retrying.")

    try:
        chunks = manifest["chunks"]
        fd = os.open(part_path, os.O_RDWR | os.O_CREAT | getattr(os, "O_NOFOLLOW", 0), 0o600)
        with os.fdopen(fd, "r+b") as f:
            if not stat.S_ISREG(os.fstat(f.fileno()).st_mode) or os.fstat(f.fileno()).st_nlink != 1:
                die("Partial download must be a regular file with one link")
            start_idx = scan_existing_part(f, chunks)
            offset = sum(c["size"] for c in chunks[:start_idx])
            if start_idx:
                print(f"Resuming {start_idx}/{len(chunks)} verified pieces ({offset}/{manifest['size']} bytes)")
            for i in range(start_idx, len(chunks)):
                c = chunks[i]
                download_chunk(args.origin, c, f, offset)
                offset += c["size"]
                print(f"[{i + 1}/{len(chunks)}] {c['path']} OK ({offset}/{manifest['size']} bytes)")
        if not verify_file_sha256(part_path, manifest["size"], manifest["sha256"]):
            die("Final bundle SHA-256 or size verification failed; keeping .part file")
        os.link(part_path, final_path)
        part_path.unlink()
        print(f"Download complete and verified: {final_path}")
    except OSError as error:
        die(f"Download failed; progress retained: {error}")
    except KeyboardInterrupt:
        print("\nInterrupted; keeping .part file for resume.", file=sys.stderr)
        sys.exit(1)
    finally:
        try:
            if not lock_path.is_symlink() and lock_path.exists():
                lock_path.unlink()
        except OSError:
            pass


if __name__ == "__main__":
    main()