import sys

from PIL import Image

from augmenator.embedding_planner import SIMILARITY_THRESHOLD
from augmenator.planner import plan_from_instruction, warmup
from augmenator.spatial import apply_spatial_ops

warmup()

planning_cases = [
    ("rotate randomly", ["rotate"]),
    ("rotate left", ["rotate_left"]),
    ("rotate right", ["rotate_right"]),
    ("flip upside down", ["flip_vertical"]),
    ("rotate 180 degrees", ["rotate_180"]),
    ("gamma brighten", ["gamma_up"]),
    ("make it brighter", ["brighten"]),
    ("vintage warm look", ["sepia", "warmer"]),
    ("replace background with beach sunset", ["replace_background"]),
    ("put a table behind it", None),
    ("blur everything softly", ["blur"]),
    ("cartoon style", ["style_candy"]),
    ("anime style", ["style_candy"]),
    ("starry night style", ["style_starry_night"]),
    ("pencil sketch", ["style_sketch"]),
    ("stylize", ["style_random"]),
    ("mosaic painting", ["style_mosaic"]),
    ("impressionist look", ["style_rain_princess"]),
    ("asdf qwerty nonsense", None),
    ("more contrast", ["contrast_up"]),
    ("transparent cutout", ["add_cutout_transparent"]),
    ("cover a piece but not text", ["cover_avoid_text"]),
    ("blur everything softly", []),  # must not trigger spatial ops
]


def spatial_ids(plan):
    ids = []
    for s in plan.get("spatial", []):
        op = s["op"]
        if op == "add_cutout":
            style = s.get("cutout_style", "solid")
            if style in ("transparent", "solid"):
                ids.append(f"add_cutout_{style}")
            else:
                ids.append("add_cutout")
        else:
            ids.append(op)
    return ids


def run_planning_tests():
    print(f"Threshold: {SIMILARITY_THRESHOLD}\n")
    failed = 0
    for instr, expected_ids in planning_cases:
        plan = plan_from_instruction(instr)
        got_ids = [s[0] for s in plan.get("scores", [])]
        all_matched = got_ids + plan["tags"] + spatial_ids(plan)

    if expected_ids is None:
        ok = not plan["supported"]
    elif expected_ids == []:
        ok = plan["supported"] and not plan.get("spatial")
    else:
        ok = all(eid in all_matched for eid in expected_ids)

        status = "OK" if ok else "FAIL"
        if not ok:
            failed += 1
        print(f"{status}: {instr!r}")
        print(f"  supported={plan['supported']} scores={plan.get('scores', [])}")
        print(f"  tags={plan['tags']} spatial={spatial_ids(plan)}")
        if not ok:
            print(f"  expected: {expected_ids}")
        print()

    return failed


def run_cutout_execution_tests():
    print("Cutout execution tests:\n")
    failed = 0
    img = Image.new("RGB", (400, 300), color=(100, 150, 200))

    execution_cases = [
        ("add cutout", 1, 4),
        ("cover and add cutout", 1, 4),
        ("cover and several cutouts avoid text", 1, 4),
        ("add cutout avoid text", 1, 4),
        ("3 cutouts", 3, 3),
        ("one cutout", 1, 1),
    ]

    for instr, min_cutouts, max_cutouts in execution_cases:
        plan = plan_from_instruction(instr)
        _, applied = apply_spatial_ops(img.copy(), plan["spatial"], 1.0, instr)

        cutout_count = 0
        for meta in applied:
            if meta.get("cutouts"):
                cutout_count += len(meta["cutouts"])
            elif meta.get("op") in ("add_cutout", "cutout") and meta.get("cutout_count"):
                cutout_count += meta["cutout_count"]

        ok = min_cutouts <= cutout_count <= max_cutouts
        status = "OK" if ok else "FAIL"
        if not ok:
            failed += 1
        print(f"{status}: {instr!r} -> {cutout_count} cutout(s) (expected {min_cutouts}-{max_cutouts})")
        for meta in applied:
            if meta.get("cutouts"):
                shapes = [c.get("shape") for c in meta["cutouts"]]
                print(f"  shapes: {shapes}")
        print()

    return failed


failed = run_planning_tests()
failed += run_cutout_execution_tests()

if failed:
    raise SystemExit(f"{failed} test(s) failed")

print("All validation tests passed")