import cv2
import numpy as np
from PIL import Image, ImageDraw
from huggingface_hub import hf_hub_download, login
import os


class SignatureDetector:
    def __init__(self, conf_threshold=0.2, iou_threshold=0.45):
        self.conf_threshold = conf_threshold
        self.iou_threshold = iou_threshold

        token = os.environ.get("HF_TOKEN")
        if token:
            login(token=token)

        print("📥 Downloading model from HuggingFace Hub...")
        self.model = self._load_model()
        print("✅ Model loaded successfully!")

    def _load_model(self):
        possible_files = ["model.onnx", "saved_model", "yolov8s.pt", "model.pt"]

        for filename in possible_files:
            try:
                model_path = hf_hub_download(
                    repo_id="tech4humans/yolov8s-signature-detector",
                    filename=filename
                )
                print(f"✅ Found model: {filename}")
                return self._init_model(model_path, filename)
            except Exception:
                continue

        raise RuntimeError("❌ Tidak bisa menemukan file model di repo.")

    def _init_model(self, model_path, filename):
        if filename.endswith(".onnx"):
            import onnxruntime as ort
            session = ort.InferenceSession(model_path)
            return ("onnx", session)
        elif filename.endswith(".pt"):
            from ultralytics import YOLO
            return ("yolo", YOLO(model_path))
        else:
            import tensorflow as tf
            model = tf.saved_model.load(model_path)
            return ("tf", model)

    def encode_image(self, image_path):
        image_data = np.fromfile(image_path, dtype="uint8")
        image_data = np.expand_dims(image_data, axis=0)
        return image_data

    def draw_result(self, image_path, result):
        image = Image.open(image_path)
        draw = ImageDraw.Draw(image)
        img_width, img_height = image.size

        for box, score in zip(result["detection_boxes"], result["detection_scores"]):
            if score >= self.conf_threshold:        
                x1, y1, w, h = box
                x2, y2 = x1 + w, y1 + h

                x1 = int(x1 * img_width / 640)
                y1 = int(y1 * img_height / 640)
                x2 = int(x2 * img_width / 640)
                y2 = int(y2 * img_height / 640)

                color = tuple(np.random.randint(0, 256, size=3).tolist())
                draw.rectangle([x1, y1, x2, y2], outline=color, width=3)
                draw.text((x1, y1 - 10), f"{score:.2f}", fill=color)

        return cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    def _apply_nms(self, boxes, scores):
        if len(boxes) == 0:
            return [], []

        boxes_xywh = []
        for box in boxes:
            x1, y1, w, h = [float(v) for v in box]
            boxes_xywh.append([int(x1), int(y1), int(w), int(h)])

        indices = cv2.dnn.NMSBoxes(
            bboxes=boxes_xywh,
            scores=[float(s) for s in scores],
            score_threshold=self.conf_threshold,
            nms_threshold=self.iou_threshold
        )

        if len(indices) == 0:
            return [], []

        indices = indices.flatten()
        filtered_boxes = [boxes[i] for i in indices]
        filtered_scores = [scores[i] for i in indices]
        return filtered_boxes, filtered_scores

    def detect_from_path(self, image_path):
        engine_type, model = self.model
        image_data = self.encode_image(image_path)

        if engine_type == "onnx":
            input_name = model.get_inputs()[0].name
            outputs = model.run(None, {input_name: image_data})
            raw_boxes  = outputs[0][0].tolist()
            raw_scores = outputs[1][0].tolist()

            filtered_boxes, filtered_scores = self._apply_nms(raw_boxes, raw_scores)
            result = {
                "detection_boxes":  filtered_boxes,
                "detection_scores": filtered_scores
            }

        elif engine_type == "yolo":
            results = model(
                image_path,
                verbose=False,
                conf=self.conf_threshold,   
                iou=self.iou_threshold      
            )
            boxes  = results[0].boxes
            xyxyn  = boxes.xyxyn.cpu().numpy()
            scores = boxes.conf.cpu().numpy()

            converted_boxes = []
            for b in xyxyn:
                x1, y1, x2, y2 = b
                converted_boxes.append([x1*640, y1*640, (x2-x1)*640, (y2-y1)*640])

            result = {
                "detection_boxes":  converted_boxes,
                "detection_scores": scores
            }

        else:  
            infer  = model.signatures["serving_default"]
            outputs = infer(image_data)
            raw_boxes  = outputs["detection_boxes"].numpy()[0].tolist()
            raw_scores = outputs["detection_scores"].numpy()[0].tolist()

            filtered_boxes, filtered_scores = self._apply_nms(raw_boxes, raw_scores)
            result = {
                "detection_boxes":  filtered_boxes,
                "detection_scores": filtered_scores
            }

        num_signatures = sum(1 for s in result["detection_scores"] if s >= self.conf_threshold)
        annotated = self.draw_result(image_path, result)
        return annotated, num_signatures

    def detect(self, image_bgr, temp_path):
        cv2.imwrite(temp_path, image_bgr)
        return self.detect_from_path(temp_path)