#!/usr/bin/env python3
"""Exercise camera defaults/errors and single-frame/movie agreement (CPU builds)."""
import os
from pathlib import Path
import struct
import subprocess
import sys
import tempfile
import zlib

BUILD = Path(sys.argv[1] if len(sys.argv) > 1 else 'build/Release').resolve()
TESTDIR = Path(sys.argv[2]).resolve() if len(sys.argv) > 2 else BUILD
ENV = dict(os.environ, OMP_NUM_THREADS='4')
FITSDIFF = Path(__file__).resolve().parent.parent / 'scripts' / 'fits_floatdiff.py'
PATH_MAX = os.pathconf('/', 'PC_PATH_MAX')


def run(binary, *args, ok=True, env=ENV):
    result = subprocess.run([str(binary), *map(str, args)], env=env,
                            capture_output=True, text=True)
    if (result.returncode == 0) != ok:
        raise AssertionError(f'{binary.name} {args}: {result.returncode}\n{result.stderr}')
    return result


def png_payload(path, dimensions):
    data = path.read_bytes()
    assert data[:8] == b'\x89PNG\r\n\x1a\n'
    offset, compressed = 8, bytearray()
    while offset < len(data):
        count, kind = struct.unpack_from('>I4s', data, offset)
        payload = data[offset + 8:offset + 8 + count]
        if kind == b'IHDR':
            assert struct.unpack_from('>II', payload) == dimensions
        if kind == b'IDAT':
            compressed.extend(payload)
        offset += count + 12
    return zlib.decompress(compressed)


def ppm_payload(path, dimensions):
    data = path.read_bytes()
    assert data[:2] == b'P6'
    fields, index = [], 2
    while len(fields) < 3:
        while data[index:index + 1].isspace():
            index += 1
        if data[index:index + 1] == b'#':
            while data[index:index + 1] not in (b'\n', b''):
                index += 1
            continue
        start = index
        while not data[index:index + 1].isspace():
            index += 1
        fields.append(int(data[start:index]))
        index += 1
    assert (fields[0], fields[1]) == dimensions, (fields, dimensions)
    return data[index:index + fields[0] * fields[1] * 3]


def image_payload(path, dimensions=(64, 48), allow_black=False):
    data = path.read_bytes()
    raw = ppm_payload(path, dimensions) if data[:2] == b'P6' else png_payload(path, dimensions)
    assert raw, f'missing image data: {path}'
    assert allow_black or any(raw), f'empty image: {path}'
    return raw


def map_vertices(path):
    data = path.read_bytes()
    assert data[:8] == b'GRLENS\x01\x00'
    assert struct.unpack_from('<Q', data, 32)[0] == 1
    vertices, triangles = struct.unpack_from('<QQ', data, 64)
    offset = 80
    values = []
    for _ in range(vertices):
        values.append(struct.unpack_from('<9dI', data, offset))
        offset += 76
    return values, data[offset:offset + triangles * 28]


with tempfile.TemporaryDirectory(prefix='gr-camera-cli-') as directory:
    tmp = Path(directory)
    for backend in ('minkowski', 'schwarzschild'):
        binary = BUILD / f'{backend}_sky'
        help_text = run(binary, '--help').stdout
        ext = 'png' if '.png' in help_text else 'ppm'
        hdr_available = '--hdr-output' in help_text
        for option in ('--observer-position', '--observer-velocity', '--camera-roll-deg'):
            assert option in help_text
        assert '--observer-inward-speed' not in help_text
        assert '_mesh.' in help_text, help_text
        common = ['--catalog', 'assets/sky_grid_5deg.csv', '--width', 64,
                  '--height', 48, '--fov-deg', 80, '--exposure', 0.1,
                  '--coarse-cell-pixels', 8, '--refine-max-level', 0, '--psf-relative-tail', 1e-4]

        def render(name, *options):
            path = tmp / f'{backend}_{name}.{ext}'
            run(binary, *common, '--output', path, *options)
            return image_payload(path)

        # CPU fast-mode CLI smoke test: the FFTW resolve must run and report its
        # one-time setup line.
        if backend == 'minkowski':
            fast_path = tmp / f'minkowski_fast.{ext}'
            fast = run(binary, *common, '--fast-mode', '--fast-supersample', 2,
                       '--output', fast_path)
            assert 'Fast FFTW:' in fast.stderr, fast.stderr
            assert image_payload(fast_path)

        # Equivalent independently specified and inferred camera geometry.
        inferred = render('position', '--observer-position', -30, 0, 0)
        explicit = render('explicit', '--observer-position', -30, 0, 0,
                          '--look-ra-deg', 0, '--look-dec-deg', 0)
        angled = render('angle', '--look-ra-deg', 0, '--look-dec-deg', 0)
        assert inferred == explicit == angled
        pole = render('pole', '--observer-position', 0, 0, 30)
        assert pole == render('pole_explicit', '--observer-position', 0, 0, 30,
                              '--look-ra-deg', 0, '--look-dec-deg', -90)
        default = render('default')
        pos = (0, 0, 0) if backend == 'minkowski' else (0, 0, 30)
        assert default == render('default_explicit', '--observer-position', *pos,
                                 '--look-ra-deg', 90, '--look-dec-deg', -90)
        assert render('radius', '--observer-radius', 40) == render(
            'radius_explicit', '--observer-radius', 40, '--look-ra-deg', 90,
            '--look-dec-deg', -90)
        for partial, value, ra, dec in [('--look-ra-deg', 37, 37, -90),
                                       ('--look-dec-deg', -23, 90, -23)]:
            assert render('partial', partial, value) == render(
                'complete', '--look-ra-deg', ra, '--look-dec-deg', dec)

        # --draw-mesh must leave the primary image untouched and only add a
        # mesh-overlay sibling.
        baseline = render('mesh_base')
        mesh_output = tmp / f'{backend}_mesh_overlay.{ext}'
        mesh_run = run(binary, *common, '--draw-mesh', '--output', mesh_output)
        assert image_payload(mesh_output) == baseline
        mesh_sibling = tmp / f'{backend}_mesh_overlay_mesh.{ext}'
        assert mesh_sibling.exists(), mesh_run.stderr
        assert image_payload(mesh_sibling) != baseline
        assert not (tmp / f'{backend}_mesh_overlay_mesh_mesh.{ext}').exists()

        # Only the final filename extension is replaced, dots in directories are
        # not touched.
        dotted_dir = tmp / 'dir.v1'
        dotted_dir.mkdir(exist_ok=True)
        dotted_output = dotted_dir / f'{backend}_image.final.{ext}'
        run(binary, *common, '--draw-mesh', '--output', dotted_output)
        assert dotted_output.exists()
        dotted_mesh = dotted_dir / f'{backend}_image.final_mesh.{ext}'
        assert dotted_mesh.exists(), dotted_output
        assert not (dotted_dir / f'{backend}_image.final.{ext}_mesh.{ext}').exists()
        assert image_payload(dotted_output) == baseline

        # A mesh sibling that overflows PATH_MAX must fail before catalog or
        # spacetime initialization.
        target = PATH_MAX - 3
        long_name = f'{backend}_image.{ext}'
        long_path = '/' + 'd' * (target - len(long_name) - 2) + '/' + long_name
        assert len(long_path) == target
        assert target + len('_mesh') > PATH_MAX
        too_long = run(binary, *common, '--draw-mesh', '--catalog',
                       tmp / 'missing_catalog.csv', '--output', long_path,
                       ok=False)
        assert 'Mesh overlay output path is too long' in too_long.stderr, too_long.stderr
        assert 'Blackbody backend' not in too_long.stderr, too_long.stderr
        assert 'PSF cache ready' not in too_long.stderr

        errors = [
            (['--observer-position', 1, 2], None),
            (['--observer-position', 1, 2, 'nan'], None),
            (['--observer-velocity', 0, 0, 'inf'], None),
            (['--look-ra-deg', 'nan'], None),
            (['--look-dec-deg', 'inf'], None),
            (['--observer-radius', 'nan'], None),
            (['--observer-radius', 0], None),
            (['--camera-roll-deg', 'nan'], None),
            (['--observer-position', 0, 0, 0], 'Cannot infer'),
            (['--observer-position', 3, 4, 5, '--observer-radius', 30], 'mutually exclusive'),
            (['--observer-velocity', 10, 0, 0], 'not timelike'),
            (['--observer-inward-speed', 0], None),
            (['--observer-track', 'missing.csv', '--observer-velocity', 0, 0, 0], 'cannot be combined'),
            (['--frames-dir', tmp, '--look-ra-deg', 0], 'cannot be combined'),
            (['--lens-map-input', 'missing.grlens', '--camera-roll-deg', 0], 'cannot be combined'),
        ]
        if backend == 'schwarzschild':
            errors += [(['--observer-position', 1.5, 0, 0, '--observer-velocity', -0.5, 0, 0], 'capture cutoff'),
                       (['--observer-position', 1.75, 0, 0], 'not timelike')]
            render('inside', '--observer-position', 1.75, 0, 0,
                   '--observer-velocity', -0.5, 0, 0, '--look-ra-deg', 0, '--look-dec-deg', 0)
        for options, message in errors:
            missing_catalog = tmp / 'should_not_be_created.csv'
            result = run(binary, '--catalog', missing_catalog, *options, ok=False)
            if message:
                assert message in result.stderr, result.stderr
            assert not missing_catalog.exists(), result.stderr
            assert 'PSF cache ready' not in result.stderr

        track = tmp / f'{backend}.csv'
        run(TESTDIR / f'test_observer_{backend}', track)
        single_map, movie_map = tmp / 'single.grlens', tmp / 'movie.grlens'
        single = render('moving', '--observer-position', 3, -4, 5,
                        '--observer-velocity', 0.2, -0.1, 0.3,
                        '--look-ra-deg', 37, '--look-dec-deg', -23,
                        '--camera-roll-deg', 19, '--lens-map-output', single_map)
        run(binary, *common, '--observer-track', track, '--frames-dir', tmp,
            '--frames-prefix', backend, '--duration', 0, '--fps', 1,
            '--lens-map-output', movie_map)
        movie = image_payload(tmp / f'{backend}_000000.{ext}')
        assert single == movie, f'{backend}: single/movie image mismatch'
        a, ta = map_vertices(single_map)
        b, tb = map_vertices(movie_map)
        assert len(a) == len(b) and ta == tb
        max_error = 0
        for x, y in zip(a, b):
            assert x[-1] == y[-1], 'ray classification mismatch'
            max_error = max(max_error, *(abs(v - w) for v, w in zip(x[:-1], y[:-1])))
        assert max_error < 1e-9, max_error

        # Movie frames keep the clean primary image and gain correctly named
        # mesh siblings; never "frame_000000.png_mesh.png".
        movie_mesh_dir = tmp / f'{backend}_movie_mesh'
        movie_mesh_dir.mkdir()
        run(binary, *common, '--observer-track', track, '--frames-dir',
            movie_mesh_dir, '--frames-prefix', 'frame', '--duration', 0,
            '--fps', 1, '--draw-mesh')
        assert image_payload(movie_mesh_dir / f'frame_000000.{ext}') == movie
        assert (movie_mesh_dir / f'frame_000000_mesh.{ext}').exists()
        assert not (movie_mesh_dir / f'frame_000000.{ext}_mesh.{ext}').exists()

        # A map import must still work without evaluating a camera/metric, and
        # must follow the same clean-main plus mesh-sibling rule.
        assert single == render('import', '--lens-map-input', single_map)
        imported_mesh = tmp / f'{backend}_import_mesh.{ext}'
        run(binary, *common, '--lens-map-input', single_map, '--draw-mesh',
            '--output', imported_mesh)
        assert image_payload(imported_mesh) == single
        assert (tmp / f'{backend}_import_mesh_mesh.{ext}').exists()

        if hdr_available:
            # The clean HDR FITS payload and primary image are bit-identical
            # with and without --draw-mesh; only the mesh sibling differs.
            hdr_output = tmp / f'{backend}_hdr.{ext}'
            hdr_mesh_output = tmp / f'{backend}_hdr_mesh.{ext}'
            run(binary, *common, '--hdr-output', '--output', hdr_output)
            run(binary, *common, '--hdr-output', '--draw-mesh', '--output',
                hdr_mesh_output)
            base_fits = tmp / f'{backend}_hdr_HDR.fits'
            mesh_fits = tmp / f'{backend}_hdr_mesh_HDR.fits'
            assert base_fits.exists() and mesh_fits.exists()
            diff = subprocess.run([sys.executable, str(FITSDIFF), str(base_fits),
                                   str(mesh_fits)], capture_output=True, text=True)
            assert diff.returncode == 0, diff.stdout + diff.stderr
            assert 'mismatches=0 max_abs=0 max_rel=0' in diff.stdout, diff.stdout
            assert image_payload(hdr_output) == baseline
            assert image_payload(hdr_mesh_output) == baseline
            assert image_payload(tmp / f'{backend}_hdr_mesh_mesh.{ext}') != baseline

        if backend == 'schwarzschild':
            # Two inward-looking free-fall samples at r=6.2696 and r=3.1593.
            # At 16:9 the latter frame finishes in generation 0, while the
            # former still needs refinement. Finishing the empty batch used
            # to abort the whole movie in generation 1.
            mixed = Path(__file__).parent / 'fixtures/schwarzschild_mixed_refinement.csv'
            parallel_map = tmp / 'mixed-parallel.grlens'
            result = run(binary, *common, '--height', 36, '--fov-deg', 60,
                         '--refine-max-level', 3, '--observer-track', mixed,
                         '--movie-track-samples', '--frames-dir', tmp,
                         '--frames-prefix', 'mixed', '--verbose',
                         '--lens-map-output', parallel_map)
            assert 'Ray trace generation 1: frame 0 added' in result.stderr
            assert 'Ray trace generation 0: frame 1 added' not in result.stderr
            for frame in range(2):
                image_payload(tmp / f'mixed_{frame:06d}.{ext}',
                              dimensions=(64, 36), allow_black=True)
            # Thread scheduling must preserve endpoints, frame/sample IDs and
            # the resulting adaptive mesh across the entire slab sweep.
            for threads in (1, 16):
                comparison_map = tmp / f'mixed-{threads}-threads.grlens'
                run(binary, *common, '--height', 36, '--fov-deg', 60,
                    '--refine-max-level', 3, '--observer-track', mixed,
                    '--movie-track-samples', '--frames-dir', tmp,
                    '--frames-prefix', f'mixed-{threads}-threads',
                    '--lens-map-output', comparison_map,
                    env=dict(ENV, OMP_NUM_THREADS=str(threads)))
                assert parallel_map.read_bytes() == comparison_map.read_bytes(), \
                    f'movie lens map changed with {threads} threads'
            print('schwarzschild: movie lens map identical with 1, 4 and 16 threads', flush=True)
        print(f'{backend}: CLI checks passed; single/movie image identical, map max error {max_error:.3g}', flush=True)