import json
import os
import sys
import tempfile

import gradio as gr
import numpy as np
import supervision as sv
import torch
from groundingdino.util.inference import Model as DinoModel
from PIL import Image
from segment_anything import build_sam
from segment_anything import SamAutomaticMaskGenerator
from segment_anything import SamPredictor
from supervision.detection.utils import mask_to_polygons
from supervision.detection.utils import xywh_to_xyxy

# segment anything
# Grounding DINO

sys.path.append("tag2text")

from tag2text.models import tag2text
from config import *
from utils import download_file_hf, detect, segment, show_anns, generate_tags

if not os.path.exists(abs_weight_dir):
    os.makedirs(abs_weight_dir, exist_ok=True)

sam_checkpoint = os.path.join(abs_weight_dir, sam_dict[default_sam]["checkpoint_file"])
if not os.path.exists(sam_checkpoint):
    os.system(f"wget {sam_dict[default_sam]['checkpoint_url']} -O {sam_checkpoint}")

tag2text_checkpoint = os.path.join(
    abs_weight_dir, tag2text_dict[default_tag2text]["checkpoint_file"]
)
if not os.path.exists(tag2text_checkpoint):
    os.system(
        f"wget {tag2text_dict[default_tag2text]['checkpoint_url']} -O {tag2text_checkpoint}"
    )

dino_checkpoint = os.path.join(
    abs_weight_dir, dino_dict[default_dino]["checkpoint_file"]
)
dino_config_file = os.path.join(abs_weight_dir, dino_dict[default_dino]["config_file"])
if not os.path.exists(dino_checkpoint):
    dino_repo_id = dino_dict[default_dino]["repo_id"]
    download_file_hf(
        repo_id=dino_repo_id,
        filename=dino_dict[default_dino]["config_file"],
        cache_dir=weight_dir,
    )
    download_file_hf(
        repo_id=dino_repo_id,
        filename=dino_dict[default_dino]["checkpoint_file"],
        cache_dir=weight_dir,
    )

# load model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tag2text_model = tag2text.tag2text_caption(
    pretrained=tag2text_checkpoint,
    image_size=384,
    vit="swin_b",
    delete_tag_index=delete_tag_index,
)
# threshold for tagging
# we reduce the threshold to obtain more tags
tag2text_model.threshold = 0.64
tag2text_model.to(device)
tag2text_model.eval()


sam = build_sam(checkpoint=sam_checkpoint)
sam.to(device=device)
sam_predictor = SamPredictor(sam)
sam_automask_generator = SamAutomaticMaskGenerator(sam)

grounding_dino_model = DinoModel(
    model_config_path=dino_config_file, model_checkpoint_path=dino_checkpoint
)


def process(image_path, task, prompt, box_threshold, text_threshold, iou_threshold):
    global tag2text_model, sam_predictor, sam_automask_generator, grounding_dino_model, device
    output_gallery = []
    detections = None
    metadata = {"image": {}, "annotations": []}

    try:
        # Load image
        image = Image.open(image_path)
        image_pil = image.convert("RGB")
        image = np.array(image_pil)

        # Extract image metadata
        filename = os.path.basename(image_path)
        h, w = image.shape[:2]
        metadata["image"]["file_name"] = filename
        metadata["image"]["width"] = w
        metadata["image"]["height"] = h

        # Generate tags
        if task in ["auto", "detection"] and prompt == "":
            tags, caption = generate_tags(tag2text_model, image_pil, "None", device)
            prompt = " . ".join(tags)
            print(f"Caption: {caption}")
            print(f"Tags: {tags}")

            # ToDo: Extract metadata
            metadata["image"]["caption"] = caption
            metadata["image"]["tags"] = tags

        if prompt:
            metadata["prompt"] = prompt
            print(f"Prompt: {prompt}")

        # Detect boxes
        if prompt != "":
            detections, phrases, classes = detect(
                grounding_dino_model,
                image,
                caption=prompt,
                box_threshold=box_threshold,
                text_threshold=text_threshold,
                iou_threshold=iou_threshold,
                post_process=True,
            )

            # Draw boxes
            box_annotator = sv.BoxAnnotator()
            labels = [
                f"{classes[class_id] if class_id else 'Unkown'} {confidence:0.2f}"
                for _, _, confidence, class_id, _ in detections
            ]
            image = box_annotator.annotate(
                scene=image, detections=detections, labels=labels
            )
            output_gallery.append(image)

        # Segmentation
        if task in ["auto", "segment"]:
            if detections:
                masks, scores = segment(
                    sam_predictor, image=image, boxes=detections.xyxy
                )
                detections.mask = masks
            else:
                masks = sam_automask_generator.generate(image)
                sorted_generated_masks = sorted(
                    masks, key=lambda x: x["area"], reverse=True
                )

                xywh = np.array([mask["bbox"] for mask in sorted_generated_masks])
                mask = np.array(
                    [mask["segmentation"] for mask in sorted_generated_masks]
                )
                scores = np.array(
                    [mask["predicted_iou"] for mask in sorted_generated_masks]
                )
                detections = sv.Detections(
                    xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask
                )
                # opacity = 0.4
                # mask_image, _ = show_anns_sam(masks)
                # annotated_image = np.uint8(mask_image * opacity + image * (1 - opacity))

            mask_annotator = sv.MaskAnnotator()
            mask_image = np.zeros_like(image, dtype=np.uint8)
            mask_image = mask_annotator.annotate(
                mask_image, detections=detections, opacity=1
            )
            annotated_image = mask_annotator.annotate(image, detections=detections)
            output_gallery.append(mask_image)
            output_gallery.append(annotated_image)

        # ToDo: Extract metadata
        if detections:
            id = 1
            for (xyxy, mask, confidence, class_id, _), area, box_area, score in zip(
                detections, detections.area, detections.box_area, scores
            ):
                annotation = {
                    "id": id,
                    "bbox": [int(x) for x in xyxy],
                    "box_area": float(box_area),
                }
                if class_id:
                    annotation["box_confidence"] = float(confidence)
                    annotation["label"] = classes[class_id] if class_id else "Unkown"
                if mask is not None:
                    # annotation["segmentation"] = mask_to_polygons(mask)
                    annotation["area"] = int(area)
                    annotation["predicted_iou"] = float(score)
                metadata["annotations"].append(annotation)
                id += 1

        meta_file = tempfile.NamedTemporaryFile(delete=False, suffix=".json")
        meta_file_path = meta_file.name
        with open(meta_file_path, "w") as fp:
            json.dump(metadata, fp)

        return output_gallery, meta_file_path
    except Exception as error:
        raise gr.Error(f"global exception: {error}")


title = "Annotate Anything"

with gr.Blocks(css="style.css", title=title) as demo:
    with gr.Row(elem_classes=["container"]):
        with gr.Column(scale=1):
            input_image = gr.Image(type="filepath", label="Input")
            task = gr.Dropdown(
                ["detect", "segment", "auto"], value="auto", label="task_type"
            )
            text_prompt = gr.Textbox(label="Detection Prompt")
            with gr.Accordion("Advanced parameters", open=False):
                box_threshold = gr.Slider(
                    minimum=0,
                    maximum=1,
                    value=0.3,
                    step=0.05,
                    label="Box threshold",
                    info="Hash size to use for image hashing",
                )
                text_threshold = gr.Slider(
                    minimum=0,
                    maximum=1,
                    value=0.25,
                    step=0.05,
                    label="Text threshold",
                    info="Number of history images used to find out duplicate image",
                )
                iou_threshold = gr.Slider(
                    minimum=0,
                    maximum=1,
                    value=0.5,
                    step=0.05,
                    label="IOU threshold",
                    info="Minimum similarity threshold (in percent) to consider 2 images to be similar",
                )
            run_button = gr.Button(label="Run")

        with gr.Column(scale=2):
            gallery = gr.Gallery(
                label="Generated images", show_label=False, elem_id="gallery"
            ).style(preview=True, grid=2, object_fit="scale-down")
            meta_file = gr.File(label="Metadata file")

    with gr.Row(elem_classes=["container"]):
        gr.Examples(
            [
                ["examples/dog.png", "auto", ""],
                ["examples/eiffel.png", "auto", ""],
                ["examples/eiffel.png", "segment", ""],
                ["examples/girl.png", "auto", "girl . face"],
                ["examples/horse.png", "detect", "horse"],
                ["examples/horses.jpg", "auto", "horse"],
                ["examples/traffic.jpg", "auto", ""],
            ],
            [input_image, task, text_prompt],
        )
    run_button.click(
        fn=process,
        inputs=[
            input_image,
            task,
            text_prompt,
            box_threshold,
            text_threshold,
            iou_threshold,
        ],
        outputs=[gallery, meta_file],
    )

demo.queue(concurrency_count=2).launch()