#!/usr/bin/env python3
"""Test script for scout and zoom workflow.

This script tests the new scout_regions_of_interest and detect_in_region tools
to verify they work correctly before running the full pipeline.
"""
import os
import sys
from pathlib import Path
import asyncio

# Set up environment
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "")
os.environ["DETECTOR_BACKEND"] = "grounded-sam-2"

# Set GROUNDED_SAM2_DIR if not already set
if not os.getenv("GROUNDED_SAM2_DIR"):
    default_path = Path(__file__).parent / "models" / "Grounded-SAM-2"
    if default_path.exists():
        os.environ["GROUNDED_SAM2_DIR"] = str(default_path)

from mcp.clients.agents.main import (
    PipelineContext,
    _crop_image_region,
    _map_mask_to_full_image,
    _mask_is_empty,
)

async def test_scout_workflow(image_path: str):
    """Test the scout and zoom workflow on an image."""
    print(f"\n{'='*60}")
    print(f"Testing Scout and Zoom Workflow")
    print(f"Image: {image_path}")
    print(f"{'='*60}\n")
    
    # Import required modules
    from PIL import Image
    import base64
    from io import BytesIO
    import json
    import openai
    
    # Check if image exists
    if not Path(image_path).exists():
        print(f"ERROR: Image not found: {image_path}")
        return False
    
    # Load image
    img = Image.open(image_path)
    img_w, img_h = img.size
    print(f"Image size: {img_w} x {img_h}")
    
    # Step 1: Scout for regions of interest
    print("\n[STEP 1] Scouting for regions of interest...")
    
    # Resize for VLM efficiency
    max_size = 1024
    if max(img_w, img_h) > max_size:
        scale = max_size / max(img_w, img_h)
        new_w = int(img_w * scale)
        new_h = int(img_h * scale)
        img_resized = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
    else:
        img_resized = img
        scale = 1.0
    
    buffer = BytesIO()
    img_resized.save(buffer, format='PNG')
    img_b64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
    
    prompt = f"""Analyze this street scene image and identify regions that may contain PII (personally identifiable information) such as text, logos, stickers, signs.

For each region, provide:
1. Bounding box as [x, y, width, height] in pixels (relative to image size {img_resized.size[0]}x{img_resized.size[1]})
2. Description of what PII might be there
3. Confidence score (0-1)

Focus on areas with:
- Text on vehicles, buildings, signs
- Logos and stickers on surfaces
- License plates
- Street signs with identifying information
- Posters or notices with text

Return JSON array: [{{"bbox": [x,y,w,h], "description": "...", "confidence": 0.0-1.0}}]

If no PII regions found, return empty array []."""
    
    try:
        client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "user", "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_b64}"}}
                ]}
            ],
            max_tokens=500,
            response_format={"type": "json_object"}
        )
        
        content = response.choices[0].message.content.strip()
        result = json.loads(content)
        regions = result.get("regions", result) if isinstance(result, dict) else result
        if not isinstance(regions, list):
            regions = []
        
        print(f"VLM found {len(regions)} raw regions")
        
        # Scale bboxes back to original image size
        scaled_regions = []
        for i, r in enumerate(regions):
            bbox = r.get("bbox", [])
            if len(bbox) != 4:
                continue
            
            # Scale coordinates
            x, y, w, h = bbox
            x = int(x / scale)
            y = int(y / scale)
            w = int(w / scale)
            h = int(h / scale)
            
            # Add margin (20% on each side)
            margin_w = int(w * 0.2)
            margin_h = int(h * 0.2)
            x = max(0, x - margin_w)
            y = max(0, y - margin_h)
            w = min(img_w - x, w + 2 * margin_w)
            h = min(img_h - y, h + 2 * margin_h)
            
            # Filter out very small regions
            if w < 32 or h < 32:
                print(f"  Region {i+1}: SKIPPED (too small: {w}x{h})")
                continue
            
            # Filter out regions that cover >60% of image
            region_area = w * h
            image_area = img_w * img_h
            if region_area / image_area > 0.6:
                print(f"  Region {i+1}: SKIPPED (too large: {region_area/image_area*100:.1f}% of image)")
                continue
            
            scaled_regions.append({
                "bbox": (x, y, w, h),
                "description": r.get("description", "potential PII"),
                "confidence": float(r.get("confidence", 0.5))
            })
            
            print(f"  Region {i+1}: bbox=({x},{y},{w},{h}) conf={scaled_regions[-1]['confidence']:.2f}")
            print(f"            desc='{scaled_regions[-1]['description']}'")
        
        print(f"\nFiltered to {len(scaled_regions)} valid regions")
        
        if not scaled_regions:
            print("\nNo regions found. Scout workflow complete (no PII detected).")
            return True
        
        # Step 2: Test cropping for first region
        print(f"\n[STEP 2] Testing crop and detection on first region...")
        test_region = scaled_regions[0]
        bbox = test_region["bbox"]
        
        # Create crop
        crop_dir = Path("artifacts") / "tmp" / "crops"
        crop_dir.mkdir(parents=True, exist_ok=True)
        crop_path = str(crop_dir / "test_crop.png")
        
        if _crop_image_region(image_path, bbox, crop_path):
            print(f"✓ Successfully cropped region to: {crop_path}")
            crop_img = Image.open(crop_path)
            print(f"  Crop size: {crop_img.size[0]} x {crop_img.size[1]}")
        else:
            print(f"✗ Failed to crop region")
            return False
        
        # Step 3: Test mask mapping (simulate with a dummy mask)
        print(f"\n[STEP 3] Testing mask mapping back to full image...")
        
        # Create a simple test mask (just a white rectangle in crop coordinates)
        test_mask = Image.new('L', crop_img.size, 0)
        from PIL import ImageDraw
        draw = ImageDraw.Draw(test_mask)
        # Draw white rectangle in center 50% of crop
        w, h = crop_img.size
        margin = 0.25
        draw.rectangle([
            (int(w*margin), int(h*margin)),
            (int(w*(1-margin)), int(h*(1-margin)))
        ], fill=255)
        
        test_mask_path = str(crop_dir / "test_mask.png")
        test_mask.save(test_mask_path)
        print(f"✓ Created test mask: {test_mask_path}")
        
        # Map back to full image
        mapped_mask_path = str(crop_dir / "test_mapped_mask.png")
        if _map_mask_to_full_image(test_mask_path, bbox, (img_w, img_h), mapped_mask_path):
            print(f"✓ Successfully mapped mask to full image: {mapped_mask_path}")
            mapped_img = Image.open(mapped_mask_path)
            print(f"  Mapped mask size: {mapped_img.size[0]} x {mapped_img.size[1]}")
            
            # Check if mask is empty
            is_empty = _mask_is_empty(mapped_mask_path)
            print(f"  Mask is empty: {is_empty}")
        else:
            print(f"✗ Failed to map mask to full image")
            return False
        
        print(f"\n{'='*60}")
        print(f"Scout and Zoom Workflow Test: PASSED")
        print(f"{'='*60}\n")
        return True
        
    except Exception as e:
        print(f"\nERROR: {e}")
        import traceback
        traceback.print_exc()
        return False


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python test_scout_zoom.py <image_path>")
        print("Example: python test_scout_zoom.py ./test_data/0000000165.png")
        sys.exit(1)
    
    image_path = sys.argv[1]
    success = asyncio.run(test_scout_workflow(image_path))
    sys.exit(0 if success else 1)
