---

#!/usr/bin/env python3
"""Restore the first Camera 08 YOLO labels from the companion Pastebin TXT.

Usage:
    python3 restore_carried_box_labels.py /path/carried_boxes_parts1-2_labels.txt
    python3 restore_carried_box_labels.py /path/labels.txt --output /path/to/labels

By default, creates a new "labels" directory beside the input TXT.
"""

import argparse
import hashlib
import json
import math
import os
import shutil
import tempfile
from pathlib import Path


FIRST_FRAME = 1
LAST_FRAME = 1536
FORMAT = "carried_box_yolo_labels_v1"


def decode_file(input_path):
    lines = [line for line in input_path.read_text(encoding="utf-8-sig").splitlines() if line.strip()]
    if not lines:
        raise ValueError("Input TXT is empty")
    header = json.loads(lines[0])
    if not isinstance(header, dict) or header.get("format") != FORMAT:
        raise ValueError("Wrong file format: use the raw labels TXT, not a Pastebin HTML page")
    if (header.get("first_frame"), header.get("last_frame")) != (FIRST_FRAME, LAST_FRAME):
        raise ValueError("This script expects frames 1 through 1536")
    if len(lines) != LAST_FRAME - FIRST_FRAME + 2:
        raise ValueError(f"Expected 1536 frame records, found {len(lines) - 1}")

    files = []
    digest = hashlib.sha256()
    positives = boxes = 0
    epsilon = 1e-7  # YOLO values rounded to 8 decimals can exceed an edge by a few nanounits
    for expected_frame, line in enumerate(lines[1:], FIRST_FRAME):
        record = json.loads(line)
        if not isinstance(record, dict) or record.get("frame") != expected_frame:
            raise ValueError(f"Missing, duplicate or out-of-order frame near {expected_frame}")
        label_lines = record.get("labels")
        if not isinstance(label_lines, list) or any(not isinstance(item, str) for item in label_lines):
            raise ValueError(f"Invalid labels for frame {expected_frame}")
        for label in label_lines:
            fields = label.split()
            if len(fields) != 5 or fields[0] != "0":
                raise ValueError(f"Invalid YOLO line at frame {expected_frame}: {label!r}")
            cx, cy, width, height = map(float, fields[1:])
            if not all(math.isfinite(value) for value in (cx, cy, width, height)):
                raise ValueError(f"Non-finite coordinate at frame {expected_frame}")
            if not (
                0 < width <= 1
                and 0 < height <= 1
                and -epsilon <= cx - width / 2 < cx + width / 2 <= 1 + epsilon
                and -epsilon <= cy - height / 2 < cy + height / 2 <= 1 + epsilon
            ):
                raise ValueError(f"Out-of-range box at frame {expected_frame}")
            boxes += 1
        if label_lines:
            positives += 1
        filename = f"frame_{expected_frame:06d}.txt"
        content = ("\n".join(label_lines) + ("\n" if label_lines else "")).encode("utf-8")
        digest.update(filename.encode("ascii") + b"\0" + content + b"\0")
        files.append((filename, content))

    expected = (header.get("files"), header.get("positives"), header.get("boxes"))
    actual = (len(files), positives, boxes)
    if expected != actual:
        raise ValueError(f"Count mismatch: expected {expected}, found {actual}")
    if digest.hexdigest() != header.get("sha256"):
        raise ValueError("SHA-256 mismatch: the Pastebin text changed or is incomplete")
    return files, positives, boxes, digest.hexdigest()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("input", type=Path, help="Raw TXT downloaded or copied from Pastebin")
    parser.add_argument("--output", type=Path, help="Output labels folder (default: beside TXT)")
    args = parser.parse_args()
    input_path = args.input.expanduser().resolve()
    output = (args.output or input_path.parent / "labels").expanduser().resolve()
    if output.exists():
        raise FileExistsError(f"Output already exists; refusing to overwrite: {output}")
    files, positives, boxes, checksum = decode_file(input_path)

    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = Path(tempfile.mkdtemp(prefix=".labels-restore-", dir=output.parent))
    try:
        for filename, content in files:
            (temporary / filename).write_bytes(content)
        if len(list(temporary.glob("*.txt"))) != len(files):
            raise ValueError("File count changed during write")
        os.replace(temporary, output)
    finally:
        if temporary.exists():
            shutil.rmtree(temporary)
    print(f"Restored {len(files)} labels to {output}")
    print(f"Positive: {positives}; empty: {len(files) - positives}; boxes: {boxes}")
    print(f"SHA-256: {checksum}")


if __name__ == "__main__":
    main()