import cv2
from ultralytics import YOLO
import os
from datetime import datetime

class GasBottleCounter:
    def __init__(self, model_path):
        self.model = YOLO(model_path)
        self.IOU_THRESHOLD = 0.5
        self.counts = {}
        self.total = 0

    def iou(self, boxA, boxB):
        xA = max(boxA[0], boxB[0])
        yA = max(boxA[1], boxB[1])
        xB = min(boxA[2], boxB[2])
        yB = min(boxA[3], boxB[3])
        inter = max(0, xB - xA) * max(0, yB - yA)
        if inter == 0:
            return 0.0
        areaA = (boxA[2]-boxA[0]) * (boxA[3]-boxA[1])
        areaB = (boxB[2]-boxB[0]) * (boxB[3]-boxB[1])
        return inter / float(areaA + areaB - inter)

    def generate_frames(self, input_path):
        """Generator that yields processed frames and updates internal count state."""
        cap = cv2.VideoCapture(input_path)
        if not cap.isOpened():
            return

        self.counts = {}
        self.total = 0
        detected_ids = set()
        id_to_box    = {}

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            results = self.model.track(frame, persist=True, conf=0.4, verbose=False)

            for r in results:
                if r.boxes is None:
                    continue

                for box in r.boxes:
                    if box.id is None:
                        continue

                    track_id = int(box.id[0])
                    cls_id   = int(box.cls[0])
                    cls_name = r.names[cls_id]
                    conf     = float(box.conf[0])
                    x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
                    curr_box = (x1, y1, x2, y2)

                    # Professional visualization (Blue/White)
                    cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 78, 3), 2) # Deep blue/orange accent
                    cv2.putText(
                        frame,
                        f"{cls_name} #{track_id}",
                        (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 78, 3), 2
                    )

                    if track_id not in detected_ids:
                        is_duplicate = False
                        for known_id, known_box in id_to_box.items():
                            if known_id not in detected_ids:
                                continue
                            if self.iou(curr_box, known_box) > self.IOU_THRESHOLD:
                                is_duplicate = True
                                break

                        if not is_duplicate:
                            detected_ids.add(track_id)
                            self.counts[cls_name] = self.counts.get(cls_name, 0) + 1

                    id_to_box[track_id] = curr_box

            # Build Overlay
            self.total = sum(self.counts.values())
            
            # Simple header overlay
            overlay = frame.copy()
            cv2.rectangle(overlay, (0, 0), (frame.shape[1], 40), (255, 255, 255), -1)
            cv2.addWeighted(overlay, 0.8, frame, 0.2, 0, frame)
            
            cv2.putText(
                frame,
                f"LIVE ANALYTICS | TOTAL DETECTED: {self.total}",
                (20, 28),
                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (3, 78, 255), 2 # Professional blue
            )

            # Encode frame to JPEG
            ret, buffer = cv2.imencode('.jpg', frame)
            frame_bytes = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')

        cap.release()

        # Save results to file
        os.makedirs("results", exist_ok=True)
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        report_path = os.path.join("results", f"report_{timestamp}.txt")
        with open(report_path, "w") as f:
            f.write(f"GazCounter Industrial Report - {timestamp}\n")
            f.write("=" * 40 + "\n")
            for brand, count in sorted(self.counts.items()):
                f.write(f"  {brand}: {count}\n")
            f.write("=" * 40 + "\n")
            f.write(f"  TOTAL: {self.total}\n")

    def get_results(self):
        return self.counts, self.total