import os
import sys
from pathlib import Path
from PIL import Image

# Add core directory to python path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "core")))
# Add newcolor directory to python path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "newcolor")))

import crop
import process_images
import color_steal
from inference import run_inference

def main():
    input_image_path = r"d:\Projects\hg spaces\id - Copy\id-maker\DSC_0001.JPG"
    output_dir = Path(r"d:\Projects\hg spaces\id - Copy\id-maker\comparison_results")
    output_dir.mkdir(parents=True, exist_ok=True)
    
    print(f"Creating comparison directory: {output_dir}")
    
    # Paths for intermediate and final outputs
    cropped_path = output_dir / "temp_crop.png"
    cutout_path = output_dir / "temp_cutout.png"
    old_corrected_path = output_dir / "old_corrected.png"
    new_corrected_path = output_dir / "new_corrected.png"
    
    # Step 1: Crop the image
    print("\n--- STEP 1: Auto-cropping image ---")
    if not os.path.exists(input_image_path):
        print(f"Error: Input image not found at {input_image_path}")
        return
        
    print(f"Cropping image {input_image_path}...")
    success = crop.crop_to_4x6_opencv(input_image_path, str(cropped_path))
    if not success:
        print("Error: Cropping failed.")
        return
    print(f"Cropped image saved to {cropped_path}")
    
    # Step 2: Background Removal (RMBG)
    print("\n--- STEP 2: Removing background ---")
    print("Loading RMBG model...")
    model, device = process_images.setup_model()
    transform = process_images.get_transform()
    
    print("Running background removal...")
    cropped_img = Image.open(cropped_path)
    cutout_img = process_images.remove_background(model, cropped_img, transform)
    cutout_img.save(cutout_path, "PNG")
    print(f"Cutout image saved to {cutout_path}")
    
    # Step 3: Run the Old Mechanism (color_steal LUT)
    print("\n--- STEP 3: Running Old Color Grading Mechanism ---")
    luts = color_steal.load_trained_curves()
    if luts is None:
        print("Warning: No pre-trained curves found. Attempting to load default or fallback.")
        # Try loading from the root of workspace or core folder
        luts = color_steal.load_trained_curves(os.path.join("core", "trained_curves.npz"))
        
    if luts is not None:
        print("Applying old color grading curves...")
        old_corrected_img = color_steal.apply_to_image(luts, cutout_img)
        old_corrected_img.save(old_corrected_path, "PNG")
        print(f"Old mechanism result saved to {old_corrected_path}")
    else:
        print("Error: Could not load LUT curves for the old mechanism.")
        
    # Step 4: Run the New Mechanism (ColorUNet)
    print("\n--- STEP 4: Running New Color Correction Mechanism ---")
    model_path = r"d:\Projects\hg spaces\id - Copy\id-maker\newcolor\color_model_best.pth"
    print(f"Running inference using model {model_path}...")
    try:
        run_inference(str(cutout_path), model_path, str(new_corrected_path))
        print(f"New mechanism result saved to {new_corrected_path}")
    except Exception as e:
        print(f"Error running new color correction model: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()