"""Decode untrusted entropy payloads outside the web-server process."""
import csv
import json
import resource
import sys
import time
from pathlib import Path

import torch
import backend
from bitstream import parse_bitstream


def main():
    resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
    directory = Path(sys.argv[1])
    data = (directory / 'reconstruction.afp').read_bytes()
    strings, info = parse_bitstream(data)
    model = backend.get_model()
    from src.utils import img_utils
    with torch.no_grad():
        if backend.DEVICE.startswith('cuda'):
            torch.cuda.synchronize()
        start = time.perf_counter()
        decoded, _, _ = backend.decode_strings(model, strings)
        if backend.DEVICE.startswith('cuda'):
            torch.cuda.synchronize()
        elapsed = (time.perf_counter()-start)*1000
        if tuple(decoded.shape[-2:]) != (info['height'], info['width']) or not torch.isfinite(decoded).all():
            raise ValueError('Invalid reconstruction.')
        img_utils.imwrite(str(directory/'reconstruction.png'), decoded)
    result = dict(dataset='user_bitstream', image_name='reconstruction.png', mode='decode',
                  **info, nominal_bpp=backend.TARGETS[info['quality_index']],
                  bpp=len(data)*8/(info['width']*info['height']), bitstream_bytes=len(data),
                  decode_ms=elapsed, device=backend.DEVICE,
                  metric_protocol='No reference image supplied; PSNR and SSIM are not computed.',
                  timing_note='Single-request decoder timing excludes model loading.')
    (directory/'metrics.json').write_text(json.dumps(result,indent=2))
    with (directory/'metrics.csv').open('w',newline='') as f:
        w=csv.DictWriter(f,fieldnames=list(result));w.writeheader();w.writerow(result)


if __name__ == '__main__':
    main()