import gradio as gr
import numpy as np
import torch
import cv2
from PIL import Image
import os
from dataclasses import dataclass
from typing import List, Dict, Any, Optional
from tqdm import tqdm
import time
import platform
# Assuming these are available in your environment
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
from GroundingDINO.groundingdino.util.inference import load_model, load_image, predict
from torchvision.ops import box_convert

# --- Configuration (Adjust paths as needed) ---
SAM2_CHECKPOINT = os.path.abspath("models/sam2.1_hiera_large.pt")
SAM2_MODEL_CONFIG = os.path.abspath("models/sam2.1_hiera_l.yaml")
GROUNDING_DINO_CONFIG = os.path.abspath("models/GroundingDINO_SwinT_OGC.py")
GROUNDING_DINO_CHECKPOINT = os.path.abspath("models/groundingdino_swint_ogc.pth")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
OUTPUT_DIR = "outputs"
MASKS_DIR = "generated_masks"

def open_folder():
    open_folder_path = os.path.abspath("outputs")
    if platform.system() == "Windows":
        os.startfile(open_folder_path)
    elif platform.system() == "Linux":
        os.system(f'xdg-open "{open_folder_path}"')

def open_folder_mask():
    open_folder_path = os.path.abspath("generated_masks")
    if platform.system() == "Windows":
        os.startfile(open_folder_path)
    elif platform.system() == "Linux":
        os.system(f'xdg-open "{open_folder_path}"')

# --- Model Loading ---
# Check if files exist
if not os.path.exists(SAM2_CHECKPOINT):
    print(f"Error: SAM2 checkpoint file not found at {SAM2_CHECKPOINT}")
    exit()
if not os.path.exists(SAM2_MODEL_CONFIG):
    print(f"Error: SAM2 model config file not found at {SAM2_MODEL_CONFIG}")
    exit()
if not os.path.exists(GROUNDING_DINO_CONFIG):
    print(f"Error: Grounding DINO config file not found at {GROUNDING_DINO_CONFIG}")
    exit()
if not os.path.exists(GROUNDING_DINO_CHECKPOINT):
    print(f"Error: Grounding DINO checkpoint file not found at {GROUNDING_DINO_CHECKPOINT}")
    exit()

# Build SAM2 image predictor
sam2_model = build_sam2(SAM2_MODEL_CONFIG, SAM2_CHECKPOINT, device=DEVICE)
sam2_predictor = SAM2ImagePredictor(sam2_model)

# Build grounding dino model
grounding_model = load_model(
    model_config_path=GROUNDING_DINO_CONFIG,
    model_checkpoint_path=GROUNDING_DINO_CHECKPOINT,
    device=DEVICE
)

# --- Data Classes ---

@dataclass
class BoundingBox:
    xmin: int
    ymin: int
    ymax: int
    xmax: int

    @property
    def xyxy(self) -> List[float]:
        return [self.xmin, self.ymin, self.xmax, self.ymax]

@dataclass
class DetectionResult:
    score: float
    label: str
    box: BoundingBox
    mask: Optional[np.array] = None

# --- Helper Functions ---

def mask_to_polygon(mask: np.ndarray) -> List[List[int]]:
    mask = (mask > 0).astype(np.uint8)
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return []
    largest_contour = max(contours, key=cv2.contourArea)
    epsilon = 0.001 * cv2.arcLength(largest_contour, True)
    approx_contour = cv2.approxPolyDP(largest_contour, epsilon, True)
    return approx_contour.reshape(-1, 2).tolist()

def polygon_to_mask(polygon: List, image_shape: tuple) -> np.ndarray:
    if not polygon:
        return np.zeros(image_shape, dtype=np.uint8)
    mask = np.zeros(image_shape, dtype=np.uint8)
    pts = np.array(polygon, dtype=np.int32)
    cv2.fillPoly(mask, [pts], color=(1,))
    return mask

def get_boxes_from_gdino(results: tuple) -> List[List[float]]:
    boxes = []
    for box in results[0]:
        xyxy = box.tolist()
        boxes.append(xyxy)
    return boxes
    
def refine_masks(masks: np.ndarray, image_shape: tuple) -> List[np.ndarray]:
    refined_masks = []
    for mask in masks:
        # Ensure the mask is a boolean array
        mask_bool = mask > 0

        # Resize the mask to the original image shape
        mask_resized = cv2.resize(mask_bool.astype(np.uint8) * 255, (image_shape[1], image_shape[0]))

        # Threshold the resized mask to ensure binary values (0 or 255)
        _, mask_thresholded = cv2.threshold(mask_resized, 127, 255, cv2.THRESH_BINARY)

        # Convert the mask back to boolean array (True/False)
        mask_bool = mask_thresholded > 0
        
        refined_masks.append(mask_bool)

    return refined_masks

# --- Main Processing Functions ---
def detect(
    temp_image_path: str,
    labels: List[str],
    box_threshold: float = 0.3,
    text_threshold: float = 0.25,
) -> List[DetectionResult]:
    image_source, image = load_image(temp_image_path)
    h, w, _ = image_source.shape
    
    text_prompt = ". ".join(labels) + "."
    
    boxes, confidences, phrases = predict(
        model=grounding_model,
        image=image,
        caption=text_prompt,
        box_threshold=box_threshold,
        text_threshold=text_threshold,
    )
    
    # Scale boxes by image dimensions
    boxes = boxes * torch.Tensor([w, h, w, h])
    boxes = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()
    
    detections = []
    for box, score, label in zip(boxes, confidences, phrases):
        detections.append(
            DetectionResult(
                score=score,
                label=label,
                box=BoundingBox(
                    xmin=int(box[0]),
                    ymin=int(box[1]),
                    xmax=int(box[2]),
                    ymax=int(box[3])
                )
            )
        )
    return detections

def segment(
    temp_image_path: str,
    detection_results: List[DetectionResult],
) -> List[DetectionResult]:
    image_source, _ = load_image(temp_image_path)
    sam2_predictor.set_image(image_source)

    # Get boxes in correct format
    boxes = np.array([det.box.xyxy for det in detection_results])
    
    # Use SAM2 predictor
    masks, _, _ = sam2_predictor.predict(
        point_coords=None,
        point_labels=None,
        box=boxes,
        multimask_output=False,
    )
    
    # Handle mask dimensionality
    if masks.ndim == 4:
        masks = masks.squeeze(1)
    
    # Assign masks to detection results
    for detection_result, mask in zip(detection_results, masks):
        detection_result.mask = mask > 0  # Convert to boolean mask
        
    return detection_results

def process_image(input_image, text_prompt, box_threshold=0.3, text_threshold=0.25, save_mask=False):
    try:
        input_image = Image.fromarray(input_image)
        timestamp = int(time.time())
        temp_filename = f"temp_image_{timestamp}.png"
        temp_image_path = os.path.join(OUTPUT_DIR, temp_filename)
        
        os.makedirs(OUTPUT_DIR, exist_ok=True)
        input_image.save(temp_image_path)
        
        # Detect and segment
        labels = [text_prompt]
        detections = detect(temp_image_path, labels, box_threshold, text_threshold)
        
        if not detections:
            print(f"No detections found for image with prompt '{text_prompt}'")
            return np.array(input_image), None, []
        
        detections = segment(temp_image_path, detections)
        
        # Create visualization
        output_image = annotate(input_image, detections)

        # Save the output image with a unique name
        output_image_filename = f"output_{timestamp}_{text_prompt.replace(' ', '_')}_{box_threshold:.2f}.png"
        output_image_path = os.path.join(OUTPUT_DIR, output_image_filename)
        Image.fromarray(output_image).save(output_image_path)
        
        # Handle mask generation and coordinates
        mask_image = None
        mask_coordinates = []
        
        if save_mask and detections:
            # Create binary mask
            combined_mask = np.zeros(np.array(input_image).shape[:2], dtype=np.uint8)
            for detection in detections:
                if detection.mask is not None:
                    combined_mask |= detection.mask.astype(np.uint8)
            
            # Convert to image
            mask_image = Image.fromarray(combined_mask * 255)
            
            # Get coordinates
            mask_coords = np.where(combined_mask > 0)
            mask_coordinates = list(zip(mask_coords[0].tolist(), mask_coords[1].tolist()))
            
            # Save mask
            os.makedirs(MASKS_DIR, exist_ok=True)
            mask_filename = f"mask_{timestamp}_{text_prompt.replace(' ', '_')}_{box_threshold:.2f}.png"
            mask_image.save(os.path.join(MASKS_DIR, mask_filename))
        
        # Cleanup
        os.remove(temp_image_path)
        
        return output_image, mask_image, mask_coordinates
        
    except Exception as e:
        print(f"Error processing image: {e}")
        return np.array(input_image), None, []

def annotate(image: Image.Image, detection_results: List[DetectionResult]) -> np.ndarray:
    image_cv2 = np.array(image)
    image_cv2 = cv2.cvtColor(image_cv2, cv2.COLOR_RGB2BGR)

    for detection in detection_results:
        label = detection.label
        score = detection.score
        box = detection.box
        mask = detection.mask

        color = np.random.randint(0, 256, size=3)

        if mask is not None:
            overlay = image_cv2.copy()
            mask_display = (mask * 255).astype(np.uint8)
            overlay[mask == 1] = color
            image_cv2 = cv2.addWeighted(image_cv2, 0.7, overlay, 0.3, 0)

            contours, _ = cv2.findContours(mask_display, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cv2.drawContours(image_cv2, contours, -1, color.tolist(), 2)

        cv2.rectangle(image_cv2, (box.xmin, box.ymin), (box.xmax, box.ymax), color.tolist(), 2)
        cv2.putText(
            image_cv2,
            f'{label}: {score:.2f}',
            (box.xmin, box.ymin - 10),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            color.tolist(),
            2
        )

    return cv2.cvtColor(image_cv2, cv2.COLOR_BGR2RGB)



def batch_process_images(
        input_folder,
        output_segmented_folder,
        output_masks_folder,
        text_prompt,
        box_threshold=0.3,
        text_threshold=0.25,
):
    os.makedirs(output_segmented_folder, exist_ok=True)
    os.makedirs(output_masks_folder, exist_ok=True)

    image_files = [
        f
        for f in os.listdir(input_folder)
        if os.path.isfile(os.path.join(input_folder, f))
           and f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.webp'))
    ]

    for image_file in tqdm(image_files, desc="Processing images"):
        try:
            input_image_path = os.path.join(input_folder, image_file)
            
            input_image = Image.open(input_image_path).convert("RGB")
            
            segmented_image, mask_image, _ = process_image(
                np.array(input_image),
                text_prompt,
                box_threshold,
                text_threshold,
                save_mask=True
            )
            
            # Save segmented image with unique naming
            timestamp = int(time.time())
            segmented_image_filename = f"segmented_{timestamp}_{os.path.splitext(image_file)[0]}.png"
            segmented_image_path = os.path.join(
                output_segmented_folder, segmented_image_filename
            )
            Image.fromarray(segmented_image).save(segmented_image_path)

            if mask_image:
                # Save mask image with unique naming
                mask_image_filename = f"mask_{timestamp}_{os.path.splitext(image_file)[0]}.png"
                mask_image_path = os.path.join(
                    output_masks_folder, mask_image_filename
                )
                mask_image.save(mask_image_path)
        except Exception as e:
            print(f"Error processing {image_file}: {e}")

def create_gradio_interface():
    with gr.Blocks() as interface:
        with gr.Tab("Single Image Processing"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type="numpy", label="Input Image", height=500)
                    text_prompt = gr.Textbox(
                        label="Enter object to segment. Make sure it is lowercase and each word ends with . Also plural or singular makes difference",
                        value="eye. lips. ear.",
                        placeholder="Enter object to segment. Make sure it is lowercase and each word ends with . ",
                    )
                    box_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.35,
                        label="Box Threshold",
                    )
                    text_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.25,
                        label="Text Threshold",
                    )
                    run_button = gr.Button("Process Image")
                with gr.Column():
                    output_image = gr.Image(type="numpy", label="Output Image", height=500)
                    btn_open_outputs = gr.Button("Open Outputs Folder")
                    btn_open_outputs.click(fn=open_folder)
                    output_mask = gr.Image(type="numpy", label="Output Mask", height=500)
                    btn_open_outputs = gr.Button("Open Generated Masks Folder")
                    btn_open_outputs.click(fn=open_folder_mask)
                    mask_coordinates_output = gr.Textbox(label="Mask Coordinates")

            def process_and_return_visible(input_image, prompt, box_threshold, text_threshold, save_mask):
                output_img, mask_img, mask_coords = process_image(input_image, prompt, box_threshold, text_threshold, save_mask)
                mask_coords_str = str(mask_coords)
                return output_img, mask_img, mask_coords_str

            run_button.click(
                fn=process_and_return_visible,
                inputs=[input_image, text_prompt, box_threshold, text_threshold, gr.Checkbox(True, label="Save Mask")],
                outputs=[output_image, output_mask, mask_coordinates_output],
            )

        with gr.Tab("Batch Processing"):
            with gr.Row():
                with gr.Column():
                    input_folder = gr.Textbox(label="Input Folder", placeholder="Enter path to input folder")
                    output_segmented_folder = gr.Textbox(label="Output Segmented Images Folder", placeholder="Enter path to output segmented images folder")
                    output_masks_folder = gr.Textbox(label="Output Masks Folder", placeholder="Enter path to output masks folder")
                    batch_text_prompt = gr.Textbox(
                        label="Enter object to segment. Make sure it is lowercase and each word ends with . Also plural or singular makes difference",
                        value="eye. lips. ear.",
                        placeholder="Enter object to segment (e.g., 'eyes', 'face')",
                    )
                    batch_box_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.35,
                        label="Box Threshold",
                    )
                    batch_text_threshold = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.25,
                        label="Text Threshold",
                    )
                    batch_run_button = gr.Button("Batch Process Images")
                with gr.Column():
                    batch_status = gr.Textbox(label="Status", interactive=False)

            batch_run_button.click(
                fn=lambda *args: batch_process_images(*args) or "Batch processing completed",
                inputs=[input_folder, output_segmented_folder, output_masks_folder, batch_text_prompt, batch_box_threshold, batch_text_threshold],
                outputs=batch_status
            )

    return interface

if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(share=False, inbrowser=True)