#!/usr/bin/env python3
import json
import argparse
import sys
import os

def main():
    parser = argparse.ArgumentParser(description="Render a structured UI layout JSON to a monochrome PNG image.")
    parser.add_argument("-i", "--input", required=True, help="Path to input JSON layout file")
    parser.add_argument("-o", "--output", required=True, help="Path to save output PNG image")
    args = parser.parse_args()

    if not os.path.exists(args.input):
        print(f"Error: Input JSON file {args.input} does not exist.", file=sys.stderr)
        sys.exit(1)

    try:
        from PIL import Image, ImageDraw, ImageFont
    except ImportError:
        print("Error: Pillow library is required. Install with: pip install Pillow", file=sys.stderr)
        sys.exit(1)

    try:
        with open(args.input, "r") as f:
            layout = json.load(f)
    except Exception as e:
        print(f"Error parsing JSON: {e}", file=sys.stderr)
        sys.exit(1)

    canvas_cfg = layout.get("canvas", {})
    width = canvas_cfg.get("width", 256)
    height = canvas_cfg.get("height", 256)
    bg_color = canvas_cfg.get("background", "white")

    # Map colors to binary (0 = Black, 255 = White) for 1-bit mode
    bg_val = 255 if bg_color.lower() in ("white", "#ffffff", "255") else 0
    fg_val = 0 if bg_val == 255 else 255

    # Create 1-bit image (mode "1")
    img = Image.new("1", (width, height), bg_val)
    draw = ImageDraw.Draw(img)

    # Attempt to load a default font
    try:
        # Try to use a basic system font
        font = ImageFont.load_default()
    except Exception:
        font = None

    elements = layout.get("elements", [])
    for el in elements:
        el_type = el.get("type")
        fill_color = fg_val if el.get("fill") else None
        outline_color = fg_val if el.get("outline", True) else None
        width_val = el.get("width", 1)

        try:
            if el_type == "rect":
                x1, y1 = el["x1"], el["y1"]
                x2, y2 = el["x2"], el["y2"]
                draw.rectangle([x1, y1, x2, y2], fill=fill_color, outline=outline_color, width=width_val)
            elif el_type == "line":
                x1, y1 = el["x1"], el["y1"]
                x2, y2 = el["x2"], el["y2"]
                draw.line([x1, y1, x2, y2], fill=fg_val, width=width_val)
            elif el_type == "circle" or el_type == "ellipse":
                x1, y1 = el["x1"], el["y1"]
                x2, y2 = el["x2"], el["y2"]
                draw.ellipse([x1, y1, x2, y2], fill=fill_color, outline=outline_color, width=width_val)
            elif el_type == "text":
                x, y = el["x"], el["y"]
                txt = el["text"]
                draw.text((x, y), txt, fill=fg_val, font=font)
        except KeyError as ke:
            print(f"Warning: Missing key {ke} in element: {el}", file=sys.stderr)

    try:
        # Ensure output directory exists
        out_dir = os.path.dirname(args.output)
        if out_dir:
            os.makedirs(out_dir, exist_ok=True)
            
        img.save(args.output, format="PNG", optimize=True)
        print(f"Success: Rendered layout to {args.output} ({os.path.getsize(args.output) / 1024:.2f} KB)")
    except Exception as e:
        print(f"Error saving image: {e}", file=sys.stderr)
        sys.exit(1)

if __name__ == "__main__":
    main()