import os
import torch
import numpy as np
import cv2
import requests
import gradio as gr
from segment_anything import sam_model_registry, SamPredictor
import supervision as sv

# ------------------------------
# 1. Setup & Model Loading
# ------------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device used: {device}")

model_type = "vit_b"
checkpoint_path = "sam_vit_b_01ec64.pth"

# Download model if needed
if not os.path.exists(checkpoint_path):
    print("Downloading SAM checkpoint...")
    url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
    r = requests.get(url, stream=True)
    with open(checkpoint_path, "wb") as f:
        for chunk in r.iter_content(1024 * 1024):
            if chunk:
                f.write(chunk)

sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
sam.to(device)

# Note: We kept the model in float32. 
# Using sam.half() often causes runtime type mismatch errors with SamPredictor 
# unless the input image is explicitly cast to half-precision manually.
predictor = SamPredictor(sam)

# ------------------------------
# 2. The API Function
# ------------------------------
def run_sam_api(image_url, box_coords):
    """
    Args:
        image_url (str): HTTP link to the image.
        box_coords (list): A list of 4 integers [x1, y1, x2, y2].
    """
    print(f"Received URL: {image_url}")
    print(f"Received Box: {box_coords}")

    # 1. Download the Image
    try:
        # Use a user-agent to avoid 403 Forbidden errors on some sites
        headers = {'User-Agent': 'Mozilla/5.0'}
        resp = requests.get(image_url, stream=True, headers=headers).raw
        image_array = np.asarray(bytearray(resp.read()), dtype="uint8")
        image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
        
        if image is None:
            raise ValueError("Could not decode image.")
            
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # Convert to RGB
    except Exception as e:
        print(f"Error downloading image: {e}")
        raise gr.Error(f"Failed to load image: {e}")

    # 2. Set Image for SAM
    predictor.set_image(image)

    # 3. Prepare Box
    # Ensure it's a numpy array: [x_min, y_min, x_max, y_max]
    input_box = np.array(box_coords)

    # 4. Predict
    # SAM expects box shape (1, 4)
    masks, scores, logits = predictor.predict(
        point_coords=None,
        point_labels=None,
        box=input_box[None, :], 
        multimask_output=False
    )

    # 5. Annotate / Visualize
    # Convert SAM masks to Supervision Detections
    # masks shape is (1, H, W), we need to ensure it fits supervision expectations
    detections = sv.Detections(
        xyxy=sv.mask_to_xyxy(masks=masks),
        mask=masks
    )
    
    # Create annotators
    mask_annotator = sv.MaskAnnotator(color=sv.Color.RED)
    box_annotator = sv.BoxAnnotator(color=sv.Color.RED)

    # Apply annotations
    annotated_image = mask_annotator.annotate(scene=image.copy(), detections=detections)
    annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections)
        
    return annotated_image

demo = gr.Interface(
    fn=run_sam_api,
    inputs=[
        gr.Textbox(label="Image URL", placeholder="http://..."), 
        gr.JSON(label="Box Coords [x1, y1, x2, y2]", value=[100, 100, 200, 200])
    ],
    outputs=gr.Image(type="numpy", label="Segmented Output"),
    title="SAM API via Gradio",
    description="Send an image URL and bounding box coordinates to segment objects.",
    api_name="predict_api"
)

if __name__ == "__main__":
    demo.queue().launch(share=True)