import os
import json
import math
import random
import shutil
import tempfile
from urllib.parse import urlparse
from PIL import Image
import cv2
import numpy as np
import gradio as gr
from roboflow import Roboflow


def parse_roboflow_url(url):
    """Extract workspace, project name, and version from a Roboflow URL."""
    parsed = urlparse(url)
    parts = parsed.path.strip('/').split('/')
    # Expect at least [workspace, project, ..., version]
    workspace = parts[0]
    project = parts[1]
    try:
        version = int(parts[-1])
    except ValueError:
        version = int(parts[-2])
    return workspace, project, version


def convert_seg_to_bbox(api_key, dataset_url):
    # Initialize Roboflow client
    rf = Roboflow(api_key=api_key)
    workspace, project_name, version = parse_roboflow_url(dataset_url)
    project = rf.workspace(workspace).project(project_name)
    version_obj = project.version(version)

    # Download the segmentation dataset in COCO format
    dataset = version_obj.download("coco-segmentation")
    root = dataset.location  # root of downloaded dataset

    # Load COCO train annotations
    ann_dir = os.path.join(root, "coco-annotations")
    ann_file = os.path.join(ann_dir, "train.json")
    with open(ann_file, 'r') as f:
        coco = json.load(f)
    images_info = {img['id']: img for img in coco['images']}

    # Map original category IDs to contiguous YOLO class indices
    categories = coco.get('categories', [])
    cat_ids = sorted(cat['id'] for cat in categories)
    id_to_index = {cid: idx for idx, cid in enumerate(cat_ids)}

    # Prepare output directories for YOLOv8 dataset
    out_root = tempfile.mkdtemp(prefix="yolov8_")
    img_out = os.path.join(out_root, "images")
    lbl_out = os.path.join(out_root, "labels")
    os.makedirs(img_out, exist_ok=True)
    os.makedirs(lbl_out, exist_ok=True)

    # Build YOLO annotation strings grouped by image
    annos = {}
    for anno in coco['annotations']:
        img_id = anno['image_id']
        poly = anno['segmentation'][0]
        xs = poly[0::2]
        ys = poly[1::2]
        x_min, x_max = min(xs), max(xs)
        y_min, y_max = min(ys), max(ys)
        width = x_max - x_min
        height = y_max - y_min
        cx = x_min + width / 2
        cy = y_min + height / 2
        info = images_info[img_id]
        img_w, img_h = info['width'], info['height']
        cxn = cx / img_w
        cyn = cy / img_h
        wnorm = width / img_w
        hnorm = height / img_h
        cls_idx = id_to_index[anno['category_id']]
        line = f"{cls_idx} {cxn:.6f} {cyn:.6f} {wnorm:.6f} {hnorm:.6f}"
        annos.setdefault(img_id, []).append(line)

    # Determine train images directory
    train_img_dir = os.path.join(root, "train", "images")
    if not os.path.isdir(train_img_dir):
        train_img_dir = os.path.join(root, "train")

    # Map filenames to image IDs
    name_to_id = {img['file_name']: img['id'] for img in coco['images']}

    # Copy images and write YOLO label files
    for fname, img_id in name_to_id.items():
        src_img = os.path.join(train_img_dir, fname)
        if not os.path.isfile(src_img):
            continue
        dst_img = os.path.join(img_out, fname)
        shutil.copy(src_img, dst_img)
        lbl_path = os.path.join(lbl_out, os.path.splitext(fname)[0] + ".txt")
        with open(lbl_path, 'w') as lf:
            for line in annos.get(img_id, []):
                lf.write(line + '\n')

    # Prepare before/after example galleries
    before_imgs, after_imgs = [], []
    example_files = random.sample(list(name_to_id.keys()), min(5, len(name_to_id)))

    for fname in example_files:
        src_img = os.path.join(train_img_dir, fname)
        img = cv2.cvtColor(cv2.imread(src_img), cv2.COLOR_BGR2RGB)

        # Overlay segmentation polygons
        seg_vis = img.copy()
        img_id = name_to_id[fname]
        for anno in coco['annotations']:
            if anno['image_id'] != img_id:
                continue
            poly = anno['segmentation'][0]
            pts = np.array(poly, dtype=np.int32).reshape(-1, 2)
            cv2.polylines(seg_vis, [pts], True, (255, 0, 0), 2)

        # Overlay bounding boxes
        box_vis = img.copy()
        for line in annos.get(img_id, []):
            _, cxn, cyn, wnorm, hnorm = line.split()
            cxn, cyn, wnorm, hnorm = map(float, (cxn, cyn, wnorm, hnorm))
            iw, ih = images_info[img_id]['width'], images_info[img_id]['height']
            w0 = int(wnorm * iw)
            h0 = int(hnorm * ih)
            x0 = int(cxn * iw - w0/2)
            y0 = int(cyn * ih - h0/2)
            cv2.rectangle(box_vis, (x0, y0), (x0 + w0, y0 + h0), (0, 255, 0), 2)

        before_imgs.append(Image.fromarray(seg_vis))
        after_imgs.append(Image.fromarray(box_vis))

    return before_imgs, after_imgs


# Build Gradio interface
with gr.Blocks() as app:
    gr.Markdown("# Segmentation → YOLOv8 Converter")
    api_input = gr.Textbox(label="Roboflow API Key", type="password")
    url_input = gr.Textbox(label="Roboflow Dataset URL (Segmentation)")
    run_btn = gr.Button("Convert")
    before_gallery = gr.Gallery(label="Before (Segmentation)").style(grid=[5], height="auto")
    after_gallery = gr.Gallery(label="After (Bounding Boxes)").style(grid=[5], height="auto")
    run_btn.click(convert_seg_to_bbox, inputs=[api_input, url_input], outputs=[before_gallery, after_gallery])

if __name__ == "__main__":
    app.launch()