#!/usr/bin/env python3
"""Download a paid datai order without keeping credentials or sessions on disk.

Uses only the Python standard library. Video files stream to .part files and
resume with HTTP Range. Server-verified per-part SHA-256 checksums are checked
before an atomic rename. A contributor-supplied whole-file checksum is never
misrepresented as independently verified.
"""
from __future__ import annotations

import argparse
import getpass
import hashlib
import http.client
import http.cookiejar
import json
import os
from pathlib import Path
import re
import sys
import tempfile
import time
from urllib.error import HTTPError, URLError
from urllib.parse import quote, urljoin, urlsplit
from urllib.request import Request, build_opener, HTTPCookieProcessor, HTTPRedirectHandler

READ_SIZE = 1024 * 1024
SHA256 = re.compile(r"^[a-fA-F0-9]{64}$")
IDENTIFIER = re.compile(r"^[a-zA-Z0-9_-]{1,128}$")


class DownloadError(Exception):
    """An actionable error without exposing authentication material."""


class NetworkError(DownloadError):
    """An interruption that can safely retry using the current .part file."""


def parsed_origin(parsed) -> str:
    host = (parsed.hostname or "").lower()
    if ":" in host:
        host = f"[{host}]"
    port = parsed.port
    if port and (parsed.scheme, port) not in (("https", 443), ("http", 80)):
        host += f":{port}"
    return f"{parsed.scheme}://{host}"


def origin_for(url: str, allow_local_http: bool = False) -> str:
    parsed = urlsplit(url)
    if parsed.username or parsed.password or parsed.query or parsed.fragment:
        raise DownloadError("Use the website origin without credentials, a query or a fragment.")
    if parsed.path not in ("", "/") or not parsed.hostname:
        raise DownloadError("Use the website origin, for example https://datai.nz.")
    local = parsed.hostname in ("localhost", "127.0.0.1", "::1")
    if parsed.scheme != "https" and not (allow_local_http and local and parsed.scheme == "http"):
        raise DownloadError("HTTPS is required. Local HTTP is available only with --allow-local-http.")
    return parsed_origin(parsed)


def same_origin_url(base: str, target: str) -> str:
    absolute = urljoin(base + "/", target)
    parsed = urlsplit(absolute)
    if parsed.username or parsed.password or parsed_origin(parsed) != base:
        raise DownloadError("The dataset contains a download URL outside this website.")
    if parsed.fragment:
        raise DownloadError("The dataset contains an invalid download URL.")
    return absolute


class SameOriginRedirects(HTTPRedirectHandler):
    def __init__(self, base: str):
        self.base = base

    def redirect_request(self, req, fp, code, msg, headers, newurl):
        safe = same_origin_url(self.base, newurl)
        return super().redirect_request(req, fp, code, msg, headers, safe)


class Client:
    def __init__(self, base: str):
        self.base = base
        # Deliberately not a FileCookieJar. Nothing is saved by this jar.
        self.cookies = http.cookiejar.CookieJar()
        self.opener = build_opener(HTTPCookieProcessor(self.cookies), SameOriginRedirects(base))

    def open(self, path: str, *, payload=None, headers=None):
        url = same_origin_url(self.base, path)
        request_headers = {"User-Agent": "datai-order-downloader/1.0", **(headers or {})}
        body = None
        if payload is not None:
            body = json.dumps(payload).encode("utf-8")
            request_headers.update({"Content-Type": "application/json", "Origin": self.base})
        request = Request(url, data=body, headers=request_headers)
        try:
            return self.opener.open(request, timeout=60)
        except HTTPError as error:
            message = f"The server returned HTTP {error.code}."
            try:
                details = json.loads(error.read(16384).decode("utf-8"))
                if isinstance(details.get("error"), str):
                    message += " " + details["error"][:500]
            except (ValueError, UnicodeError, AttributeError):
                pass
            raise (NetworkError if error.code >= 500 or error.code == 429 else DownloadError)(message) from None
        except URLError as error:
            raise NetworkError(f"Could not reach the website: {error.reason}") from None

    def login(self, email: str, password: str):
        with self.open("/api/auth/login", payload={"email": email, "password": password}) as response:
            result = json.load(response)
        if result.get("ok") is not True or not result.get("user"):
            raise DownloadError("Sign-in could not be confirmed.")

    def logout(self):
        if not list(self.cookies):
            return
        try:
            with self.open("/api/auth/logout", payload={}) as response:
                response.read(1024)
        except (DownloadError, OSError):
            pass
        self.cookies.clear()


def destination_for(root: Path, entry: dict) -> Path:
    video_id = entry.get("video_id", "")
    if not isinstance(video_id, str) or not IDENTIFIER.fullmatch(video_id):
        raise DownloadError("The manifest contains an invalid video identifier.")
    original = str(entry.get("filename") or "recording").replace("\\", "/").rsplit("/", 1)[-1]
    cleaned = re.sub(r'[<>:"/\\|?*\x00-\x1f]', "_", original).strip(" .")[:100] or "recording"
    extension = ".webm" if entry.get("content_type") == "video/webm" else ".mp4"
    if not cleaned.lower().endswith(extension):
        cleaned += extension
    target = root / f"{video_id}__{cleaned}"
    partial = target.with_name(target.name + ".part")
    if target.is_symlink() or partial.is_symlink():
        raise DownloadError("Refusing to write through an existing symbolic link.")
    if target.resolve().parent != root.resolve() or partial.resolve().parent != root.resolve():
        raise DownloadError("The manifest contains an unsafe filename.")
    return target


def verified_parts(entry: dict) -> list[dict]:
    metadata = entry.get("metadata") or {}
    if metadata.get("upload_integrity") != "server-verified-size-and-part-sha256":
        return []
    parts = metadata.get("part_checksums")
    if not isinstance(parts, list) or not parts:
        raise DownloadError("The server-verified part checksum manifest is incomplete.")
    for index, part in enumerate(parts, 1):
        if (not isinstance(part, dict) or part.get("part_number") != index
                or type(part.get("size_bytes")) is not int or part["size_bytes"] <= 0
                or not SHA256.fullmatch(str(part.get("sha256", "")))):
            raise DownloadError("The server-verified part checksum manifest is invalid.")
    return parts


def expected_size(entry: dict) -> int | None:
    parts = verified_parts(entry)
    if parts:
        return sum(part["size_bytes"] for part in parts)
    value = entry.get("size_bytes")
    return value if isinstance(value, int) and value > 0 else None


def verify_file(path: Path, entry: dict) -> dict:
    size = expected_size(entry)
    if size is not None and path.stat().st_size != size:
        raise DownloadError("The downloaded size does not match the purchased snapshot.")
    parts = verified_parts(entry)
    whole = hashlib.sha256()
    with path.open("rb") as stream:
        if parts:
            for part in parts:
                remaining = part["size_bytes"]
                current = hashlib.sha256()
                while remaining:
                    chunk = stream.read(min(READ_SIZE, remaining))
                    if not chunk:
                        raise DownloadError("The recording is incomplete.")
                    remaining -= len(chunk)
                    current.update(chunk)
                    whole.update(chunk)
                if current.hexdigest().lower() != part["sha256"].lower():
                    raise DownloadError(f"SHA-256 mismatch in part {part['part_number']}. Retry with --restart.")
            if stream.read(1):
                raise DownloadError("The recording contains more bytes than the purchased snapshot.")
        else:
            for chunk in iter(lambda: stream.read(READ_SIZE), b""):
                whole.update(chunk)
    local_hash = whole.hexdigest()
    metadata = entry.get("metadata") or {}
    whole_verified = metadata.get("sha256_verified") is True
    if whole_verified:
        trusted = entry.get("sha256")
        if not isinstance(trusted, str) or not SHA256.fullmatch(trusted):
            raise DownloadError("The independently verified SHA-256 value is invalid.")
        if local_hash.lower() != trusted.lower():
            raise DownloadError("The recording does not match its independently verified SHA-256 checksum.")
    return {"local_sha256": local_hash, "size_bytes": path.stat().st_size,
            "server_parts_verified": bool(parts), "server_whole_sha256_verified": whole_verified}


def atomic_write(path: Path, data: bytes):
    if path.is_symlink():
        raise DownloadError("Refusing to replace an existing symbolic link.")
    descriptor, temporary = tempfile.mkstemp(prefix=".datai-", dir=path.parent)
    try:
        with os.fdopen(descriptor, "wb") as stream:
            stream.write(data)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def save_response(client: Client, route: str, path: Path):
    """Stream a potentially large manifest to disk before publishing it."""
    if path.is_symlink():
        raise DownloadError("Refusing to replace an existing symbolic link.")
    descriptor, temporary = tempfile.mkstemp(prefix=".datai-", dir=path.parent)
    try:
        with os.fdopen(descriptor, "wb") as output, client.open(route) as response:
            for chunk in iter(lambda: response.read(READ_SIZE), b""):
                output.write(chunk)
            output.flush()
            os.fsync(output.fileno())
        os.replace(temporary, path)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def manifest_entries(path: Path, order_id: str, base: str, output: Path):
    """Read one bounded JSONL record at a time, not the complete order."""
    with path.open("r", encoding="utf-8") as source:
        while True:
            line = source.readline(4 * 1024 * 1024 + 1)
            if not line:
                return
            if len(line) > 4 * 1024 * 1024:
                raise DownloadError("A manifest record exceeds the supported size.")
            if not line.strip():
                continue
            try:
                entry = json.loads(line)
            except ValueError:
                raise DownloadError("The server returned an invalid JSONL record.") from None
            if not isinstance(entry, dict) or entry.get("order_id") != order_id or entry.get("schema_version") != "datai.video.v1":
                raise DownloadError("The manifest does not match this order or schema.")
            destination_for(output, entry)
            download_url = same_origin_url(base, entry.get("download_url", ""))
            if urlsplit(download_url).path != f"/api/videos/{entry['video_id']}/content":
                raise DownloadError("The manifest contains an unexpected recording URL.")
            if entry.get("analysis_url"):
                analysis_url = same_origin_url(base, entry["analysis_url"])
                if urlsplit(analysis_url).path != f"/api/orders/{order_id}/analysis/{entry['video_id']}" or urlsplit(analysis_url).query:
                    raise DownloadError("The manifest contains an unexpected analysis URL.")
            verified_parts(entry)
            yield entry


def download_analysis(client: Client, root: Path, entry: dict) -> Path | None:
    """Save the exact frozen private analysis, rather than a regenerated label set."""
    if not entry.get("analysis_url"):
        return None
    video_id = entry.get("video_id", "")
    if not isinstance(video_id, str) or not IDENTIFIER.fullmatch(video_id):
        raise DownloadError("The analysis recording identifier is invalid.")
    summary = (entry.get("metadata") or {}).get("analysis") or {}
    expected_hash, expected_bytes = summary.get("sha256"), summary.get("size_bytes")
    if not isinstance(expected_hash, str) or not SHA256.fullmatch(expected_hash) or type(expected_bytes) is not int or not 0 < expected_bytes <= 32 * 1024 * 1024:
        raise DownloadError("The frozen analysis checksum or byte size is invalid.")
    with client.open(entry["analysis_url"]) as response:
        payload = response.read(expected_bytes + 1)
        if len(payload) != expected_bytes or hashlib.sha256(payload).hexdigest().lower() != expected_hash.lower():
            raise DownloadError("The analysis package differs from the purchased snapshot.")
    try:
        analysis = json.loads(payload)
    except (ValueError, UnicodeError):
        raise DownloadError("The analysis package contains invalid JSON.") from None
    if not isinstance(analysis, dict) or analysis.get("schema_version") != "datai.video-analysis.v1" or analysis.get("recording_id") != video_id or (analysis.get("source") or {}).get("sha256") != entry.get("sha256"):
        raise DownloadError("The analysis package does not match this recording.")
    target = root / f"{video_id}.analysis.json"
    if target.is_symlink() or target.resolve().parent != root.resolve():
        raise DownloadError("Refusing an unsafe analysis destination.")
    atomic_write(target, payload)
    return target


def download_video(client: Client, root: Path, entry: dict, restart: bool = False) -> dict:
    target = destination_for(root, entry)
    partial = target.with_name(target.name + ".part")
    wanted = expected_size(entry)
    if target.exists() and not restart:
        verified = verify_file(target, entry)
        print(f"Already complete: {target.name}", flush=True)
        return {"video_id": entry["video_id"], "file": target.name, **verified}
    if restart and partial.exists():
        partial.unlink()
    offset = partial.stat().st_size if partial.exists() else 0
    if wanted is not None and offset > wanted:
        raise DownloadError("The existing partial file is too large. Retry with --restart.")
    if wanted is None or offset < wanted:
        headers = {"Range": f"bytes={offset}-"} if offset else {}
        with client.open(entry["download_url"], headers=headers) as response:
            if offset and response.status == 206:
                match = re.fullmatch(r"bytes (\d+)-(\d+)/(\d+)", response.headers.get("Content-Range", ""))
                if not match or int(match[1]) != offset or (wanted is not None and int(match[3]) != wanted):
                    raise DownloadError("The server returned an invalid resume range.")
            elif offset and response.status == 200:
                offset = 0  # The server ignored Range: replace our own partial file.
            elif response.status != 200:
                raise DownloadError(f"Unexpected download response HTTP {response.status}.")
            if wanted is None:
                length = response.headers.get("Content-Length")
                if length and length.isdigit():
                    wanted = offset + int(length)
            mode = "ab" if offset else "wb"
            with partial.open(mode) as output:
                received = offset
                last_report = 0.0
                for chunk in iter(lambda: response.read(READ_SIZE), b""):
                    received += len(chunk)
                    if wanted is not None and received > wanted:
                        raise DownloadError("The server sent more bytes than expected.")
                    output.write(chunk)
                    if time.monotonic() - last_report > 3:
                        suffix = f" / {wanted / 1048576:.1f} MiB" if wanted else ""
                        print(f"  {received / 1048576:.1f} MiB{suffix}", flush=True)
                        last_report = time.monotonic()
                output.flush()
                os.fsync(output.fileno())
    if wanted is not None and partial.stat().st_size < wanted:
        raise NetworkError("The transfer ended early. The partial file can be resumed.")
    verified = verify_file(partial, entry)
    os.replace(partial, target)
    label = "server part checksums verified" if verified["server_parts_verified"] else "local SHA-256 recorded"
    print(f"Complete: {target.name} ({label})", flush=True)
    return {"video_id": entry["video_id"], "file": target.name, **verified}


def parser():
    result = argparse.ArgumentParser(description=__doc__)
    result.add_argument("--base-url", default="https://datai.nz")
    result.add_argument("--email", required=True, help="Your buyer account email; the password is prompted securely.")
    result.add_argument("--order", required=True, help="The ID of a paid order.")
    result.add_argument("--output", type=Path, help="Destination folder; defaults to datai-<order-id>.")
    result.add_argument("--restart", action="store_true", help="Replace partial downloads and redownload existing files.")
    result.add_argument("--manifest-only", action="store_true", help="Download JSONL, CSV and license without recordings.")
    result.add_argument("--allow-local-http", action="store_true", help="Development only: allow loopback HTTP.")
    return result


def main(argv=None):
    args = parser().parse_args(argv)
    client = None
    try:
        if not IDENTIFIER.fullmatch(args.order):
            raise DownloadError("The order identifier is invalid.")
        base = origin_for(args.base_url, args.allow_local_http)
        output = (args.output or Path(f"datai-{args.order}")).resolve()
        output.mkdir(parents=True, exist_ok=True)
        client = Client(base)
        password = getpass.getpass("datai password: ")
        client.login(args.email, password)
        del password
        order_path = f"/api/orders/{quote(args.order, safe='')}/export"
        manifest_path = output / f"order-{args.order}.jsonl"
        save_response(client, order_path + ".jsonl", manifest_path)
        count = sum(1 for _ in manifest_entries(manifest_path, args.order, base, output))
        if not count:
            raise DownloadError("The purchased order has no recordings.")
        save_response(client, order_path + ".csv", output / f"order-{args.order}.csv")
        with client.open("/api/marketplace/license") as response:
            license_info = json.load(response)
        if license_info.get("id") != "non-exclusive-ai-training-v1" or not isinstance(license_info.get('text'), str):
            raise DownloadError("The dataset license version could not be confirmed.")
        atomic_write(output / "license.json", json.dumps(license_info, indent=2).encode("utf-8"))
        atomic_write(output / "LICENSE.txt", license_info["text"].encode("utf-8"))
        print(f"Saved {count} recording records and the dataset license to {output}", flush=True)
        if args.manifest_only:
            return 0
        reports_path = output / "verified-local.jsonl"
        if reports_path.is_symlink():
            raise DownloadError("Refusing to replace an existing symbolic link.")
        completed = 0
        with reports_path.open("w", encoding="utf-8") as reports:
          for number, entry in enumerate(manifest_entries(manifest_path, args.order, base, output), 1):
            print(f"[{number}/{count}] {destination_for(output, entry).name}", flush=True)
            for attempt in range(3):
                try:
                    report = download_video(client, output, entry, restart=args.restart and attempt == 0)
                    analysis_path = download_analysis(client, output, entry)
                    if analysis_path is not None:
                        report["analysis_file"] = analysis_path.name
                    reports.write(json.dumps(report) + "\n")
                    reports.flush()
                    completed += 1
                    break
                except (NetworkError, OSError, URLError, TimeoutError, http.client.HTTPException) as error:
                    if attempt == 2:
                        raise DownloadError(f"The download was interrupted. Rerun the same command to resume. {error}") from None
                    time.sleep(attempt + 1)
        print(f"Finished: {completed} recordings. Cookies and password were not saved.", flush=True)
        return 0
    except (DownloadError, OSError, ValueError, http.client.HTTPException) as error:
        print(f"Error: {error}", file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        print("\nInterrupted. Rerun the same command to resume partial files.", file=sys.stderr)
        return 130
    finally:
        if client is not None:
            client.logout()


if __name__ == "__main__":
    raise SystemExit(main())