"""Bar-model (Singapore model-method) renderer -> PNG.

Reads a structured spec and draws unit boxes with labels, a "?" bracket,
and an optional total brace. Used by Workstream B to attach drawn model
diagrams to solutions (the Flutter solution modal renders only text/LaTeX,
so the bars have to be a real image).

Spec shape:
{
  "unit_px": 90,          # width of ONE unit box
  "bars": [
    {"label": "Alan",
     "segments": [{"u": 1, "text": "", "fill": "#dbe9ff"},
                  {"u": 1, "text": "", "fill": "#dbe9ff"},
                  {"u": 1, "text": "", "fill": "#dbe9ff"}],
     "brace": {"text": "?"}},          # brace under whole bar
    {"label": "Ben",
     "segments": [{"u": 1, "text": "", "fill": "#ffe7c2"}]}
  ],
  "total": {"text": "$640", "bars": [0,1]}   # right-side brace spanning bars
}
"""
import json, sys, os
from PIL import Image, ImageDraw, ImageFont

BASE = os.path.dirname(os.path.abspath(__file__))

def _font(sz, bold=True):
    p = "C:/Windows/Fonts/" + ("arialbd.ttf" if bold else "arial.ttf")
    try:
        return ImageFont.truetype(p, sz)
    except Exception:
        return ImageFont.load_default()

def _hex(c):
    c = (c or "#ffffff").lstrip("#")
    return tuple(int(c[i:i+2], 16) for i in (0, 2, 4))

def render(spec, out_path):
    U      = int(spec.get("unit_px", 90))
    BH     = 56            # bar height
    GAP    = 26            # vertical gap between bars
    LBL_W  = 130           # left label gutter
    PAD    = 24
    BRACE  = 30            # space under a bar for its "?" brace
    f      = _font(26); fs = _font(22, bold=False); fb = _font(26)

    bars = spec["bars"]
    has_total = bool(spec.get("total"))
    max_units = max(sum(s["u"] for s in b["segments"]) for b in bars)

    # Right gutter sized to fit the total-brace label (was a fixed 120 that
    # clipped long labels like "4y + 9 = 53 km").
    RTOTAL = 0
    if has_total:
        ttxt = str(spec["total"].get("text", ""))
        _tmp = ImageDraw.Draw(Image.new("RGB", (8, 8)))
        tw = _tmp.textbbox((0, 0), ttxt, font=fb)[2]
        RTOTAL = 16 + 16 + tw + 24   # brace gap + tick + label + pad

    W = LBL_W + max_units * U + RTOTAL + PAD * 2
    rowh = [BH + (BRACE if b.get("brace") else 0) for b in bars]
    H = PAD * 2 + sum(rowh) + GAP * (len(bars) - 1)

    img = Image.new("RGB", (W, H), (255, 255, 255))
    d = ImageDraw.Draw(img)

    def ctext(cx, cy, t, font, fill=(0, 0, 0)):
        l, tp, r, b = d.textbbox((0, 0), t, font=font)
        d.text((cx - (r - l) / 2, cy - (b - tp) / 2), t, font=font, fill=fill)

    y = PAD
    bar_top = {}
    for i, bar in enumerate(bars):
        x = PAD + LBL_W
        bar_top[i] = y
        # left label
        d.text((PAD, y + BH / 2 - 13), bar.get("label", ""), font=f, fill=(0, 0, 0))
        seg_x = x
        for s in bar["segments"]:
            w = s["u"] * U
            d.rectangle([seg_x, y, seg_x + w, y + BH], outline=(0, 0, 0),
                        width=3, fill=_hex(s.get("fill", "#ffffff")))
            if s.get("text"):
                ctext(seg_x + w / 2, y + BH / 2, str(s["text"]), fb)
            seg_x += w
        full_w = sum(s["u"] for s in bar["segments"]) * U
        # under-bar brace with "?" (or custom)
        if bar.get("brace"):
            by = y + BH + 8
            d.line([x, by, x + full_w, by], fill=(0, 0, 0), width=2)
            d.line([x, by, x, by - 6], fill=(0, 0, 0), width=2)
            d.line([x + full_w, by, x + full_w, by - 6], fill=(0, 0, 0), width=2)
            ctext(x + full_w / 2, by + 14, str(bar["brace"].get("text", "?")), f)
        y += rowh[i] + GAP

    # right-side total brace spanning given bars
    if has_total:
        tb = spec["total"]
        idxs = tb.get("bars", list(range(len(bars))))
        spanw = max(sum(s["u"] for s in bars[i]["segments"]) for i in idxs) * U
        bx = PAD + LBL_W + spanw + 16
        ytop = bar_top[min(idxs)]
        ybot = bar_top[max(idxs)] + BH
        d.line([bx, ytop, bx, ybot], fill=(0, 0, 0), width=2)
        d.line([bx, ytop, bx - 6, ytop], fill=(0, 0, 0), width=2)
        d.line([bx, ybot, bx - 6, ybot], fill=(0, 0, 0), width=2)
        d.line([bx, (ytop + ybot) / 2, bx + 10, (ytop + ybot) / 2], fill=(0, 0, 0), width=2)
        d.text((bx + 16, (ytop + ybot) / 2 - 13), str(tb.get("text", "")), font=fb, fill=(0, 0, 0))

    img.save(out_path)
    return out_path

if __name__ == "__main__":
    # demo spec: "Alan and Ben shared $640. Alan got 3x as much as Ben."
    demo = {
        "unit_px": 90,
        "bars": [
            {"label": "Alan",
             "segments": [{"u": 1, "fill": "#dbe9ff"}, {"u": 1, "fill": "#dbe9ff"},
                          {"u": 1, "fill": "#dbe9ff"}],
             "brace": {"text": "?"}},
            {"label": "Ben",
             "segments": [{"u": 1, "fill": "#ffe7c2"}]},
        ],
        "total": {"text": "$640", "bars": [0, 1]},
    }
    out = sys.argv[1] if len(sys.argv) > 1 else os.path.join(BASE, "_barmodel_sample.png")
    spec = demo
    if len(sys.argv) > 2:
        spec = json.load(open(sys.argv[2], encoding="utf-8"))
    print(render(spec, out))