#!/usr/bin/env python3
"""
Funky comparison tests: PIL vs TextStrip system.
Tests colors, opacity, fonts, sizes, edge positions, clipping, overlaps, etc.
"""

import numpy as np
import jax.numpy as jnp
from PIL import Image, ImageDraw, ImageFont
from streaming.jax_typography import (
    render_text_strip, place_text_strip_jax, _load_font
)

FONTS = {
    'sans': '/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf',
    'bold': '/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf',
    'serif': '/usr/share/fonts/truetype/dejavu/DejaVuSerif.ttf',
    'serif_bold': '/usr/share/fonts/truetype/dejavu/DejaVuSerif-Bold.ttf',
    'mono': '/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf',
    'mono_bold': '/usr/share/fonts/truetype/dejavu/DejaVuSansMono-Bold.ttf',
    'narrow': '/usr/share/fonts/truetype/liberation/LiberationSansNarrow-Regular.ttf',
    'italic': '/usr/share/fonts/truetype/freefont/FreeSansOblique.ttf',
}


def render_pil(text, x, y, font_path=None, font_size=36, frame_size=(400, 100),
               fill=(255, 255, 255), bg=(0, 0, 0), opacity=1.0,
               stroke_width=0, stroke_fill=None, anchor="la",
               multiline=False, line_spacing=4, align="left"):
    """Render with PIL directly, including color/opacity."""
    frame = np.full((frame_size[1], frame_size[0], 3), bg, dtype=np.uint8)
    # For opacity, render to RGBA then composite
    txt_layer = Image.new('RGBA', frame_size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(txt_layer)
    font = _load_font(font_path, font_size)

    if stroke_fill is None:
        stroke_fill = (0, 0, 0)

    # PIL fill with alpha for opacity
    alpha_int = int(round(opacity * 255))
    fill_rgba = (*fill, alpha_int)
    stroke_rgba = (*stroke_fill, alpha_int) if stroke_width > 0 else None

    if multiline:
        draw.multiline_text((x, y), text, fill=fill_rgba, font=font,
                           stroke_width=stroke_width, stroke_fill=stroke_rgba,
                           spacing=line_spacing, align=align, anchor=anchor)
    else:
        draw.text((x, y), text, fill=fill_rgba, font=font,
                  stroke_width=stroke_width, stroke_fill=stroke_rgba, anchor=anchor)

    # Composite onto background
    bg_img = Image.fromarray(frame).convert('RGBA')
    result = Image.alpha_composite(bg_img, txt_layer)
    return np.array(result.convert('RGB'))


def render_strip(text, x, y, font_path=None, font_size=36, frame_size=(400, 100),
                 fill=(255, 255, 255), bg=(0, 0, 0), opacity=1.0,
                 stroke_width=0, stroke_fill=None, anchor="la",
                 multiline=False, line_spacing=4, align="left"):
    """Render with TextStrip system."""
    frame = jnp.full((frame_size[1], frame_size[0], 3), jnp.array(bg, dtype=jnp.uint8), dtype=jnp.uint8)

    strip = render_text_strip(
        text, font_path, font_size,
        stroke_width=stroke_width, stroke_fill=stroke_fill,
        anchor=anchor, multiline=multiline, line_spacing=line_spacing, align=align
    )
    strip_img = jnp.asarray(strip.image)
    color = jnp.array(fill, dtype=jnp.float32)

    result = place_text_strip_jax(
        frame, strip_img, x, y,
        strip.baseline_y, strip.bearing_x,
        color, opacity,
        anchor_x=strip.anchor_x, anchor_y=strip.anchor_y,
        stroke_width=strip.stroke_width
    )
    return np.array(result)


def compare(name, tolerance=0, **kwargs):
    """Compare PIL and TextStrip rendering."""
    pil = render_pil(**kwargs)
    strip = render_strip(**kwargs)

    diff = np.abs(pil.astype(np.int16) - strip.astype(np.int16))
    max_diff = diff.max()
    pixels_diff = (diff > 0).any(axis=2).sum()

    if max_diff == 0:
        print(f"  PASS: {name}")
        return True

    if tolerance > 0:
        best_diff = diff.copy()
        for dy in range(-tolerance, tolerance + 1):
            for dx in range(-tolerance, tolerance + 1):
                if dy == 0 and dx == 0:
                    continue
                shifted = np.roll(np.roll(strip, dy, axis=0), dx, axis=1)
                sdiff = np.abs(pil.astype(np.int16) - shifted.astype(np.int16))
                best_diff = np.minimum(best_diff, sdiff)
        if best_diff.max() == 0:
            print(f"  PASS: {name} (within {tolerance}px tolerance)")
            return True

    print(f"  FAIL: {name}")
    print(f"        Max diff: {max_diff}, Pixels different: {pixels_diff}")
    Image.fromarray(pil).save(f"/tmp/pil_{name}.png")
    Image.fromarray(strip).save(f"/tmp/strip_{name}.png")
    diff_vis = np.clip(diff * 10, 0, 255).astype(np.uint8)
    Image.fromarray(diff_vis).save(f"/tmp/diff_{name}.png")
    print(f"        Saved: /tmp/pil_{name}.png /tmp/strip_{name}.png /tmp/diff_{name}.png")
    return False


def test_colors():
    """Test various text colors on various backgrounds."""
    print("\n--- Colors ---")
    results = []

    # White on black (baseline)
    results.append(compare("color_white_on_black",
        text="Hello", x=20, y=30, fill=(255, 255, 255), bg=(0, 0, 0)))

    # Red on black
    results.append(compare("color_red",
        text="Red Text", x=20, y=30, fill=(255, 0, 0), bg=(0, 0, 0)))

    # Green on black
    results.append(compare("color_green",
        text="Green!", x=20, y=30, fill=(0, 255, 0), bg=(0, 0, 0)))

    # Blue on black
    results.append(compare("color_blue",
        text="Blue sky", x=20, y=30, fill=(0, 100, 255), bg=(0, 0, 0)))

    # Yellow on dark gray
    results.append(compare("color_yellow_on_gray",
        text="Yellow", x=20, y=30, fill=(255, 255, 0), bg=(40, 40, 40)))

    # Magenta on white
    results.append(compare("color_magenta_on_white",
        text="Magenta", x=20, y=30, fill=(255, 0, 255), bg=(255, 255, 255)))

    # Subtle: gray text on slightly lighter gray
    results.append(compare("color_subtle_gray",
        text="Subtle", x=20, y=30, fill=(128, 128, 128), bg=(64, 64, 64)))

    # Orange on deep blue
    results.append(compare("color_orange_on_blue",
        text="Warm", x=20, y=30, fill=(255, 165, 0), bg=(0, 0, 80)))

    return results


def test_opacity():
    """Test different opacity levels."""
    print("\n--- Opacity ---")
    results = []

    results.append(compare("opacity_100",
        text="Full", x=20, y=30, opacity=1.0))

    results.append(compare("opacity_75",
        text="75%", x=20, y=30, opacity=0.75))

    results.append(compare("opacity_50",
        text="Half", x=20, y=30, opacity=0.5))

    results.append(compare("opacity_25",
        text="Quarter", x=20, y=30, opacity=0.25))

    results.append(compare("opacity_10",
        text="Ghost", x=20, y=30, opacity=0.1))

    # Opacity on colored background
    results.append(compare("opacity_on_colored_bg",
        text="Overlay", x=20, y=30, fill=(255, 255, 255), bg=(100, 0, 0),
        opacity=0.5))

    # Color + opacity combo
    results.append(compare("opacity_red_on_green",
        text="Blend", x=20, y=30, fill=(255, 0, 0), bg=(0, 100, 0),
        opacity=0.6))

    return results


def test_fonts():
    """Test different fonts and sizes."""
    print("\n--- Fonts & Sizes ---")
    results = []

    for label, path in FONTS.items():
        results.append(compare(f"font_{label}",
            text="Quick Fox", x=20, y=30, font_path=path, font_size=28,
            frame_size=(300, 80)))

    # Tiny text
    results.append(compare("size_tiny",
        text="Tiny text at 12px", x=10, y=15, font_size=12,
        frame_size=(200, 40)))

    # Big text
    results.append(compare("size_big",
        text="BIG", x=20, y=10, font_size=72,
        frame_size=(300, 100)))

    # Huge text
    results.append(compare("size_huge",
        text="XL", x=10, y=10, font_size=120,
        frame_size=(300, 160)))

    return results


def test_anchors():
    """Test all anchor combinations."""
    print("\n--- Anchors ---")
    results = []

    # All horizontal x vertical combos
    for h in ['l', 'm', 'r']:
        for v in ['a', 'm', 's', 'd']:
            anchor = h + v
            results.append(compare(f"anchor_{anchor}",
                text="Anchor", x=200, y=50, anchor=anchor,
                frame_size=(400, 100)))

    return results


def test_strokes():
    """Test various stroke widths and colors."""
    print("\n--- Strokes ---")
    results = []

    for sw in [1, 2, 3, 4, 6, 8]:
        results.append(compare(f"stroke_w{sw}",
            text="Stroke", x=30, y=20, font_size=40,
            stroke_width=sw, stroke_fill=(0, 0, 0),
            frame_size=(300, 80)))

    # Colored strokes
    results.append(compare("stroke_red",
        text="Red outline", x=20, y=20, font_size=36,
        stroke_width=3, stroke_fill=(255, 0, 0),
        frame_size=(350, 80)))

    results.append(compare("stroke_white_on_black",
        text="Glow", x=20, y=20, font_size=40,
        fill=(255, 255, 255), stroke_width=4, stroke_fill=(0, 0, 255),
        frame_size=(250, 80)))

    # Stroke with bold font
    results.append(compare("stroke_bold",
        text="Bold+Stroke", x=20, y=20,
        font_path=FONTS['bold'], font_size=36,
        stroke_width=3, stroke_fill=(0, 0, 0),
        frame_size=(400, 80)))

    # Stroke + colored text on colored bg
    results.append(compare("stroke_colored_on_bg",
        text="Party", x=20, y=20, font_size=48,
        fill=(255, 255, 0), bg=(50, 0, 80),
        stroke_width=3, stroke_fill=(255, 0, 0),
        frame_size=(300, 80)))

    return results


def test_edge_clipping():
    """Test text at frame edges - clipping behavior."""
    print("\n--- Edge Clipping ---")
    results = []

    # Text at very left edge
    results.append(compare("clip_left_edge",
        text="LEFT", x=0, y=30, frame_size=(200, 80)))

    # Text partially off right edge
    results.append(compare("clip_right_edge",
        text="RIGHT SIDE", x=150, y=30, frame_size=(200, 80)))

    # Text at top edge
    results.append(compare("clip_top",
        text="TOP", x=20, y=0, frame_size=(200, 80)))

    # Text at bottom edge - partially clipped
    results.append(compare("clip_bottom",
        text="BOTTOM", x=20, y=55, font_size=40,
        frame_size=(200, 80)))

    # Large text overflowing all sides from center
    results.append(compare("clip_overflow_center",
        text="HUGE", x=75, y=40, font_size=100, anchor="mm",
        frame_size=(150, 80)))

    # Corner placement
    results.append(compare("clip_corner_br",
        text="Corner", x=350, y=70, font_size=36,
        frame_size=(400, 100)))

    return results


def test_multiline_fancy():
    """Test multiline with various styles."""
    print("\n--- Multiline Fancy ---")
    results = []

    # Right-aligned (1px tolerance: same sub-pixel issue as center alignment)
    results.append(compare("multi_right",
        text="Right\nAligned\nText", x=380, y=20,
        frame_size=(400, 150),
        multiline=True, anchor="ra", align="right",
        tolerance=1))

    # Center + stroke
    results.append(compare("multi_center_stroke",
        text="Title\nSubtitle", x=200, y=20,
        font_size=32, frame_size=(400, 120),
        multiline=True, anchor="ma", align="center",
        stroke_width=2, stroke_fill=(0, 0, 0),
        tolerance=1))

    # Wide line spacing
    results.append(compare("multi_wide_spacing",
        text="Line A\nLine B\nLine C", x=20, y=10,
        frame_size=(300, 200),
        multiline=True, line_spacing=20))

    # Zero extra spacing
    results.append(compare("multi_tight_spacing",
        text="Tight\nPacked\nLines", x=20, y=10,
        frame_size=(300, 150),
        multiline=True, line_spacing=0))

    # Many lines
    results.append(compare("multi_many_lines",
        text="One\nTwo\nThree\nFour\nFive\nSix", x=20, y=5,
        font_size=20, frame_size=(200, 200),
        multiline=True, line_spacing=4))

    # Bold multiline with stroke
    results.append(compare("multi_bold_stroke",
        text="BOLD\nSTROKE", x=20, y=10,
        font_path=FONTS['bold'], font_size=48,
        stroke_width=3, stroke_fill=(200, 0, 0),
        frame_size=(350, 150), multiline=True))

    # Multiline on colored bg with opacity
    results.append(compare("multi_opacity_on_bg",
        text="Semi\nTransparent", x=20, y=10,
        fill=(255, 255, 0), bg=(0, 50, 100), opacity=0.7,
        frame_size=(300, 120), multiline=True))

    return results


def test_special_chars():
    """Test special characters and edge cases."""
    print("\n--- Special Characters ---")
    results = []

    # Numbers and symbols
    results.append(compare("chars_numbers",
        text="0123456789", x=20, y=30, frame_size=(300, 80)))

    results.append(compare("chars_punctuation",
        text="Hello, World! (v2.0)", x=10, y=30, frame_size=(350, 80)))

    results.append(compare("chars_symbols",
        text="@#$%^&*+=", x=20, y=30, frame_size=(300, 80)))

    # Single character
    results.append(compare("chars_single",
        text="X", x=50, y=30, font_size=48, frame_size=(100, 80)))

    # Very long text (clipped)
    results.append(compare("chars_long",
        text="The quick brown fox jumps over the lazy dog", x=10, y=30,
        font_size=24, frame_size=(400, 80)))

    # Mixed case
    results.append(compare("chars_mixed_case",
        text="AaBbCcDdEeFf", x=10, y=30, frame_size=(350, 80)))

    return results


def test_combos():
    """Complex combinations of features."""
    print("\n--- Combos ---")
    results = []

    # Big bold stroke + color + opacity
    results.append(compare("combo_all_features",
        text="EPIC", x=20, y=10,
        font_path=FONTS['bold'], font_size=64,
        fill=(255, 200, 0), bg=(20, 0, 40), opacity=0.85,
        stroke_width=4, stroke_fill=(180, 0, 0),
        frame_size=(350, 100)))

    # Small mono on busy background
    results.append(compare("combo_mono_code",
        text="fn main() {}", x=10, y=15,
        font_path=FONTS['mono'], font_size=16,
        fill=(0, 255, 100), bg=(30, 30, 30),
        frame_size=(250, 50)))

    # Serif italic multiline with stroke
    results.append(compare("combo_serif_italic_multi",
        text="Once upon\na time...", x=20, y=10,
        font_path=FONTS['italic'], font_size=28,
        stroke_width=1, stroke_fill=(80, 80, 80),
        frame_size=(300, 120), multiline=True))

    # Narrow font, big stroke, center anchored
    results.append(compare("combo_narrow_stroke_center",
        text="NARROW", x=150, y=40,
        font_path=FONTS['narrow'], font_size=44,
        stroke_width=5, stroke_fill=(0, 0, 0),
        anchor="mm", frame_size=(300, 80)))

    # Multiple strips on same frame (simulated by sequential placement)
    results.append(compare("combo_opacity_blend",
        text="Layered", x=20, y=30,
        fill=(255, 0, 0), bg=(0, 0, 255), opacity=0.5,
        font_size=48, frame_size=(300, 80)))

    return results


def test_multi_strip_overlay():
    """Test placing multiple strips on the same frame."""
    print("\n--- Multi-Strip Overlay ---")
    results = []

    frame_size = (400, 150)
    bg = (20, 20, 40)

    # PIL version - multiple draw calls
    font1 = _load_font(FONTS['bold'], 48)
    font2 = _load_font(None, 24)
    font3 = _load_font(FONTS['mono'], 18)

    pil_frame = np.full((frame_size[1], frame_size[0], 3), bg, dtype=np.uint8)
    txt_layer = Image.new('RGBA', frame_size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(txt_layer)
    draw.text((20, 10), "TITLE", fill=(255, 255, 0, 255), font=font1,
              stroke_width=2, stroke_fill=(0, 0, 0, 255))
    draw.text((20, 70), "Subtitle here", fill=(200, 200, 200, 255), font=font2)
    draw.text((20, 110), "code_snippet()", fill=(0, 255, 128, 200), font=font3)
    bg_img = Image.fromarray(pil_frame).convert('RGBA')
    pil_result = np.array(Image.alpha_composite(bg_img, txt_layer).convert('RGB'))

    # Strip version - multiple placements
    frame = jnp.full((frame_size[1], frame_size[0], 3), jnp.array(bg, dtype=jnp.uint8), dtype=jnp.uint8)

    s1 = render_text_strip("TITLE", FONTS['bold'], 48, stroke_width=2, stroke_fill=(0, 0, 0))
    s2 = render_text_strip("Subtitle here", None, 24)
    s3 = render_text_strip("code_snippet()", FONTS['mono'], 18)

    frame = place_text_strip_jax(
        frame, jnp.asarray(s1.image), 20, 10,
        s1.baseline_y, s1.bearing_x,
        jnp.array([255, 255, 0], dtype=jnp.float32), 1.0,
        anchor_x=s1.anchor_x, anchor_y=s1.anchor_y,
        stroke_width=s1.stroke_width)

    frame = place_text_strip_jax(
        frame, jnp.asarray(s2.image), 20, 70,
        s2.baseline_y, s2.bearing_x,
        jnp.array([200, 200, 200], dtype=jnp.float32), 1.0,
        anchor_x=s2.anchor_x, anchor_y=s2.anchor_y,
        stroke_width=s2.stroke_width)

    frame = place_text_strip_jax(
        frame, jnp.asarray(s3.image), 20, 110,
        s3.baseline_y, s3.bearing_x,
        jnp.array([0, 255, 128], dtype=jnp.float32), 200/255,
        anchor_x=s3.anchor_x, anchor_y=s3.anchor_y,
        stroke_width=s3.stroke_width)

    strip_result = np.array(frame)

    diff = np.abs(pil_result.astype(np.int16) - strip_result.astype(np.int16))
    max_diff = diff.max()
    pixels_diff = (diff > 0).any(axis=2).sum()

    if max_diff <= 1:
        print(f"  PASS: multi_overlay (max_diff={max_diff})")
        results.append(True)
    else:
        print(f"  FAIL: multi_overlay")
        print(f"        Max diff: {max_diff}, Pixels different: {pixels_diff}")
        Image.fromarray(pil_result).save("/tmp/pil_multi_overlay.png")
        Image.fromarray(strip_result).save("/tmp/strip_multi_overlay.png")
        diff_vis = np.clip(diff * 10, 0, 255).astype(np.uint8)
        Image.fromarray(diff_vis).save("/tmp/diff_multi_overlay.png")
        print(f"        Saved: /tmp/pil_multi_overlay.png /tmp/strip_multi_overlay.png /tmp/diff_multi_overlay.png")
        results.append(False)

    return results


def main():
    print("=" * 60)
    print("Funky TextStrip vs PIL Comparison")
    print("=" * 60)

    all_results = []
    all_results.extend(test_colors())
    all_results.extend(test_opacity())
    all_results.extend(test_fonts())
    all_results.extend(test_anchors())
    all_results.extend(test_strokes())
    all_results.extend(test_edge_clipping())
    all_results.extend(test_multiline_fancy())
    all_results.extend(test_special_chars())
    all_results.extend(test_combos())
    all_results.extend(test_multi_strip_overlay())

    print("\n" + "=" * 60)
    passed = sum(all_results)
    total = len(all_results)
    print(f"Results: {passed}/{total} passed")
    if passed == total:
        print("ALL TESTS PASSED!")
    else:
        failed = [i for i, r in enumerate(all_results) if not r]
        print(f"FAILED: {total - passed} tests (indices: {failed})")
    print("=" * 60)


if __name__ == "__main__":
    main()