#!/usr/bin/env python3
"""Regenerate PII mask from individual masks in companion folder.

This script regenerates the combined PII mask by collecting all individual
mask files from the companion folder's masks/ directory.

Usage:
    python tools/regenerate_pii_mask.py <companion_folder>

Example:
    python tools/regenerate_pii_mask.py artifacts/data/CityScapes/leftImg8bit/test/berlin/berlin_000323_000019_leftImg8bit
"""

import sys
from pathlib import Path
import numpy as np
from PIL import Image
import json


def regenerate_pii_mask(companion_folder: Path) -> None:
    """Regenerate PII mask from all individual masks in companion folder."""
    
    print(f"\n{'='*80}")
    print(f"[PII Mask] Regenerating combined PII mask")
    print(f"{'='*80}")
    print(f"Companion folder: {companion_folder}")
    
    masks_dir = companion_folder / "masks"
    if not masks_dir.exists():
        print(f"ERROR: Masks directory not found: {masks_dir}")
        return
    
    # Find the original image to get dimensions
    json_files = list(companion_folder.glob("*.json"))
    if not json_files:
        print(f"ERROR: No JSON file found in {companion_folder}")
        return
    
    json_path = json_files[0]
    with json_path.open('r') as f:
        data = json.load(f)
    
    # Get image dimensions from first step or output path
    output_path = data.get('output_path')
    if output_path and Path(output_path).exists():
        img = Image.open(output_path)
        img_width, img_height = img.size
        print(f"Image dimensions: {img_width}x{img_height}")
    else:
        print("ERROR: Cannot determine image dimensions")
        return
    
    # Create blank mask (boolean array, same as in main.py)
    combined_mask = np.zeros((img_height, img_width), dtype=bool)
    
    # Collect all PII masks (exclude traffic signs and the pii_mask itself)
    pii_masks = []
    all_mask_files = list(masks_dir.glob('*.png'))
    
    print(f"\nFound {len(all_mask_files)} mask files:")
    for mask_file in sorted(all_mask_files):
        print(f"  - {mask_file.name}")
        
        # Skip the combined PII mask itself
        if mask_file.name.endswith('_pii_mask.png'):
            print(f"    → SKIP (combined PII mask)")
            continue
        
        # Skip traffic signs (public infrastructure, not PII)
        if 'traffic_sign' in mask_file.name:
            print(f"    → SKIP (traffic sign - public infrastructure)")
            continue
        
        # Include this mask
        pii_masks.append(mask_file)
        print(f"    → INCLUDE")
    
    print(f"\nCombining {len(pii_masks)} PII mask(s) (excluding traffic signs)")
    
    # Combine all PII masks (same logic as in main.py)
    mask_count = 0
    for mask_path in pii_masks:
        try:
            mask = Image.open(mask_path).convert('L')
            
            # Resize if dimensions don't match
            if mask.size != (img_width, img_height):
                print(f"  Resizing {mask_path.name} from {mask.size} to {img_width}x{img_height}")
                mask = mask.resize((img_width, img_height), Image.NEAREST)
            
            # Convert to boolean array
            mask_arr = np.array(mask) > 127
            
            # Verify shape matches
            if mask_arr.shape != combined_mask.shape:
                print(f"  WARNING: Shape mismatch for {mask_path.name}, skipping")
                continue
            
            # Track pixels before combining
            pixels_before = combined_mask.sum()
            
            # Combine using bitwise OR (same as main.py and audit_output)
            combined_mask |= mask_arr
            
            pixels_after = combined_mask.sum()
            added_pixels = pixels_after - pixels_before
            
            print(f"  ✓ {mask_path.name}: added {added_pixels:,} pixels (total: {pixels_after:,})")
            mask_count += 1
        except Exception as e:
            print(f"  ✗ ERROR processing {mask_path.name}: {e}")
    
    # Convert boolean mask to 0-255 uint8 for saving
    combined_mask_uint8 = (combined_mask.astype(np.uint8) * 255)
    combined_mask_img = Image.fromarray(combined_mask_uint8, mode='L')
    
    # Save PII mask
    input_stem = companion_folder.name
    pii_mask_path = masks_dir / f"{input_stem}_pii_mask.png"
    combined_mask_img.save(pii_mask_path)
    
    # Print statistics
    total_pixels = combined_mask_uint8.size
    masked_pixels = np.count_nonzero(combined_mask_uint8)
    coverage_pct = 100 * masked_pixels / total_pixels if total_pixels > 0 else 0
    
    print(f"\n{'='*80}")
    print(f"SUCCESS: Created combined PII mask from {mask_count} individual mask(s)")
    print(f"Combined {mask_count} mask(s): {masked_pixels:,} pixels ({coverage_pct:.2f}% coverage)")
    print(f"Saved to: {pii_mask_path}")
    print(f"{'='*80}")
    
    # Update JSON with mask path
    data['pii_mask'] = str(pii_mask_path)
    with json_path.open('w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    print(f"Updated JSON: {json_path}")


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        return 1
    
    companion_folder = Path(sys.argv[1])
    if not companion_folder.exists():
        print(f"ERROR: Companion folder not found: {companion_folder}")
        return 1
    
    regenerate_pii_mask(companion_folder)
    return 0


if __name__ == "__main__":
    sys.exit(main())
