#!/usr/bin/env python3
"""
Example usage of the scout-and-zoom visualization tool.

This script demonstrates the complete scout-and-zoom process:
1. Runs PII classification to dynamically find bounding boxes in the image
2. Performs bbox-guided segmentation on each detected instance
3. Creates comprehensive visualizations of the results

Unlike the static version, this script runs the actual detection pipeline
to find PII instances and generate their segmentation masks in real-time.
"""

import json
import sys
import asyncio
import os
import subprocess
from pathlib import Path

# Add the project root to the path
sys.path.insert(0, str(Path(__file__).parent.parent))

from tools.visualize_scout_zoom import create_comprehensive_visualization
from mcp.clients.agents.main import _detect, _map_mask_to_full_image, _mask_is_empty

def run_bbox_segmentation(image_path: str, instance: dict) -> str:
    """Run bbox-guided segmentation for a single instance.
    
    Returns the path to the generated mask, or empty string if failed.
    """
    from PIL import Image
    import hashlib
    from mcp.clients.agents.main import _detect, _map_mask_to_full_image, _mask_is_empty, CANONICAL_OUTPUT
    
    detection_prompt = instance.get('det_prompt', instance.get('description', 'text.'))
    bbox = instance.get('bbox')
    
    if not bbox or not isinstance(bbox, list) or len(bbox) != 4:
        return ""
    
    try:
        # Get image dimensions
        img = Image.open(image_path)
        img_w, img_h = img.size
        
        # Validate and clamp bbox coordinates
        x, y, w, h = bbox
        if w <= 0 or h <= 0:
            return ""
        
        # Clamp coordinates to image bounds
        x = max(0, min(x, img_w - 1))
        y = max(0, min(y, img_h - 1))
        w = min(w, img_w - x)
        h = min(h, img_h - y)
        
        if w <= 0 or h <= 0:
            return ""
        
        clamped_bbox = [x, y, w, h]
        
        # Create crop from bbox region
        crop_dir = Path("artifacts") / "tmp" / "crops"
        crop_dir.mkdir(parents=True, exist_ok=True)
        
        bbox_str = f"{x}_{y}_{w}_{h}"
        crop_hash = hashlib.md5(bbox_str.encode()).hexdigest()[:8]
        crop_path = str(crop_dir / f"crop_{crop_hash}.png")
        
        # Crop the region
        crop_img = img.crop((x, y, x + w, y + h))
        crop_img.save(crop_path)
        
        # Create a minimal context for detection
        from mcp.clients.agents.context import PipelineContext
        ctx = PipelineContext(image_path=image_path, output_path=image_path)
        
        # Run detection on the cropped region
        prev_backend = os.getenv("DETECTOR_BACKEND")
        try:
            os.environ["DETECTOR_BACKEND"] = "grounded-sam-2"
            crop_detect_result = asyncio.run(_detect(ctx, crop_path, detection_prompt))
        finally:
            if prev_backend:
                os.environ["DETECTOR_BACKEND"] = prev_backend
            else:
                os.environ.pop("DETECTOR_BACKEND", None)
        
        # Extract crop mask
        crop_mask_path = crop_detect_result.get("mask_path")
        if crop_mask_path and Path(crop_mask_path).exists() and not _mask_is_empty(crop_mask_path):
            # Map crop mask back to full image coordinates
            mask_dir = Path(image_path).parent / "masks"
            mask_dir.mkdir(parents=True, exist_ok=True)
            full_mask_path = str(mask_dir / f"mapped_{crop_hash}_{Path(crop_mask_path).name}")
            
            if _map_mask_to_full_image(crop_mask_path, clamped_bbox, (img_w, img_h), full_mask_path):
                return full_mask_path
    
    except Exception as e:
        print(f"Segmentation failed for instance: {e}")
    
    return ""

def main():
    # Example data - will be dynamically generated
    image_path = "test_data/other/0000000165.png"
    
    # Set up canonical output path (same as main pipeline)
    src = Path(image_path)
    if not src.exists() or not src.is_file():
        print(f"[error] Input image not found: {image_path}", file=sys.stderr, flush=True)
        return
    
    # Set canonical output to artifacts/{relative_parent}/{input_stem}.jpeg
    input_stem = src.stem
    try:
        relative_parent = src.parent.relative_to(Path.cwd())
    except ValueError:
        relative_parent = src.parent.name if src.parent.name else "."
    
    output_dir = Path("artifacts") / relative_parent
    canonical_output = output_dir / f"{input_stem}.jpeg"
    canonical_output.parent.mkdir(parents=True, exist_ok=True)
    
    # Copy input to canonical output
    if src.resolve() != canonical_output.resolve():
        import shutil
        shutil.copyfile(src, canonical_output)
    
    print(f"Running scout-and-zoom test on: {image_path}")
    print(f"Canonical output: {canonical_output}")
    
    # Initialize pipeline context
    from mcp.clients.agents.context import PipelineContext
    ctx = PipelineContext(
        image_path=str(canonical_output),
        output_path=str(canonical_output),
    )
    
    # Step 1: Run PII classification to find bounding boxes
    print("\n=== Step 1: Classifying PII instances ===")
    try:
        # Run the classify_pii tool
        result = subprocess.run([
            sys.executable, "-X", "utf8", "tools/classify_pii.py", str(canonical_output)
        ], capture_output=True, text=True, cwd=Path(__file__).parent.parent)
        
        if result.returncode != 0:
            print(f"Classification failed: {result.stderr}")
            return
        
        # Parse the JSON output - the entire stdout should be JSON
        json_text = result.stdout.strip()
        
        # Remove any debug lines that come before the JSON
        lines = json_text.split('\n')
        json_start = 0
        for i, line in enumerate(lines):
            if line.strip().startswith('{'):
                json_start = i
                break
        
        json_text = '\n'.join(lines[json_start:])
        
        try:
            classify_result = json.loads(json_text)
        except json.JSONDecodeError as e:
            print(f"Failed to parse JSON: {e}")
            print("JSON text:", repr(json_text[:200]))
            return
            
        instances = classify_result.get("instances", [])
        print(f"Found {len(instances)} PII instances")
        
        if not instances:
            print("No PII instances found, exiting")
            return
            
    except Exception as e:
        print(f"Classification failed: {e}")
        return
    
    # Step 2: Run bbox-guided segmentation for each instance
    print("\n=== Step 2: Running bbox-guided segmentation ===")
    
    # Process each instance to get its segmentation mask
    processed_instances = []
    mask_paths = []
    
    for i, instance in enumerate(instances):
        print(f"Processing instance {i+1}/{len(instances)}: {instance.get('description', 'Unknown')[:50]}...")
        
        # Run bbox-guided segmentation
        mask_path = run_bbox_segmentation(str(canonical_output), instance)
        
        if mask_path:
            mask_paths.append(mask_path)
            processed_instances.append(instance)
            print(f"  -> Generated mask: {Path(mask_path).name}")
        else:
            print(f"  -> Failed to generate mask for instance {i+1}")
    
    if not processed_instances or not mask_paths:
        print("No instances were successfully segmented")
        return
    
    print(f"\nSuccessfully processed {len(processed_instances)} instances with {len(mask_paths)} masks")
    
    # Step 3: Create visualizations
    print("\n=== Step 3: Creating visualizations ===")
    
    # Create comprehensive visualization
    output_files = create_comprehensive_visualization(
        str(canonical_output),
        processed_instances,
        mask_paths,
        "artifacts/visualizations"
    )
    
    print("\nVisualization files created:")
    for viz_type, path in output_files.items():
        print(f"  {viz_type}: {path}")
    
    print("\nScout-and-zoom test completed successfully!")
    print(f"- Found {len(instances)} PII instances")
    print(f"- Successfully segmented {len(processed_instances)} instances")
    print(f"- Generated {len(output_files)} visualization files")


if __name__ == '__main__':
    main()