import cv2, json, os, sys, uuid, threading, time, base64
import numpy as np
from ultralytics import YOLO
from collections import defaultdict

YOLO_MODEL = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'best.pt')
TRACKER_CONFIG = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'bytetrack.yaml')
CONF_THRESH = 0.25
IOU_THRESH = 0.45
MAX_ABSENT_FRAMES = 15
SIMILARITY_IOU = 0.35

_model_lock = threading.Lock()
_model_cache = None


def _get_model():
    global _model_cache
    if _model_cache is None:
        _model_cache = YOLO(YOLO_MODEL)
    return _model_cache


def _ensure_tracker_config():
    path = TRACKER_CONFIG
    if not os.path.exists(path):
        with open(path, 'w') as f:
            f.write("""# ByteTrack
tracker_type: bytetrack
track_high_thresh: 0.25
track_low_thresh: 0.1
new_track_thresh: 0.6
track_buffer: 30
match_thresh: 0.8
fuse_score: True
# BoT-SORT (uncomment to use instead)
# tracker_type: botsort
# track_high_thresh: 0.25
# track_low_thresh: 0.1
# new_track_thresh: 0.5
# track_buffer: 50
# match_thresh: 0.8
# fuse_score: True
# gmc_method: sparseOptFlow
""")
    return path


class AppleTracker:
    def __init__(self, tracker_type='bytetrack'):
        self.model = _get_model()
        self.tracker_type = tracker_type
        self.tracker_cfg = _ensure_tracker_config()
        self.apple_tracks = {}
        self.next_id = 1
        self.frame_count = 0
        self.completed_tracks = set()

    def _get_label(self, cls_id):
        APPLE_CLASSES = {0: 'apple'}
        return APPLE_CLASSES.get(int(cls_id), 'apple')

    def process_frame(self, frame):
        h, w = frame.shape[:2]
        self.frame_count += 1

        results = self.model.track(
            frame,
            persist=True,
            conf=CONF_THRESH,
            iou=IOU_THRESH,
            tracker=self.tracker_cfg,
            verbose=False,
            device='cpu'
        )

        current_ids = set()
        detections = []
        if results and len(results) > 0:
            r = results[0]
            if r.boxes is not None and r.boxes.id is not None:
                for i, box in enumerate(r.boxes):
                    track_id = int(box.id[i].item())
                    conf = float(box.conf[i].item())
                    cls = int(box.cls[i].item())
                    x1, y1, x2, y2 = box.xyxy[i].tolist()
                    x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                    bw, bh = x2 - x1, y2 - y1
                    cx, cy = x1 + bw // 2, y1 + bh // 2
                    radius = max(bw, bh) // 2

                    current_ids.add(track_id)
                    detections.append({
                        'frame': self.frame_count,
                        'track_id': track_id,
                        'bbox': [x1, y1, x2, y2],
                        'center': [cx, cy],
                        'radius': radius,
                        'confidence': round(conf * 100, 1),
                    })

                    if track_id not in self.apple_tracks:
                        self.apple_tracks[track_id] = {
                            'track_id': track_id,
                            'first_frame': self.frame_count,
                            'last_frame': self.frame_count,
                            'total_frames': 1,
                            'max_confidence': conf,
                            'bboxes': [detections[-1]],
                            'absent_frames': 0,
                        }
                    else:
                        t = self.apple_tracks[track_id]
                        t['last_frame'] = self.frame_count
                        t['total_frames'] += 1
                        t['absent_frames'] = 0
                        if conf > t['max_confidence']:
                            t['max_confidence'] = conf
                        if len(t['bboxes']) < 60:
                            t['bboxes'].append(detections[-1])
                        else:
                            t['bboxes'][-1] = detections[-1]

        for tid in list(self.apple_tracks.keys()):
            if tid not in current_ids:
                t = self.apple_tracks[tid]
                t['absent_frames'] += 1
                if t['absent_frames'] >= MAX_ABSENT_FRAMES:
                    if tid not in self.completed_tracks:
                        self.completed_tracks.add(tid)

        return detections, current_ids

    def get_summary(self):
        active = {tid: t for tid, t in self.apple_tracks.items()
                  if tid not in self.completed_tracks and t['absent_frames'] < 5}
        completed_ids = self.completed_tracks | {
            tid for tid, t in self.apple_tracks.items()
            if t['absent_frames'] >= 5
        }
        valid_tracks = [t for tid, t in self.apple_tracks.items() if tid in completed_ids]

        total_unique = len(valid_tracks)
        if total_unique == 0:
            total_unique = len(self.apple_tracks)

        avg_conf = np.mean([t['max_confidence'] for t in self.apple_tracks.values()]) * 100 if self.apple_tracks else 0

        return {
            'total_unique_apples': total_unique,
            'active_tracks': len(active),
            'completed_tracks': len(completed_ids),
            'total_frames_processed': self.frame_count,
            'average_confidence': round(avg_conf, 1),
            'tracks': [
                {
                    'track_id': t['track_id'],
                    'first_frame': t['first_frame'],
                    'last_frame': t['last_frame'],
                    'total_frames': t['total_frames'],
                    'max_confidence': round(t['max_confidence'] * 100, 1),
                    'status': 'completed' if tid in completed_ids else 'active',
                }
                for tid, t in self.apple_tracks.items()
            ],
        }

    def reset(self):
        self.apple_tracks.clear()
        self.completed_tracks.clear()
        self.next_id = 1
        self.frame_count = 0


def create_tracked_annotation(frame, detections, active_ids, completed_count):
    annotated = frame.copy()
    h, w = frame.shape[:2]
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = max(0.5, min(h, w) / 1500.0)
    thickness = max(2, int(font_scale * 2.5))
    line_w = max(3, int(min(h, w) / 300))

    for det in detections:
        tid = det['track_id']
        x1, y1, x2, y2 = det['bbox']
        color = (0, 200, 0) if tid in active_ids else (100, 100, 100)
        cv2.rectangle(annotated, (x1, y1), (x2, y2), color, line_w)
        label = f"#{tid}"
        (tw, th), _ = cv2.getTextSize(label, font, font_scale, thickness)
        pad = 4
        ly = y1 - 8 if y1 - th - pad > 0 else y2 + th + pad
        cv2.rectangle(annotated, (x1, ly - th - pad), (x1 + tw + pad * 2, ly + pad), (0, 0, 0), -1)
        cv2.putText(annotated, label, (x1 + pad, ly), font, font_scale, (255, 255, 255), thickness, cv2.LINE_AA)

    status = f"Frame: {detections[0]['frame'] if detections else 0} | Active: {len(active_ids)} | Total unique: {completed_count}"
    (stw, sth), _ = cv2.getTextSize(status, font, font_scale * 1.1, thickness)
    bar_h = sth + 16
    cv2.rectangle(annotated, (0, 0), (stw + 24, bar_h), (0, 0, 0), -1)
    cv2.putText(annotated, status, (12, sth + 8), font, font_scale * 1.1, (0, 255, 0), thickness, cv2.LINE_AA)

    return annotated


def process_video_file(video_path, max_frames=0, progress_callback=None):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise ValueError(f"Cannot open video: {video_path}")

    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAMES))
    if max_frames > 0:
        total_frames = min(total_frames, max_frames)

    tracker = AppleTracker()
    annotated_frames = []
    frame_idx = 0

    while True:
        if max_frames > 0 and frame_idx >= max_frames:
            break
        ret, frame = cap.read()
        if not ret:
            break

        detections, active_ids = tracker.process_frame(frame)
        if frame_idx % max(1, int(fps / 3)) == 0 or len(detections) > 0:
            annotated = create_tracked_annotation(
                frame, detections, active_ids,
                len(tracker.completed_tracks)
            )
            _, buffer = cv2.imencode('.jpg', annotated, [cv2.IMWRITE_JPEG_QUALITY, 85])
            annotated_b64 = base64.b64encode(buffer).decode('utf-8')
            annotated_frames.append({
                'frame': frame_idx + 1,
                'active_tracks': len(active_ids),
                'detections': len(detections),
                'annotated_base64': annotated_b64,
            })

            if len(annotated_frames) > 20:
                annotated_frames.pop(0)

        frame_idx += 1
        if progress_callback and frame_idx % max(1, int(total_frames / 10)) == 0:
            progress_callback(int(frame_idx / total_frames * 100))

    cap.release()
    summary = tracker.get_summary()

    return {
        'summary': summary,
        'fps': fps,
        'total_frames': frame_idx,
        'annotated_previews': annotated_frames,
    }