import os
os.umask(0o077)
reader, writer = os.pipe()
os.dup2(reader, 0)
os.close(reader)
import collections
import configparser
import json
import pathlib
import threading
import time
from pyservicemaker import Pipeline, Flow, Probe, BufferOperator, BufferRetriever

R = pathlib.Path('/evidence')
cfg = next(pathlib.Path('/configs').glob('*.preprocess'))
c = configparser.ConfigParser()
c.read(cfg)
uid = int(c['property']['unique-id'])
counts = collections.Counter()
pending = collections.OrderedDict()
latencies = collections.defaultdict(lambda: collections.deque(maxlen=256))
first = threading.Event()
errors = []
first_time = None

def mark(stage, buffer):
    global first_time
    now = time.monotonic()
    counts[stage] += 1
    for frame in buffer.batch_meta.frame_items:
        key = (int(frame.frame_number), int(frame.buffer_pts))
        row = pending.setdefault(key, {})
        row[stage] = now
        if stage == 'sink':
            if first_time is None:
                first_time = now
                first.set()
            for a, b in [('mux', 'preprocess'), ('preprocess', 'infer'), ('infer', 'sink')]:
                if a in row and b in row:
                    latencies[a + '_to_' + b].append((row[b] - row[a]) * 1000)
            pending.pop(key, None)
        while len(pending) > 128:
            pending.popitem(last=False)

class Observe(BufferOperator):
    def __init__(self, stage):
        super().__init__()
        self.stage = stage
    def handle_buffer(self, buffer):
        mark(self.stage, buffer)
        return True

class Consume(BufferRetriever):
    def consume(self, buffer):
        mark('sink', buffer)
        return 1

p = Pipeline('sdk-baseline')
p.add('videotestsrc', 'source', {'is-live': True, 'pattern': 2})
p.add('capsfilter', 'cpu-caps', {'caps': 'video/x-raw,format=NV12,width=1920,height=1080,framerate=15/1'})
p.add('nvvideoconvert', 'convert', {'gpu-id': 0})
p.add('capsfilter', 'gpu-caps', {'caps': 'video/x-raw(memory:NVMM),format=NV12'})
p.add('nvstreammux', 'mux', {'width': 1920, 'height': 1080, 'batch-size': 1,
      'batched-push-timeout': 33000, 'buffer-pool-size': 4, 'drop-pipeline-eos': False,
      'live-source': True, 'gpu-id': 0, 'compute-hw': 1})
p.link('source', 'cpu-caps', 'convert', 'gpu-caps')
p.link(('gpu-caps', 'mux'), ('', 'sink_%u'))
probes = [Probe('observe-' + s, Observe(s)) for s in ['mux', 'preprocess', 'infer']]
p.attach('mux', probes[0])
f = Flow(p, streams=['mux']).preprocess(str(cfg), target_unique_ids=uid).attach(probes[1])
f = f.infer(str(cfg)[:-11], with_triton=True, batch_size=1).attach(probes[2])
f.retrieve(Consume())

def message(msg):
    name = type(msg).__name__
    if 'Error' in name:
        errors.append(name)

def wait():
    try:
        p.wait()
    except Exception as exc:
        errors.append(type(exc).__name__)

def write(data):
    (R / 'probe.json').write_text(json.dumps(data, indent=2))

start = time.monotonic()
p.start(message)
worker = threading.Thread(target=wait, daemon=True)
worker.start()
result = {'first_result': first.wait(30), 'source_nominal_fps': 15,
          'source_kind': 'synthetic_live', 'selection': 'all', 'error_types': errors}
if result['first_result']:
    result['first_result_s'] = first_time - start
    time.sleep(5)
    for values in latencies.values():
        values.clear()
    before = dict(counts)
    begin = time.monotonic()
    time.sleep(15)
    elapsed = time.monotonic() - begin
    after = dict(counts)
    result.update(measure_s=elapsed, delta={k: after[k] - before.get(k, 0) for k in after})
    result['fps'] = {k: v / elapsed for k, v in result['delta'].items()}
    result['latency_ms'] = {k: {'n': len(v), 'mean': sum(v)/len(v),
                             'max': max(v)} for k, v in latencies.items() if v}
result.update(counts=dict(counts), pending=len(pending), stop_requested=True)
write(result)
stop = time.monotonic()
p.stop()
result['stop_return_s'] = time.monotonic() - stop
write(result)
worker.join(max(0, 60 - (time.monotonic() - stop)))
result.update(wait_alive=worker.is_alive(), cleanup_s=time.monotonic() - stop, counts=dict(counts))
write(result)
os.close(writer)
os._exit(0 if not worker.is_alive() and result['first_result'] else 2)