from __future__ import annotations
from scrypted_sdk.types import ObjectDetectionResult, ObjectsDetected, Setting
import io
from PIL import Image
import re
import scrypted_sdk
from typing import Any, List, Tuple, Mapping
import asyncio
import time
from .rectangle import Rectangle, intersect_area, intersect_rect, to_bounding_box, from_bounding_box, combine_rect
import urllib.request
import os

from detect import DetectionSession, DetectPlugin

from .sort_oh import tracker
import numpy as np
import traceback

try:
    from gi.repository import Gst
except:
    pass

class PredictSession(DetectionSession):
    image: Image.Image
    tracker: sort_oh.tracker.Sort_OH

    def __init__(self, start_time: float) -> None:
        super().__init__()
        self.image = None
        self.processed = 0
        self.start_time = start_time
        self.tracker = None

def parse_label_contents(contents: str):
    lines = contents.splitlines()
    ret = {}
    for row_number, content in enumerate(lines):
        pair = re.split(r'[:\s]+', content.strip(), maxsplit=1)
        if len(pair) == 2 and pair[0].strip().isdigit():
            ret[int(pair[0])] = pair[1].strip()
        else:
            ret[row_number] = content.strip()
    return ret

class RawImage:
    jpegMediaObject: scrypted_sdk.MediaObject

    def __init__(self, image: Image.Image):
        self.image = image
        self.jpegMediaObject = None

def is_same_box(bb1, bb2, threshold = .7):
    r1 = from_bounding_box(bb1)
    r2 = from_bounding_box(bb2)
    ia = intersect_area(r1, r2)

    if not ia:
        return False, None

    a1 = bb1[2] * bb1[3]
    a2 = bb2[2] * bb2[3]

    # if area intersect area is too small, these are different boxes
    if ia / a1 < threshold or ia / a2 < threshold:
        return False, None

    l = min(bb1[0], bb2[0])
    t = min(bb1[1], bb2[1])
    r = max(bb1[0] + bb1[2], bb2[0] + bb2[2])
    b = max(bb1[1] + bb1[3], bb2[1] + bb2[3])

    w = r - l
    h = b - t

    return True, (l, t, w, h)

def is_same_detection(d1: ObjectDetectionResult, d2: ObjectDetectionResult):
    if d1['className'] != d2['className']:
        return False, None

    return is_same_box(d1['boundingBox'], d2['boundingBox'])

def dedupe_detections(input: List[ObjectDetectionResult], is_same_detection = is_same_detection):
    input = input.copy()
    detections = []
    while len(input):
        d = input.pop()
        found = False
        for c in detections:
            same, box = is_same_detection(d, c)
            if same:
                # encompass this box and score
                d['boundingBox'] = box
                d['score'] = max(d['score'], c['score'])
                # remove from current detections list
                detections = list(filter(lambda r: r != c, detections))
                # run dedupe again with this new larger item
                input.append(d)
                found = True
                break

        if not found:
            detections.append(d)
    return detections

class Prediction:
    def __init__(self, id: int, score: float, bbox: Tuple[float, float, float, float]):
        self.id = id
        self.score = score
        self.bbox = bbox

class PredictPlugin(DetectPlugin, scrypted_sdk.BufferConverter, scrypted_sdk.Settings):
    labels: dict

    def __init__(self, PLUGIN_MIME_TYPE: str, nativeId: str | None = None):
        super().__init__(nativeId=nativeId)

        self.fromMimeType = PLUGIN_MIME_TYPE
        self.toMimeType = scrypted_sdk.ScryptedMimeTypes.MediaObject.value

        self.crop = False
        self.trackers: Mapping[str, tracker.Sort_OH] = {}

        # periodic restart because there seems to be leaks in tflite or coral API.
        loop = asyncio.get_event_loop()
        loop.call_later(4 * 60 * 60, lambda: self.requestRestart())

    def downloadFile(self, url: str, filename: str):
        filesPath = os.path.join(os.environ['SCRYPTED_PLUGIN_VOLUME'], 'files')
        fullpath = os.path.join(filesPath, filename)
        if os.path.isfile(fullpath):
            return fullpath
        os.makedirs(filesPath, exist_ok=True)
        tmp = fullpath + '.tmp'
        urllib.request.urlretrieve(url, tmp)
        os.rename(tmp, fullpath)
        return fullpath

    def getClasses(self) -> list[str]:
        return list(self.labels.values())

    def getTriggerClasses(self) -> list[str]:
        return ['motion']

    async def createMedia(self, data: RawImage) -> scrypted_sdk.MediaObject:
        mo = await scrypted_sdk.mediaManager.createMediaObject(data, self.fromMimeType)
        return mo

    def end_session(self, detection_session: PredictSession):
        image = detection_session.image
        if image:
            detection_session.image = None
            image.close()

        dps = detection_session.processed / (time.time() - detection_session.start_time)
        print("Detections per second %s" % dps)
        return super().end_session(detection_session)

    def invalidateMedia(self, detection_session: PredictSession, data: RawImage):
        if not data:
            return
        image = data.image
        data.image = None
        if image:
            if not detection_session.image:
                detection_session.image = image
            else:
                image.close()
        data.jpegMediaObject = None

    async def convert(self, data: RawImage, fromMimeType: str, toMimeType: str, options: scrypted_sdk.BufferConvertorOptions = None) -> Any:
        mo = data.jpegMediaObject
        if not mo:
            image = data.image
            if not image:
                raise Exception('data is no longer valid')

            bio = io.BytesIO()
            image.save(bio, format='JPEG')
            jpegBytes = bio.getvalue()
            mo = await scrypted_sdk.mediaManager.createMediaObject(jpegBytes, 'image/jpeg')
            data.jpegMediaObject = mo
        return mo

    def requestRestart(self):
        asyncio.ensure_future(scrypted_sdk.deviceManager.requestRestart())

    # width, height, channels
    def get_input_details(self) -> Tuple[int, int, int]:
        pass

    def getModelSettings(self, settings: Any = None) -> list[Setting]:
        allowList: Setting = {
            'title': 'Detections Types',
            # 'subgroup': 'Advanced',
            'description': 'The detections that will be reported. If none are specified, all detections will be reported. Select only detection types of interest for optimal performance.',
            'choices': self.getClasses(),
            'multiple': True,
            'key': 'allowList',
            'value': [
                'person',
                'dog',
                'cat',
                'car',
                'truck',
                'bus',
                'motorcycle',
            ],
        }

        trackerWindow: Setting = {
            'title': 'Tracker Window',
            'subgroup': 'Advanced',
            'description': 'Internal Setting. Do not change.',
            'key': 'trackerWindow',
            'value': 3,
            'type': 'number',
        }
        trackerCertainty: Setting = {
            'title': 'Tracker Certainty',
            'subgroup': 'Advanced',
            'description': 'Internal Setting. Do not change.',
            'key': 'trackerCertainty',
            'value': .2,
            'type': 'number',
        }
        return [allowList, trackerWindow, trackerCertainty]

    def create_detection_result(self, objs: List[Prediction], size, allowList, convert_to_src_size=None) -> ObjectsDetected:
        detections: List[ObjectDetectionResult] = []
        detection_result: ObjectsDetected = {}
        detection_result['detections'] = detections
        detection_result['inputDimensions'] = size

        for obj in objs:
            className = self.labels.get(obj.id, obj.id)
            if allowList and len(allowList) and className not in allowList:
                continue
            detection: ObjectDetectionResult = {}
            detection['boundingBox'] = (
                obj.bbox.xmin, obj.bbox.ymin, obj.bbox.xmax - obj.bbox.xmin, obj.bbox.ymax - obj.bbox.ymin)
            detection['className'] = className
            detection['score'] = obj.score
            detections.append(detection)

        if convert_to_src_size:
            detections = detection_result['detections']
            detection_result['detections'] = []
            for detection in detections:
                bb = detection['boundingBox']
                x, y, valid = convert_to_src_size((bb[0], bb[1]), True)
                x2, y2, valid2 = convert_to_src_size(
                    (bb[0] + bb[2], bb[1] + bb[3]), True)
                if not valid or not valid2:
                    # print("filtering out", detection['className'])
                    continue
                detection['boundingBox'] = (x, y, x2 - x + 1, y2 - y + 1)
                detection_result['detections'].append(detection)

        # print(detection_result)
        return detection_result

    async def run_detection_jpeg(self, detection_session: PredictSession, image_bytes: bytes, settings: Any) -> ObjectsDetected:
        stream = io.BytesIO(image_bytes)
        image = Image.open(stream)
        if image.mode == 'RGBA':
            image = image.convert('RGB')

        detections, _ = await self.run_detection_image(detection_session, image, settings, image.size)
        return detections

    def get_detection_input_size(self, src_size):
        # signals to pipeline that any input size is fine
        # previous code used to resize to correct size and run detection that way.
        # new code will resize the frame and potentially do multiple passes.
        # this is useful for high quality thumbnails.
        return (None, None)

    def get_input_size(self) -> Tuple[int, int]:
        pass

    async def detect_once(self, input: Image.Image, settings: Any, src_size, cvss) -> ObjectsDetected:
        pass

    async def run_detection_videoframe(self, videoFrame: scrypted_sdk.VideoFrame, detection_session: PredictSession) -> ObjectsDetected:
        settings = detection_session and detection_session.settings
        src_size = videoFrame.width, videoFrame.height
        w, h = self.get_input_size()
        iw, ih = src_size
        ws = w / iw
        hs = h / ih
        s = max(ws, hs)
        if ws == 1 and hs == 1:
            def cvss(point, normalize=False):
                return point[0], point[1], True

            data = await videoFrame.toBuffer({
                'format': videoFrame.format or 'rgb',
            })
            if videoFrame.format == 'rgba':
                image = Image.frombuffer('RGBA', (w, h), data).convert('RGB')
            else:
                image = Image.frombuffer('RGB', (w, h), data)
            try:
                ret = await self.detect_once(image, settings, src_size, cvss)
                return ret
            finally:
                image.close()

        sw = int(w / s)
        sh = int(h / s)
        first_crop = (0, 0, sw, sh)


        ow = iw - sw
        oh = ih - sh
        second_crop = (ow, oh, ow + sw, oh + sh)

        firstData, secondData = await asyncio.gather(
            videoFrame.toBuffer({
                'resize': {
                    'width': w,
                    'height': h,
                },
                'crop': {
                    'left': 0,
                    'top': 0,
                    'width': sw,
                    'height': sh,
                },
                'format': videoFrame.format or 'rgb',
            }),
            videoFrame.toBuffer({
                'resize': {
                    'width': w,
                    'height': h,
                },
                'crop': {
                    'left': ow,
                    'top': oh,
                    'width': sw,
                    'height': sh,
                },
                'format': videoFrame.format or 'rgb',
            })
        )

        if videoFrame.format == 'rgba':
            first = Image.frombuffer('RGBA', (w, h), firstData).convert('RGB')
        else:
            first = Image.frombuffer('RGB', (w, h), firstData)
        if videoFrame.format == 'rgba':
            second = Image.frombuffer('RGBA', (w, h), secondData).convert('RGB')
        else:
            second = Image.frombuffer('RGB', (w, h), secondData)

        def cvss1(point, normalize=False):
            return point[0] / s, point[1] / s, True
        def cvss2(point, normalize=False):
            return point[0] / s + ow, point[1] / s + oh, True

        ret1 = await self.detect_once(first, settings, src_size, cvss1)
        first.close()
        ret2 = await self.detect_once(second, settings, src_size, cvss2)
        second.close()

        two_intersect = intersect_rect(Rectangle(*first_crop), Rectangle(*second_crop))

        def is_same_detection_middle(d1: ObjectDetectionResult, d2: ObjectDetectionResult):
            same, ret = is_same_detection(d1, d2)
            if same:
                return same, ret

            if d1['className'] != d2['className']:
                return False, None

            r1 = from_bounding_box(d1['boundingBox'])
            m1 = intersect_rect(two_intersect, r1)
            if not m1:
                return False, None

            r2 = from_bounding_box(d2['boundingBox'])
            m2 = intersect_rect(two_intersect, r2)
            if not m2:
                return False, None

            same, ret = is_same_box(to_bounding_box(m1), to_bounding_box(m2))
            if not same:
                return False, None
            c = to_bounding_box(combine_rect(r1, r2))
            return True, c

        ret = ret1
        ret['detections'] = dedupe_detections(ret1['detections'] + ret2['detections'], is_same_detection=is_same_detection_middle)
        return ret
    
    async def run_detection_image(self, detection_session: PredictSession, image: Image.Image, settings: Any, src_size, convert_to_src_size: Any = None, multipass_crop: Tuple[float, float, float, float] = None):
        (w, h) = self.get_input_size() or image.size
        (iw, ih) = image.size

        if detection_session and not detection_session.tracker:
            t = self.trackers.get(detection_session.id)
            if not t:
                t = tracker.Sort_OH(scene=np.array([iw, ih]))
                trackerCertainty = settings.get('trackerCertainty')
                if not isinstance(trackerCertainty, int):
                    trackerCertainty = .2
                t.conf_three_frame_certainty = trackerCertainty * 3
                trackerWindow = settings.get('trackerWindow')
                if not isinstance(trackerWindow, int):
                    trackerWindow = 3
                t.conf_unmatched_history_size = trackerWindow
                self.trackers[detection_session.id] = t
            detection_session.tracker = t
            # conf_trgt = 0.35
            # conf_objt = 0.75
            # detection_session.tracker.conf_trgt = conf_trgt
            # detection_session.tracker.conf_objt = conf_objt

        # this a single pass or the second pass. detect once and return results.
        if multipass_crop:
            (l, t, dx, dy) = multipass_crop

            # find center
            cx = l + dx / 2
            cy = t + dy / 2

            # fix aspect ratio on box
            if dx / w > dy / h:
                dy = dx / w * h
            else:
                dx = dy / h * w

            if dx > image.width:
                s = image.width / dx
                dx = image.width
                dy *= s

            if dy > image.height:
                s = image.height / dy
                dy = image.height
                dx *= s

            # crop size to fit input size
            if dx < w:
                dx = w
            if dy < h:
                dy = h
            
            l = cx - dx / 2
            t = cy - dy / 2
            if l < 0:
                l = 0
            if t < 0:
                t = 0
            if l + dx > iw:
                l = iw - dx
            if t + dy > ih:
                t = ih - dy
            crop_box = (l, t, l + dx, t + dy)
            if dx == w and dy == h:
                input = image.crop(crop_box)
            else:
                input = image.resize((w, h), Image.ANTIALIAS, crop_box)

            def cvss(point, normalize=False):
                unscaled = ((point[0] / w) * dx + l, (point[1] / h) * dy + t)
                converted = convert_to_src_size(unscaled, normalize) if convert_to_src_size else (unscaled[0], unscaled[1], True)
                return converted

            ret = await self.detect_once(input, settings, src_size, cvss)
            input.close()
            detection_session.processed = detection_session.processed + 1
            return ret, RawImage(image)
        
        ws = w / iw
        hs = h / ih
        s = max(ws, hs)
        if ws == 1 and hs == 1:
            def cvss(point, normalize=False):
                converted = convert_to_src_size(point, normalize) if convert_to_src_size else (point[0], point[1], True)
                return converted

            ret = await self.detect_once(image, settings, src_size, cvss)
            if detection_session:
                detection_session.processed = detection_session.processed + 1
        else:
            sw = int(w / s)
            sh = int(h / s)
            first_crop = (0, 0, sw, sh)
            first = image.resize((w, h), Image.ANTIALIAS, first_crop)
            ow = iw - sw
            oh = ih - sh
            second_crop = (ow, oh, ow + sw, oh + sh)
            second = image.resize((w, h), Image.ANTIALIAS, second_crop)

            def cvss1(point, normalize=False):
                unscaled = (point[0] / s, point[1] / s)
                converted = convert_to_src_size(unscaled, normalize) if convert_to_src_size else (unscaled[0], unscaled[1], True)
                return converted
            def cvss2(point, normalize=False):
                unscaled = (point[0] / s + ow, point[1] / s + oh)
                converted = convert_to_src_size(unscaled, normalize) if convert_to_src_size else (unscaled[0], unscaled[1], True)
                return converted

            ret1 = await self.detect_once(first, settings, src_size, cvss1)
            first.close()
            if detection_session:
                detection_session.processed = detection_session.processed + 1
            ret2 = await self.detect_once(second, settings, src_size, cvss2)
            if detection_session:
                detection_session.processed = detection_session.processed + 1
            second.close()

            two_intersect = intersect_rect(Rectangle(*first_crop), Rectangle(*second_crop))

            def is_same_detection_middle(d1: ObjectDetectionResult, d2: ObjectDetectionResult):
                same, ret = is_same_detection(d1, d2)
                if same:
                    return same, ret

                if d1['className'] != d2['className']:
                    return False, None

                r1 = from_bounding_box(d1['boundingBox'])
                m1 = intersect_rect(two_intersect, r1)
                if not m1:
                    return False, None

                r2 = from_bounding_box(d2['boundingBox'])
                m2 = intersect_rect(two_intersect, r2)
                if not m2:
                    return False, None

                same, ret = is_same_box(to_bounding_box(m1), to_bounding_box(m2))
                if not same:
                    return False, None
                c = to_bounding_box(combine_rect(r1, r2))
                return True, c

            ret = ret1
            ret['detections'] = dedupe_detections(ret1['detections'] + ret2['detections'], is_same_detection=is_same_detection_middle)

        if detection_session:
            self.track(detection_session, ret)

        if not len(ret['detections']):
            return ret, RawImage(image)

        return ret, RawImage(image)
    
    def track(self, detection_session: PredictSession, ret: ObjectsDetected):
        detections = ret['detections']
        sort_input = []
        for d in ret['detections']:
            r: ObjectDetectionResult = d
            l, t, w, h = r['boundingBox']
            sort_input.append([l, t, l + w, t + h, r['score']])
        trackers, unmatched_trckr, unmatched_gts = detection_session.tracker.update(np.array(sort_input), [])
        for td in trackers:
            x0, y0, x1, y1, trackID = td[0].item(), td[1].item(
            ), td[2].item(), td[3].item(), td[4].item()
            slop = 0
            obj: ObjectDetectionResult = None
            ta = (x1 - x0) * (y1 - y0)
            box = Rectangle(x0, y0, x1, y1)
            for d in detections:
                if d.get('id'):
                    continue
                ob: ObjectDetectionResult = d
                dx0, dy0, dw, dh = ob['boundingBox']
                dx1 = dx0 + dw
                dy1 = dy0 + dh
                da = dw * dh
                area = intersect_area(Rectangle(dx0, dy0, dx1, dy1), box)
                if not area:
                    continue
                # intersect area always gonna be smaller than
                # the detection or tracker area.
                # greater numbers, ie approaching 2, is better.
                dslop = area / ta + area / da
                if (dslop > slop):
                    slop = dslop
                    obj = ob
            if obj:
                obj['id'] = str(trackID)
            # this may happen if tracker predicts something is still in the scene
            # but was not detected
            # else:
            #     print('unresolved tracker')
        # for d in detections:
        #     if not d.get('id'):
        #         # this happens if the tracker is not confident in a new detection yet due
        #         # to low score or has not been found in enough frames
        #         if d['className'] == 'person':
        #             print('untracked %s: %s' % (d['className'], d['score']))


    async def run_detection_crop(self, detection_session: DetectionSession, sample: RawImage, settings: Any, src_size, convert_to_src_size, bounding_box: Tuple[float, float, float, float]) -> ObjectsDetected:
        (ret, _) = await self.run_detection_image(detection_session, sample.image, settings, src_size, convert_to_src_size, bounding_box)
        return ret

    async def run_detection_gstsample(self, detection_session: PredictSession, gstsample, settings: Any, src_size, convert_to_src_size) -> Tuple[ObjectsDetected, Image.Image]:
        caps = gstsample.get_caps()
        # can't trust the width value, compute the stride
        height = caps.get_structure(0).get_value('height')
        width = caps.get_structure(0).get_value('width')
        gst_buffer = gstsample.get_buffer()
        result, info = gst_buffer.map(Gst.MapFlags.READ)
        if not result:
            return
        try:
            image = detection_session.image
            detection_session.image = None

            if image and (image.width != width or image.height != height):
                image.close()
                image = None
            if image:
                image.frombytes(bytes(info.data))
            else:
                image = Image.frombuffer('RGB', (width, height), bytes(info.data))
        finally:
            gst_buffer.unmap(info)

        try:
            return await self.run_detection_image(detection_session, image, settings, src_size, convert_to_src_size)
        except:
            image.close()
            traceback.print_exc()
            raise

    def create_detection_session(self):
        return PredictSession(start_time=time.time())
