import cv2
import numpy as np
from torchvision.transforms.v2 import ToPILImage
from PIL import Image
from transformers import pipeline
import torch
from imgutils.detect import detect_heads
from src.utils.device import determine_accelerator
topil = ToPILImage()

# 1. Initialize the filtering pipeline
device = determine_accelerator()

print("Loading AI Model...")
pipe = pipeline(
    "image-text-to-text",
    model="google/gemma-3-12b-it",
    device=device,
    torch_dtype=torch.bfloat16,
)

def crop_and_mask_characters_gradio(pil_img):
    """
    Crops character regions from an image, saves them as separate files,
    and generates binary masks for each cropped region using the Gemini 2.0 Flash Exp model.

    Args:
        uploaded_file_obj (str): The path to the input image.
    """
    img = np.array(pil_img)

    # Convert the image to grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Apply thresholding to create a binary image
    _, thresh = cv2.threshold(gray, 253, 255, cv2.THRESH_BINARY_INV)

    # Find contours in the binary image
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Create output directories if they don't exist
    # cropped_dir = Path(temp_dir) / 'cropped_dir'
    # masks_dir = cropped_dir

    # os.makedirs(cropped_dir, exist_ok=True)
    # os.makedirs(masks_dir, exist_ok=True)
    coord_info_list = []
    i = 0
    # Iterate through the contours and crop the regions
    for contour in contours:
        # Get the bounding box of the contour
        x, y, w, h = cv2.boundingRect(contour)
        if w < 256 or h < 256:  # Skip small contours
            continue

        # Crop the region
        cropped_img = img[y:y+h, x:x+w]
        
        messages = [
            {
                "role": "system",
                "content": [{"type": "text", "text": "You are a helpful assistant."}]
            },
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": topil(cropped_img)},
                    {"type": "text", "text": "You are given a black-and-white line drawing as input. Please analyze the image carefully. If the drawing contains the majority of a head or face—meaning most key facial features or the overall shape of the head are visible—respond with 'True'. Otherwise, respond with 'False'. Do not contain any punctuation or extra spaces in your answer. Just respond with 'True' or 'False'"}
                ]
            }
        ]
        result = detect_heads(topil(cropped_img))
        if len(result) == 0:
            continue
        
        output = pipe(text=messages, max_new_tokens=200)
        if output[0]["generated_text"][-1]["content"] == 'False':
            # print(f"Skipping character {i+1} as it does not contain a head or face.")
            continue
        i += 1
        # Append the coordinates to the list
        coord_info_list.append((i,(x,y,w,h)))
    return coord_info_list