#!/usr/bin/env python3
"""Compare a rendered dialog with the design reference screenshot.

Renders the XML (half-scale preview to reference resolution), composites the
card over the reference backdrop, measures layout features on both images and
prints a comparison table. Saves side_by_side.png and diff.png next to it.

Usage:
    python3 compare_ref.py <xml> [--profile episode] [--focus 111] [--out-dir DIR]
"""
import argparse
import os
import sys

import numpy as np
from PIL import Image, ImageDraw

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from kodi_renderer import Renderer, parse_po, STRINGS_PO  # noqa: E402
from render_dialog import BASE_PROPS, PROFILES  # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
REF_PATH = os.path.join(HERE, 'reference.jpg')

CARD_X, CARD_Y, CARD_W, CARD_H = 50, 50, 628, 290
REF_CARD_BOX = (10, 14, 430, 208)  # measured on the reference screenshot


def measure_card(card, scale=1):
    """Measure layout features on a card image; results are in reference-scale units.

    scale = how many pixels make one reference-scale pixel (2 for the 1080p render).
    """
    s = float(scale)
    si = int(round(scale))
    a = np.asarray(card.convert('RGB')).astype(int)
    b = a.sum(axis=2) / 3.0
    sat = a.max(axis=2) - a.min(axis=2)
    h, w, _ = a.shape
    out = {}

    def q(v):
        return int(round(v / s))

    mask = (sat > 40) & (b > 45)
    mask[:, int(160 * s):] = False
    mask[:int(8 * s), :] = False
    mask[int(150 * s):, :] = False
    mask[:, :int(6 * s)] = False
    ys, xs = np.where(mask)
    out['thumb'] = (q(xs.min()), q(ys.min()), q(xs.max()), q(ys.max())) if len(xs) else None

    text = b[:, int(140 * s):w - int(4 * s)]
    rows = (text > 130).any(axis=1)
    bands, start = [], None
    for y, v in enumerate(rows):
        if v and start is None:
            start = y
        if not v and start is not None:
            bands.append((start, y - 1))
            start = None
    if start is not None:
        bands.append((start, h - 1))
    merged = []
    for band in bands:
        if merged and band[0] - merged[-1][1] <= 2 * si:
            merged[-1] = (merged[-1][0], band[1])
        else:
            merged.append(band)
    out['bands'] = []
    for y0, y1 in merged:
        if y1 - y0 < 4 * si:
            continue
        cols = (b[y0:y1 + 1, int(140 * s):w] > 130).any(axis=0)
        xs2 = np.where(cols)[0]
        out['bands'].append((q(y0), q(y1), q(xs2.min()) + 140, q(xs2.max()) + 140))

    wm = b > 195
    rows_pill = []
    for y in range(h):
        best = cur = 0
        for v in wm[y]:
            cur = cur + 1 if v else 0
            best = max(best, cur)
        if best > 60 * si:
            rows_pill.append(y)
    if rows_pill:
        out['white_pill_y'] = (q(rows_pill[0]), q(rows_pill[-1]))
        ymid = (rows_pill[0] + rows_pill[-1]) // 2
        xs3 = np.where(wm[ymid])[0]
        out['white_pill_x'] = (q(xs3.min()), q(xs3.max()))
    else:
        out['white_pill_y'] = None
        out['white_pill_x'] = None

    out['bg'] = {
        'top_left': tuple(int(v) for v in a[int(22 * s), int(40 * s)]),
        'mid_right': tuple(int(v) for v in a[int(100 * s), int(355 * s)]),
        'bottom_left': tuple(int(v) for v in a[int(145 * s), int(40 * s)]),
    }
    return out


def fmt_band(band):
    if band is None:
        return 'none'
    return f"y {band[0]}..{band[1]}  x {band[2]}..{band[3]}"


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('xml')
    parser.add_argument('--profile', default='episode')
    parser.add_argument('--focus', default='111')
    parser.add_argument('--hidden', default='')
    parser.add_argument('--out-dir', default=HERE)
    parser.add_argument('--subdir', default='premium_episode')
    args = parser.parse_args()

    props = dict(BASE_PROPS)
    props.update(PROFILES.get(args.profile, {}))
    strings = parse_po(STRINGS_PO)
    focus = [int(x) for x in args.focus.split(',') if x.strip()]
    hidden = [int(x) for x in args.hidden.split(',') if x.strip()]

    renderer = Renderer(props, focus_ids=focus, strings=strings)
    full = renderer.render(args.xml, hidden_ids=hidden)

    out_dir = os.path.join(args.out_dir, args.subdir)
    os.makedirs(out_dir, exist_ok=True)

    preview_bg = Image.new('RGBA', full.size, (12, 12, 14, 255))
    preview_bg.alpha_composite(full)
    preview_bg.convert('RGB').save(os.path.join(out_dir, 'render_1080.png'))

    card = full.crop((CARD_X, CARD_Y, CARD_X + CARD_W, CARD_Y + CARD_H))
    card_flat = preview_bg.convert('RGB').crop((CARD_X, CARD_Y, CARD_X + CARD_W, CARD_Y + CARD_H))
    ref = Image.open(REF_PATH).convert('RGBA')
    ref_card = ref.crop(REF_CARD_BOX)
    ours_small = card.resize(ref_card.size, Image.LANCZOS)
    composite = ref_card.copy()
    composite.alpha_composite(ours_small)

    side = Image.new('RGB', (ref_card.width * 2 + 12, ref_card.height), (24, 24, 24))
    side.paste(ref_card.convert('RGB'), (0, 0))
    side.paste(composite.convert('RGB'), (ref_card.width + 12, 0))
    side = side.resize((side.width * 3, side.height * 3), Image.LANCZOS)
    side.save(os.path.join(out_dir, 'side_by_side.png'))

    ref_arr = np.asarray(ref_card.convert('RGB')).astype(int)
    our_arr = np.asarray(composite.convert('RGB')).astype(int)
    diff = np.abs(ref_arr - our_arr).astype(np.uint8)
    diff_img = Image.fromarray((np.clip(diff * 3, 0, 255)).astype(np.uint8), 'RGB')
    diff_img.resize((diff_img.width * 3, diff_img.height * 3), Image.NEAREST).save(
        os.path.join(out_dir, 'diff.png'))
    mean_diff = float(np.abs(ref_arr - our_arr).mean())

    ref_f = measure_card(ref_card, scale=1)
    our_f = measure_card(card_flat, scale=CARD_W / float(ref_card.width))
    our_f['bg'] = measure_card(composite, scale=1)['bg']

    print("=" * 78)
    print(f"COMPARE: {os.path.basename(args.xml)}  (focus={args.focus or '-'})")
    print("=" * 78)
    print(f"mean abs diff (0-255): {mean_diff:.2f}")
    print("-" * 78)
    print(f"{'feature':<16}{'reference':<34}{'ours':<34}")
    def row(name, r, o):
        print(f"{name:<16}{str(r):<34}{str(o):<34}")
    row('thumb', ref_f['thumb'], our_f['thumb'])
    row('white pill y', ref_f['white_pill_y'], our_f['white_pill_y'])
    row('white pill x', ref_f['white_pill_x'], our_f['white_pill_x'])
    for i in range(max(len(ref_f['bands']), len(our_f['bands']))):
        r = ref_f['bands'][i] if i < len(ref_f['bands']) else None
        o = our_f['bands'][i] if i < len(our_f['bands']) else None
        row(f'text band {i}', fmt_band(r), fmt_band(o))
    for key in ref_f['bg']:
        row(f'bg {key}', ref_f['bg'][key], our_f['bg'][key])
    print("-" * 78)
    print(f"outputs: {out_dir}/render_1080.png  side_by_side.png  diff.png")


if __name__ == '__main__':
    main()