import os
import tempfile
import zipfile
import shutil # For make_archive
import uuid
from PIL import Image, ImageDraw
from psd_tools import PSDImage
from psd_tools.api.layers import PixelLayer
import gradio as gr
import traceback # For printing stack traces
import subprocess

def install(package):
    subprocess.check_call([os.sys.executable, "-m", "pip", "install", package])

install("timm")
install("pydantic==2.10.6")
install("dghs-imgutils==0.15.0")
install("onnxruntime >= 1.17.0")
install("psd_tools==1.10.7")
# os.environ["no_proxy"] = "localhost,127.0.0.1,::1"
# --- Attempt to import the actual function (Detector 1) ---
# Let's keep the original import name as requested in the previous version
try:
    # Assuming this is the intended "Detector 1"
    from src.wise_crop.detect_and_crop import crop_and_mask_characters_gradio
    detector_1_available = True
    print("Successfully imported 'crop_and_mask_characters_gradio' as Detector 1.")
except ImportError:
    detector_1_available = False
    print("Warning: Could not import 'crop_and_mask_characters_gradio'. Using dummy function for Detector 1.")
    # Define a dummy version for Detector 1 if import fails
    def crop_and_mask_characters_gradio(image_pil: Image.Image):
        """Dummy function 1 if import fails."""
        print("Using DUMMY Detector 1.")
        if image_pil is None: return []
        width, height = image_pil.size
        boxes = [
            (0, (int(width * 0.1), int(height * 0.1), int(width * 0.3), int(height * 0.4))),
            (1, (int(width * 0.6), int(height * 0.5), int(width * 0.25), int(height * 0.35))),
        ]
        valid_boxes = []
        for i, (x, y, w, h) in boxes:
            x1, y1, x2, y2 = max(0, x), max(0, y), min(width, x + w), min(height, y + h)
            if x2 - x1 > 0 and y2 - y1 > 0: valid_boxes.append((i, (x1, y1, x2 - x1, y2 - y1)))
        return valid_boxes

# from src.oskar_crop.detect_and_crop import process_single_image as detector_2_function
try:
    # Assuming this is the intended "Detector 2"
    # Note: Renamed the import alias to avoid conflict if both imports succeed.
    # The function call inside process_lineart still uses crop_and_mask_characters_gradio_2
    from src.oskar_crop.detect_and_crop import process_single_image as detector_2_function
    detector_2_available = True
    print("Successfully imported 'process_single_image' as Detector 2.")
    # Define the function name used in process_lineart
    def crop_and_mask_characters_gradio_2(image_pil: Image.Image):
        return detector_2_function(image_pil)

except ImportError:
    detector_2_available = False
    print("Warning: Could not import 'process_single_image'. Using dummy function for Detector 2.")
    # Define a dummy version for Detector 2 if import fails
    # --- Define the SECOND Dummy Detection Function ---
    def crop_and_mask_characters_gradio_2(image_pil: Image.Image):
        """
        SECOND Dummy function to simulate detecting objects and returning bounding boxes.
        Returns different results than the first function.
        """
        print("Using DUMMY Detector 2.")
        if image_pil is None:
            return []

        width, height = image_pil.size
        print(f"Dummy detection 2 running on image size: {width}x{height}")

        # Define DIFFERENT fixed bounding boxes for demonstration
        boxes = [
            (0, (int(width * 0.05), int(height * 0.6), int(width * 0.4), int(height * 0.3))), # Bottom-leftish, wider
            (1, (int(width * 0.7), int(height * 0.1), int(width * 0.20), int(height * 0.25))), # Top-rightish, smaller
            (2, (int(width * 0.4), int(height * 0.4), int(width * 0.15), int(height * 0.15))), # Center-ish, very small
        ]

        # Basic validation
        valid_boxes = []
        for i, (x, y, w, h) in boxes:
            x1 = max(0, x)
            y1 = max(0, y)
            x2 = min(width, x + w)
            y2 = min(height, y + h)
            new_w = x2 - x1
            new_h = y2 - y1
            if new_w > 0 and new_h > 0:
                valid_boxes.append((i, (x1, y1, new_w, new_h)))

        print(f"Dummy detection 2 found {len(valid_boxes)} boxes.")
        return valid_boxes


# --- Helper Function (make_lineart_transparent - unchanged) ---
def make_lineart_transparent(lineart_path, threshold=200):
    """Converts a lineart image file to a transparent RGBA PIL Image."""
    try:
        # Ensure we handle potential pathlib objects if Gradio passes them
        lineart_gray = Image.open(str(lineart_path)).convert('L')
        w, h = lineart_gray.size
        lineart_rgba = Image.new('RGBA', (w, h), (0, 0, 0, 0))
        gray_pixels = lineart_gray.load()
        rgba_pixels = lineart_rgba.load()
        for y in range(h):
            for x in range(w):
                gray_val = gray_pixels[x, y]
                alpha = 255 - gray_val
                if gray_val < threshold :
                     rgba_pixels[x, y] = (0, 0, 0, alpha)
                else:
                     rgba_pixels[x, y] = (0, 0, 0, 0)
        return lineart_rgba
    except FileNotFoundError:
        print(f"Helper Error: Image file not found at {lineart_path}")
        # Return a blank transparent image or None? Returning None is clearer.
        return None
    except Exception as e:
        print(f"Helper Error processing image {lineart_path}: {e}")
        return None

# --- Main Processing Function (modified for better error handling with PIL) ---
def process_lineart(input_pil_or_path, detector_choice): # Input can be PIL or path from examples
    """
    Processes the input lineart image using the selected detector.
    Detects objects (e.g., characters based on head/face), crops them,
    provides a gallery of crops, a ZIP file of crops, and a PSD file
    with the original lineart (made transparent) and bounding boxes.
    """
    # --- Initialize variables ---
    input_pil_image = None
    temp_input_path = None
    using_temp_input_path = False
    status_updates = ["Status: Initializing..."]
    psd_output_path = None # Initialize to None
    zip_output_path = None # Initialize to None
    cropped_images_for_gallery = [] # Initialize to empty list

    try:
        # --- Handle Input ---
        if input_pil_or_path is None:
            gr.Warning("Please upload a PNG image or select an example.")
            return [], None, None, "Status: No image provided."

        print(f"Input type: {type(input_pil_or_path)}")
        print(f"Input value: {input_pil_or_path}")
        # Check if input is already a PIL image (from upload) or a path (from examples)
        if isinstance(input_pil_or_path, Image.Image):
            input_pil_image = input_pil_or_path
            print("Processing PIL image from upload.")
            # Create a temporary path for make_lineart_transparent if needed later
            temp_input_fd, temp_input_path = tempfile.mkstemp(suffix=".png")
            os.close(temp_input_fd)
            input_pil_image.save(temp_input_path, "PNG")
            using_temp_input_path = True
        elif isinstance(input_pil_or_path, str) and os.path.exists(input_pil_or_path):
            print(f"Processing image from file path: {input_pil_or_path}")
            try:
                input_pil_image = Image.open(input_pil_or_path)
                # Use the example path directly for make_lineart_transparent
                temp_input_path = input_pil_or_path
                using_temp_input_path = False # Don't delete the example file later
            except Exception as e:
                 status_updates.append(f"ERROR: Could not open image file from path '{input_pil_or_path}': {e}")
                 print(status_updates[-1])
                 return [], None, None, "\n".join(status_updates) # Return error status
        else:
             status_updates.append(f"ERROR: Invalid input type received: {type(input_pil_or_path)}. Expected PIL image or file path.")
             print(status_updates[-1])
             return [], None, None, "\n".join(status_updates) # Return error status

        # --- Ensure RGBA and get dimensions ---
        try:
            input_pil_image = input_pil_image.convert("RGBA")
            width, height = input_pil_image.size
        except Exception as e:
             status_updates.append(f"ERROR: Could not process input image (convert/get size): {e}")
             print(status_updates[-1])
             # Clean up temp file if created before error
             if using_temp_input_path and temp_input_path and os.path.exists(temp_input_path):
                 try: os.remove(temp_input_path)
                 except Exception as e_rem: print(f"Warning: Could not remove temp input file {temp_input_path}: {e_rem}")
             return [], None, None, "\n".join(status_updates) # Return error status

        status_updates = [f"Status: Processing started using {detector_choice}."] # Reset status
        print("Starting processing...")

        # --- 1. Detect Objects (Conditional) ---
        print(f"Selected detector: {detector_choice}")
        if detector_choice == "Detector 1":
            if not detector_1_available:
                 status_updates.append("Warning: Using DUMMY Detector 1.")
            boxes_info = crop_and_mask_characters_gradio(input_pil_image)
        elif detector_choice == "Detector 2":
            if not detector_2_available:
                 status_updates.append("Warning: Using DUMMY Detector 2.")
            boxes_info = crop_and_mask_characters_gradio_2(input_pil_image)
        else:
            # This case should ideally not happen with Radio buttons, but good for safety
            status_updates.append(f"ERROR: Invalid detector choice received: {detector_choice}")
            print(status_updates[-1])
            # Clean up temp file if created before error
            if using_temp_input_path and temp_input_path and os.path.exists(temp_input_path):
                try: os.remove(temp_input_path)
                except Exception as e_rem: print(f"Warning: Could not remove temp input file {temp_input_path}: {e_rem}")
            return [], None, None, "\n".join(status_updates) # Return error status

        if not boxes_info:
            gr.Warning("No objects detected.")
            status_updates.append("No objects detected.")
            # Clean up temp file if created
            if using_temp_input_path and temp_input_path and os.path.exists(temp_input_path):
                try: os.remove(temp_input_path)
                except Exception as e_rem: print(f"Warning: Could not remove temp input file {temp_input_path}: {e_rem}")
            return [], None, None, "\n".join(status_updates)

        status_updates.append(f"Detected {len(boxes_info)} objects.")
        print(f"Detected boxes: {boxes_info}")

        # --- Temporary file paths (partially adjusted) ---
        temp_dir_for_outputs = tempfile.gettempdir()
        unique_id = uuid.uuid4().hex[:8]
        zip_base_name = os.path.join(temp_dir_for_outputs, f"cropped_images_{unique_id}")
        zip_output_path = f"{zip_base_name}.zip" # Path for the final zip file
        psd_output_path = os.path.join(temp_dir_for_outputs, f"lineart_boxes_{unique_id}.psd")
        # temp_input_path is already handled above based on input source

        # --- 2. Crop Images and Prepare for ZIP ---
        with tempfile.TemporaryDirectory() as temp_crop_dir:
            print(f"Saving cropped images to temporary directory: {temp_crop_dir}")
            for i, (x, y, w, h) in boxes_info:
                # Ensure box coordinates are within image bounds
                x1, y1 = max(0, x), max(0, y)
                x2, y2 = min(width, x + w), min(height, y + h)
                box = (x1, y1, x2, y2)
                if box[2] > box[0] and box[3] > box[1]: # Check if width and height are positive
                    try:
                        cropped_img = input_pil_image.crop(box)
                        cropped_images_for_gallery.append(cropped_img)
                        crop_filename = os.path.join(temp_crop_dir, f"cropped_{i}.png")
                        cropped_img.save(crop_filename, "PNG")
                    except Exception as e:
                        print(f"Error cropping or saving box {i} with coords {box}: {e}")
                        status_updates.append(f"Warning: Error processing crop {i}.")
                else:
                     print(f"Skipping invalid box {i} with coords {box}")
                     status_updates.append(f"Warning: Skipped invalid crop dimensions for box {i}.")


            # --- 3. Create ZIP File ---
            # Check if any PNG files were actually created in the temp dir
            if any(f.endswith(".png") for f in os.listdir(temp_crop_dir)):
                print(f"Creating ZIP file: {zip_output_path} from {temp_crop_dir}")
                try:
                    shutil.make_archive(zip_base_name, 'zip', temp_crop_dir)
                    status_updates.append("Cropped images ZIP created.")
                    # zip_output_path is already correctly set
                except Exception as e:
                    print(f"Error creating ZIP file: {e}")
                    status_updates.append("Error: Failed to create ZIP file.")
                    zip_output_path = None # Indicate failure
            else:
                print("No valid cropped images were saved, skipping ZIP creation.")
                status_updates.append("Skipping ZIP creation (no valid crops).")
                zip_output_path = None # No zip file to provide

        # --- 4. Prepare PSD Layers ---
        #   a) Line Layer (Use the temp_input_path which is either the original example path or a temp copy)
        print(f"Using image path for transparent layer: {temp_input_path}")
        line_layer_pil = make_lineart_transparent(temp_input_path)
        if line_layer_pil is None:
            status_updates.append("Error: Failed to create transparent lineart layer.")
            print(status_updates[-1])
            # Don't create PSD if lineart failed, return current results
            # Clean up temp file if created
            if using_temp_input_path and temp_input_path and os.path.exists(temp_input_path):
                try: os.remove(temp_input_path)
                except Exception as e_rem: print(f"Warning: Could not remove temp input file {temp_input_path}: {e_rem}")
            return cropped_images_for_gallery, zip_output_path, None, "\n".join(status_updates) # Return None for PSD

        status_updates.append("Transparent lineart layer created.")

        #   b) Box Layer
        box_layer_pil = Image.new('RGBA', (width, height), (255, 255, 255, 255)) # White background
        draw = ImageDraw.Draw(box_layer_pil)
        for i, (x, y, w, h) in boxes_info:
            # Use validated coords again, ensure they are within bounds
            x1, y1 = max(0, x), max(0, y)
            x2, y2 = min(width, x + w), min(height, y + h)
            if x2 > x1 and y2 > y1: # Check validity again just in case
                 rect = [(x1, y1), (x2, y2)]
                 # Changed to fill for solid boxes, yellow fill, semi-transparent
                 draw.rectangle(rect, fill=(255, 255, 0, 128))
        status_updates.append("Bounding box layer created.")

        # --- 5. Create PSD File ---
        print(f"Creating PSD file: {psd_output_path}")
        # Double check layer sizes before creating PSD object
        if line_layer_pil.size != (width, height) or box_layer_pil.size != (width, height):
             size_error_msg = (f"Error: Layer size mismatch during PSD creation. "
                             f"Line: {line_layer_pil.size}, Box: {box_layer_pil.size}, "
                             f"Expected: {(width, height)}")
             status_updates.append(size_error_msg)
             print(size_error_msg)
             # Clean up temp file if created
             if using_temp_input_path and temp_input_path and os.path.exists(temp_input_path):
                 try: os.remove(temp_input_path)
                 except Exception as e_rem: print(f"Warning: Could not remove temp input file {temp_input_path}: {e_rem}")
             return cropped_images_for_gallery, zip_output_path, None, "\n".join(status_updates) # No PSD

        try:
            psd = PSDImage.new(mode='RGBA', size=(width, height))
            # Add layers (order matters for visibility in PSD viewers)
            # Base layer is transparent by default with RGBA
            psd.append(PixelLayer.frompil(line_layer_pil, layer_name='line', top=0, left=0))
            psd.append(PixelLayer.frompil(box_layer_pil, layer_name='box', top=0, left=0))
            psd.save(psd_output_path)
            status_updates.append("PSD file created.")
        except Exception as e:
            print(f"Error saving PSD file: {e}")
            traceback.print_exc()
            status_updates.append("Error: Failed to save PSD file.")
            psd_output_path = None # Indicate failure


        print("Processing finished.")
        status_updates.append("Success!")
        final_status = "\n".join(status_updates)

        # Return all paths, even if None (Gradio handles None for File output)
        return cropped_images_for_gallery, zip_output_path, psd_output_path, final_status

    except Exception as e:
        print(f"An unexpected error occurred in process_lineart: {e}")
        traceback.print_exc()
        status_updates.append(f"FATAL ERROR: {e}")
        final_status = "\n".join(status_updates)
        # Return empty/None outputs and the error status
        # Ensure cleanup happens even on fatal error
        if using_temp_input_path and temp_input_path and os.path.exists(temp_input_path):
            try:
                os.remove(temp_input_path)
                print(f"Cleaned up temporary input file due to error: {temp_input_path}")
            except Exception as e_rem:
                print(f"Warning: Could not remove temp input file {temp_input_path} during error handling: {e_rem}")
        return [], None, None, final_status # Return safe defaults

    finally:
        # --- Final Cleanup (Only removes temp input if created from upload) ---
        if using_temp_input_path and temp_input_path and os.path.exists(temp_input_path):
            try:
                os.remove(temp_input_path)
                print(f"Cleaned up temporary input file: {temp_input_path}")
            except Exception as e_rem:
                # This might happen if the file was already removed in an error block
                print(f"Notice: Could not remove temp input file {temp_input_path} in finally block (may already be removed): {e_rem}")


# --- Gradio Interface Definition (modified) ---
css = '''
        .custom-gallery { 
            height: 500px !important; 
            width: 100%; 
            margin: 10px auto; 
            padding: 0px; 
            overflow-y: auto !important; 
        }
        '''
with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    gr.Markdown("# Webtoon Lineart Cropper with Filtering by Head-or-Face Detection")
    gr.Markdown("Upload a PNG lineart image of your webtoon and automatically crop the character's face or head included region. "
                "This demo leverages some detectors to precisely detect and isolate characters. "
                "The app will display cropped objects, provide a ZIP of cropped PNGs, "
                "and a PSD file with transparent lineart and half-transparent yellow-filled box layers. "
                "We provide two detectors to choose from, each with different filtering methods. ")
    gr.Markdown("- **Detector 1**: Uses [`imgutils.detect`](https://github.com/deepghs/imgutils/tree/main/imgutils/detect) and VLM-based filtering with [`google/gemma-3-12b-it`](https://huggingface.co/google/gemma-3-12b-it)")
    gr.Markdown("- **Detector 2**: Uses [`imgutils.detect`](https://github.com/deepghs/imgutils/tree/main/imgutils/detect) and tag-based filtering with [`SmilingWolf/wd-eva02-large-tagger-v3`](https://huggingface.co/SmilingWolf/wd-eva02-large-tagger-v3)")
    gr.Markdown("**Note 1:** The app may take a few seconds to process the image, depending on the size and number of characters detected. The example image below is a lineart PNG file created synthetically from images on [Danbooru](https://danbooru.donmai.us/posts?page=1&tags=dragon_ball_z) after [lineart extraction](https://huggingface.co/spaces/carolineec/informativedrawings).")
    gr.Markdown("**Note 2:** This demo is developed by [Kakao Entertainment](https://kakaoent.com/)'s AI Lab for research purposes, specifically designed to preprocess webtoon image data and is also not intended for production use. It is a research prototype and may not be suitable for all use cases. Please use it at your own risk.")

    with gr.Row():
        with gr.Column(scale=1):
            # Input type remains 'filepath' to handle examples cleanly.
            image_input = gr.Image(type="filepath", label="Upload Lineart PNG or Select Example", image_mode='RGBA', height=400)

            detector_choice_radio = gr.Radio(
                choices=["Detector 1", "Detector 2"],
                label="Choose Detection Function",
                value="Detector 1" # Default value
            )
            process_button = gr.Button("Process Uploaded/Modified Image", variant="primary")
            status_output = gr.Textbox(label="Status", interactive=False, lines=8) # Increased lines slightly more

        with gr.Column(scale=3):
            gr.Markdown("### Cropped Objects")
            # Setting height explicitly can sometimes help layout.
            gallery_output = gr.Gallery(label="Detected Objects (Cropped)", elem_id="gallery_crops", columns=4, height=500, interactive=False, elem_classes="custom-gallery")  # object_fit="contain")
            with gr.Row():
                zip_output = gr.File(label="Download Cropped Images (ZIP)")
                psd_output = gr.File(label="Download PSD (Lineart + Boxes)")

    # --- Add Examples ---
    # IMPORTANT: Make sure 'sample_img.png' exists in the same directory
    #            as this script, or provide the correct relative/absolute path.
    #            Also ensure the image is a valid PNG.
    example_image_path = "./sample_img/sample_danbooru_dragonball.png"
    if os.path.exists(example_image_path):
         gr.Examples(
            examples=[
                [example_image_path, "Detector 1"],
                [example_image_path, "Detector 2"] # Add example for detector 2 as well
                ],
            # Inputs that the examples populate
            inputs=[image_input, detector_choice_radio],
            # Outputs that are updated when an example is clicked AND run_on_click=True
            outputs=[gallery_output, zip_output, psd_output, status_output],
            # The function to call when an example is clicked
            fn=process_lineart,
            # Make clicking an example automatically run the function
            run_on_click=True,
            label="Click Example to Run Automatically", # Updated label
            cache_examples=True, # Disable caching to ensure fresh processing
            cache_mode="lazy",
         )
    else:
         gr.Markdown(f"**(Note:** Could not find `{example_image_path}` for examples. Please create it or ensure it's in the correct directory.)")


    # --- Button Click Handler (for manual uploads/changes) ---
    process_button.click(
        fn=process_lineart,
        inputs=[image_input, detector_choice_radio],
        outputs=[gallery_output, zip_output, psd_output, status_output]
    )

# --- Launch the Gradio App ---
if __name__ == "__main__":
    # Create a dummy sample image if it doesn't exist for testing
    if not os.path.exists("./sample_img/sample_danbooru_dragonball.png"):
        print("Creating a dummy 'sample_danbooru_dragonball.png' for demonstration.")
        try:
            img = Image.new('L', (300, 200), color=255) # White background (grayscale)
            draw = ImageDraw.Draw(img)
            # Draw some black lines/shapes
            draw.line((30, 30, 270, 30), fill=0, width=2)
            draw.rectangle((50, 50, 150, 150), outline=0, width=3)
            draw.ellipse((180, 70, 250, 130), outline=0, width=3)
            img.save("./sample_img/sample_danbooru_dragonball.png", "PNG")
            print("Dummy 'sample_danbooru_dragonball.png' created.")
        except Exception as e:
            print(f"Warning: Failed to create dummy sample image: {e}")

    demo.launch()
    # ssr_mode=False