"""Decode exactly the sampled PTS, then save unannotated/annotated training images."""
import json
from pathlib import Path
import sys
import time


def render(stage):
    import gi
    gi.require_version('Gst','1.0'); gi.require_version('GstVideo','1.0')
    from gi.repository import Gst, GstVideo
    from PIL import Image
    from app.artifact_media import packets, draw_frame
    import threading
    Gst.init(None)
    stage = Path(stage)
    job = json.loads((stage/'job.json').read_text())
    snapshot = job['snapshot']
    cfg = job['settings']['configured']
    first = next(packets(stage/'capture.bin'))
    caps = Gst.Caps.from_string(first[1])
    decoder = {'video/x-h264':'h264parse ! avdec_h264 max-threads=2',
               'video/x-h265':'h265parse ! avdec_h265 max-threads=2'}.get(caps.get_structure(0).get_name())
    if not decoder:
        raise ValueError('Unsupported training codec.')
    pipeline = Gst.parse_launch('appsrc name=src format=time block=true max-bytes=4194304 ! '+decoder+
        ' ! videoconvert ! video/x-raw,format=RGB ! appsink name=sink sync=false max-buffers=2 drop=false')
    src, sink = pipeline.get_by_name('src'), pipeline.get_by_name('sink')
    src.set_property('caps',caps)
    origin, target = first[0][0], snapshot['ptsNs']
    errors = []
    def feed():
        try:
            for h, description, data in packets(stage/'capture.bin'):
                if description != first[1]:
                    raise ValueError('Codec changed.')
                buffer = Gst.Buffer.new_allocate(None,len(data),None); buffer.fill(0,data)
                buffer.pts = h[0]-origin
                buffer.dts = Gst.CLOCK_TIME_NONE if h[1] == Gst.CLOCK_TIME_NONE else max(0,h[1]-origin)
                buffer.duration = h[2]
                if not h[5]: buffer.set_flags(Gst.BufferFlags.DELTA_UNIT)
                if src.emit('push-buffer',buffer) != Gst.FlowReturn.OK: break
            src.emit('end-of-stream')
        except Exception as error:
            errors.append(type(error).__name__)
            src.emit('end-of-stream')
    thread = threading.Thread(target=feed,daemon=True)
    pipeline.set_state(Gst.State.PLAYING); thread.start()
    found = False
    try:
        deadline = time.monotonic()+20
        while time.monotonic() < deadline:
            sample = sink.emit('try-pull-sample',100*Gst.MSECOND)
            if sample is None:
                if sink.get_property('eos') or errors: break
                if pipeline.get_bus().pop_filtered(Gst.MessageType.ERROR):
                    raise RuntimeError('Training decoder failed.')
                continue
            buf = sample.get_buffer()
            if int(buf.pts)+origin != target: continue
            info = GstVideo.VideoInfo.new_from_caps(sample.get_caps())
            if not 0 < info.width*info.height <= 4096*2160:
                raise ValueError('Training dimensions exceeded.')
            if (info.width,info.height) != (snapshot['frameWidth'],snapshot['frameHeight']):
                raise ValueError('Training source geometry mismatch.')
            ok, mapped = buf.map(Gst.MapFlags.READ)
            if not ok: raise RuntimeError('Training map failed.')
            try:
                image = Image.frombytes('RGB',(info.width,info.height),bytes(mapped.data),'raw','RGB',info.stride[0])
            finally:
                buf.unmap(mapped)
            size = cfg['width'],cfg['height']
            draw_frame(image,None,size).save(stage/'raw.jpg',quality=90)
            draw_frame(image,snapshot,size).save(stage/'overlay.jpg',quality=90)
            scale = min(size[0]/image.width,size[1]/image.height)
            scaled = round(image.width*scale),round(image.height*scale)
            offset = (size[0]-scaled[0])//2,(size[1]-scaled[1])//2
            meta = dict(job, savedAt=time.time_ns()//1_000_000, decodedPtsNs=target,
                        sourceSize=[image.width,image.height], storedSize=list(size),
                        transform=dict(scale=scale,resizedSize=list(scaled),paddingXY=list(offset)),
                        storedBoxes=[dict(objectRef=o['objectRef'],bboxXYXY=[
                            offset[i%2]+v*scaled[i%2] for i,v in enumerate(o['bbox'])]) for o in snapshot['objects']])
            with (stage/'metadata.json').open('x') as f:
                json.dump(meta,f,allow_nan=False)
            found = True
            break
        if not found:
            raise RuntimeError('Exact training Frame unavailable.')
    finally:
        pipeline.set_state(Gst.State.NULL)
        thread.join(2)
        if thread.is_alive(): raise RuntimeError('Training feed cleanup failed.')


if __name__ == '__main__':
    import os
    from app.channel_child import parent_death_guard
    parent_death_guard(int(os.environ['CONTEXTRA_TRAINING_PARENT_PID']))
    render(sys.argv[1])