import json
from pathlib import Path

from PIL import Image


DOCS_DIR = Path(__file__).resolve().parents[1]
SOURCE_PATH = DOCS_DIR / "CropDocs-Items_DSC_2416.jpg"
JSON_PATH = DOCS_DIR / "crop_subject_metadata" / "result.json"
OUTPUT_PATH = DOCS_DIR / "Crop_Subject_Mask_Overlay.png"


def decode_subject_mask_rle(mask_json):
    width = int(mask_json["width"])
    height = int(mask_json["height"])
    value = 1 if mask_json.get("startsWith") else 0
    pixels = []

    for run in mask_json["counts"]:
        pixels.extend([value] * int(run))
        value = 0 if value else 1

    expected = width * height
    if len(pixels) != expected:
        raise ValueError(f"Decoded {len(pixels)} pixels, expected {expected}")

    return width, height, pixels


def main():
    data = json.loads(JSON_PATH.read_text(encoding="utf-8"))
    subject_mask = data["plugins"][0]["subjectMask"]
    width, height, pixels = decode_subject_mask_rle(subject_mask)

    mask = Image.new("L", (width, height))
    mask.putdata(pixels)

    source = Image.open(SOURCE_PATH).convert("RGBA")
    mask = mask.resize(source.size, Image.Resampling.NEAREST)
    overlay = Image.new("RGBA", source.size, (0, 128, 255, 0))
    overlay.putalpha(mask.point(lambda p: 112 if p else 0))
    Image.alpha_composite(source, overlay).save(OUTPUT_PATH)


if __name__ == "__main__":
    main()