import json
import uuid
from pathlib import Path
from PIL import Image
from sahi.predict import get_sliced_prediction

default_options = {
    "is_predict_saved": True,
    "predict_output_path": str(Path.cwd()),
    "is_mask_saved": True,
    "mask_output_path": str(Path.cwd()),
}

def sliced_predict(images, model, options=None):
    """
    Use YOLO + sahi (sliced prediction) to predict objects in images
    :param images: path to images
    :param model: model instance
    :param options: is_predict_saved, predict_output_path, is_mask_saved, mask_output_path
    :return: json serialized results
    """
    if options is None:
        options = default_options

    results = predict_images(images, model, options)

    if options["is_predict_saved"]:
        save_as_json(results, options["predict_output_path"])
    return results

def predict_images(images, model, options):
    results = []
    paths = [p for p in Path(images).iterdir() if p.is_file()]
    for path in paths:
        predictions = get_sliced_prediction(Image.open(path), model, slice_height=1024, slice_width=1024)
        results.extend(serialize_predictions(predictions, path))
        if options["is_mask_saved"]:
            predictions.export_visuals(export_dir=options["mask_output_path"], file_name=f"mask_{uuid.uuid4()}.png")
    return results

def serialize_predictions(result, path):
    results = []
    predictions = result.to_coco_predictions()
    for prediction in predictions:
        serialized = serialize_prediction(path, prediction)
        results.append(serialized)
    return results

def serialize_prediction(path, result):
    return {
        "path": path,
        "class": result["category_name"],
        "confidence": result["score"],
        "bbox": result["bbox"],
        "centroid": get_centroid(result["bbox"])
    }

def get_centroid(bbox):
    x_min, y_min, width, height = bbox
    centroid_x = x_min + width / 2
    centroid_y = y_min + height / 2
    return centroid_x, centroid_y

def save_as_json(data, path):
    output_path = path / f"{uuid.uuid4()}_detections.json"
    with open(output_path, "w") as f:
        json.dump(data, f, indent=4)
    print(f"✅ Results saved to {output_path}")