"""
Local test for SmartGridCutter node logic — runs without ComfyUI.
Uses PIL (Pillow) and numpy only.
"""
import sys
sys.path.insert(0, '.')

import numpy as np
from PIL import Image
import torch
import os

# Import our node
from nodes import SmartGridCutter

def test_with_real_image():
    """Test with actual storyboard image."""
    img_path = os.path.join('..', 'test_videos', 'cf_pipeline', 'storyboard.png')
    
    if not os.path.exists(img_path):
        print(f"❌ Test image not found: {img_path}")
        return False
    
    print(f"📷 Loading: {img_path}")
    pil_img = Image.open(img_path).convert('RGB')
    print(f"   Size: {pil_img.size[0]}×{pil_img.size[1]}")
    
    # Convert to ComfyUI tensor format [B, H, W, C]
    arr = np.array(pil_img).astype(np.float32) / 255.0
    tensor = torch.from_numpy(arr).unsqueeze(0)
    print(f"   Tensor: {tensor.shape}")
    
    # Test auto mode
    cutter = SmartGridCutter()
    result = cutter.cut_grid(tensor, mode="auto")
    
    # Unpack results
    *frames, cols, rows, total, info = result
    
    print(f"\n✅ РЕЗУЛЬТАТ:")
    print(f"   {info}")
    print(f"   Grid: {cols}×{rows} = {total} frames")
    
    # Show frame sizes
    for i in range(total):
        f = frames[i]
        if f.shape[1] > 16:  # Not a placeholder
            print(f"   Frame {i}: {f.shape[2]}×{f.shape[1]}px ({f.shape})")
    
    # Save frames for visual inspection
    out_dir = os.path.join('..', 'test_videos', 'cf_pipeline', 'cut_test')
    os.makedirs(out_dir, exist_ok=True)
    
    for i in range(total):
        f = frames[i]
        if f.shape[1] > 16:
            f_arr = (f[0].numpy() * 255).clip(0, 255).astype(np.uint8)
            f_img = Image.fromarray(f_arr)
            path = os.path.join(out_dir, f'frame_{i:02d}.png')
            f_img.save(path)
    
    print(f"\n   💾 Saved {total} frames to {out_dir}/")
    
    # Test 3x3 preset mode
    print(f"\n🔧 Testing 3x3 preset mode...")
    result2 = cutter.cut_grid(tensor, mode="3x3")
    *frames2, cols2, rows2, total2, info2 = result2
    print(f"   {info2}")
    
    return True


def test_with_synthetic():
    """Test with a synthetic grid image."""
    print("🧪 Testing with synthetic 3×3 grid (900×900)...")
    
    # Create a 3×3 grid with white dividers
    cell_size = 290
    divider = 10
    total = cell_size * 3 + divider * 2  # 890
    
    img = np.zeros((total, total, 3), dtype=np.uint8)
    
    colors = [
        (255, 0, 0), (0, 255, 0), (0, 0, 255),
        (255, 255, 0), (255, 0, 255), (0, 255, 255),
        (128, 0, 0), (0, 128, 0), (0, 0, 128),
    ]
    
    for r in range(3):
        for c in range(3):
            y = r * (cell_size + divider)
            x = c * (cell_size + divider)
            idx = r * 3 + c
            img[y:y+cell_size, x:x+cell_size] = colors[idx]
    
    # White dividers
    for i in range(2):
        pos = (i + 1) * cell_size + i * divider
        img[pos:pos+divider, :] = 255  # Horizontal
        img[:, pos:pos+divider] = 255  # Vertical
    
    tensor = torch.from_numpy(img.astype(np.float32) / 255.0).unsqueeze(0)
    
    cutter = SmartGridCutter()
    result = cutter.cut_grid(tensor, mode="auto")
    *frames, cols, rows, total_frames, info = result
    
    print(f"   {info}")
    assert cols == 3, f"Expected 3 cols, got {cols}"
    assert rows == 3, f"Expected 3 rows, got {rows}"
    print("   ✅ Synthetic test PASSED")
    return True


if __name__ == "__main__":
    print("=" * 60)
    print("  SmartGridCutter — Local Test Suite")
    print("=" * 60)
    
    ok1 = test_with_synthetic()
    print()
    ok2 = test_with_real_image()
    
    print("\n" + "=" * 60)
    if ok1 and ok2:
        print("  ✅ ALL TESTS PASSED")
    else:
        print("  ⚠️ Some tests failed")
    print("=" * 60)