"""Paired graybox / styled team-battle clips from matrix-shooter-js.

usage: matrix_capture.py plan --game CAPTURE_DIR --out DIR [--clips 300] [--block 10]
       matrix_capture.py work --game CAPTURE_DIR --out DIR --map paris [--shard 0 --shards 1] [--port 8131]

CAPTURE_DIR is the output of matrix_capture_game.py. One Chromium (GPU, ANGLE)
opens the game four times, once per look (original graybox, cartoon, war,
punk), and keeps those pages for every clip of its map. Every frame the
scripted player decides its input from the graybox page's state; the graybox
page steps 1/24 s and renders with its own camera, then the three styled pages
step with the same input and that camera. The simulation is deterministic, so
all four show the same moment; every page's state is compared each frame and
the largest drift is recorded.

plan writes DIR/jobs.jsonl: per map --clips clips spread over the walkable
--block-metre blocks (every block before any repeats), 20% at landmarks, eight
yaw sectors and three pitch tiers, team sizes, fight distance and the side the
enemy comes from, all seeded by the clip number.

DIR/<map>/clip_NNNNNN/
  original.mp4 cartoon.mp4 war.mp4 punk.mp4   960x544, 24 fps, x264 yuv444 crf 12
  cond/LLLLLL.png    graybox frame for latent L (frame 0, then 4L-3), lossless
  frames.jsonl       per frame: input, graybox state, camera, drift across pages
  meta.json          the job, renderer, drift and fight statistics
"""
import argparse
import asyncio
import base64
import json
import math
import os
import random
import subprocess
import time
from pathlib import Path

from playwright.async_api import async_playwright

LOOKS = ('original', 'cartoon', 'war', 'punk')
MAPS = ('paris', 'tsinghua')
FPS = 24
SECONDS = 10
STYLED_FRAMES = FPS * SECONDS
W, H = 960, 544
CHROME = os.environ.get('GBDATA_CHROMIUM', '/root/autodl-tmp/edenlive/cache/chrome-156/chrome-linux64/chrome')
# ANGLE on GL/EGL drops geometry at random in headless Chrome here (the same state rendered twice differs);
# the Vulkan backend renders it identically every time.
GPU_ARGS = ['--enable-gpu', '--ignore-gpu-blocklist', '--use-angle=vulkan', '--enable-features=Vulkan']
# Shoulder camera offsets (game.mjs updateCamera): right of the eye, and up per map.
RIGHT = {False: 1.9, True: 1.1}
UP = {'paris': .65, 'tsinghua': .15}
TURN = 2.6 / FPS          # rad per frame
TAPS = {'shift', ' ', 'r', 'f', '1', '2', '3'}
# (allies besides the player, enemies, weight): 1v1, outnumbered, even squads, big fights, lopsided, nobody.
TEAMS = [((0, 1), 15), ((0, 2), 10), ((0, 3), 10), ((1, 2), 15), ((2, 3), 15), ((3, 4), 8), ((4, 5), 7),
         ((1, 5), 10), ((0, 0), 10)]
DISTANCE = {'close': ((3, 6), (4, 9), .3), 'mid': ((8, 15), (10, 20), .35), 'far': ((16, 26), (20, 32), .35)}
# The player only closes in beyond the fight's own range and backs off inside its near edge.
APPROACH = {'close': (3, 8), 'mid': (6, 16), 'far': (14, 28)}
SIDES = {'front': (0.0, .5), 'left': (-math.pi / 2, .15), 'right': (math.pi / 2, .15), 'back': (math.pi, .2)}


def wrap(a):
    return (a + math.pi) % (2 * math.pi) - math.pi


def pick(rng, table):
    keys = list(table)
    return rng.choices(keys, [table[k][-1] for k in keys])[0]


class CombatPolicy:
    """A third-person shooter player: face the nearest enemy, line the crosshair
    up (the camera sits right of the eye, so the aim leads by asin(offset/d)),
    fire in bursts, strafe, dash, jump, crouch, switch guns, reload, throw
    grenades, look away and back, and back off when a wall squeezes the camera."""

    def __init__(self, seed, map_id, pitch_bias, distance='mid'):
        self.rng = random.Random(seed)
        self.map = map_id
        self.retreat, self.approach = APPROACH[distance]
        self.pitch_bias = pitch_bias
        self.held = set()
        self.strafe, self.strafe_left = None, 0
        self.look_away, self.look_left = 0.0, 0
        self.burst_left, self.pause_left = 0, 0
        self.ads_left, self.squeezed, self.escape_left = 0, 0, 0
        self.yaw = self.pitch = None

    def step(self, state, camera, frame):
        rng = self.rng
        p = state['player']
        if self.yaw is None:
            self.yaw, self.pitch = state['yaw'], state['pitch']
        foes = [a for a in state['actors'] if a['team'] == 1 and a['alive']]
        want, taps = set(), []
        ads = self.ads_left > 0
        self.ads_left = max(0, self.ads_left - 1)
        # A wall right behind the shoulder pulls the camera in: turn round and walk out.
        close = camera and camera.get('s') and camera['s']['distance'] < 2.4
        self.squeezed = self.squeezed + 1 if close else 0
        if self.squeezed > 8 and self.escape_left <= 0:
            self.escape_left, self.look_away = 30, rng.choice([-1, 1]) * rng.uniform(2.2, 3.1)
            self.look_left = 30
        if self.look_left <= 0 and rng.random() < .005:
            self.look_away = rng.choice([-1, 1]) * rng.uniform(.8, 2.6)
            self.look_left = rng.randint(18, 40)
        away = 0.0
        if self.look_left > 0:
            self.look_left -= 1
            away = self.look_away
        if foes and self.escape_left <= 0:
            foe = min(foes, key=lambda a: math.hypot(a['x'] - p['x'], a['z'] - p['z']))
            d = math.hypot(foe['x'] - p['x'], foe['z'] - p['z'])
            bearing = math.atan2(foe['x'] - p['x'], foe['z'] - p['z'])
            aim_yaw = bearing - math.asin(min(.95, RIGHT[ads] / max(d, 1e-3))) + away
            aim_pitch = math.atan2(foe['y'] + 1.1 - (p['y'] + p['h'] + 1.65 + UP[self.map]), d + 4.5)
            if away:
                aim_pitch += self.pitch_bias
            err = wrap(aim_yaw - self.yaw)
            self.yaw += max(-TURN, min(TURN, err))
            self.pitch += max(-.04, min(.04, aim_pitch - self.pitch))
            # Fire in bursts once lined up, and suppress when nearly there.
            if self.burst_left > 0:
                self.burst_left -= 1
                if abs(err) < .2:
                    want.add('j')
            elif self.pause_left > 0:
                self.pause_left -= 1
            elif abs(err) < .12:
                self.burst_left, self.pause_left = rng.randint(10, 40), rng.randint(2, 10)
            if ads or (abs(err) < .1 and d > 15 and rng.random() < .03):
                want.add('k')
                if not ads:
                    self.ads_left = rng.randint(24, 60)
            if d > self.approach:
                want.add('w')
            elif d < self.retreat:
                want.add('s')
            if rng.random() < .012 and 6 < d < 22 and state['grenades'] == 0:
                taps.append('f')
        else:
            # Nobody to fight, or walking out of a corner: walk and look around.
            self.escape_left = max(0, self.escape_left - 1)
            target = self.yaw + (away if self.escape_left else TURN * .5 * math.sin(frame / 31))
            self.yaw += max(-TURN, min(TURN, wrap(target - self.yaw)))
            self.pitch += max(-.03, min(.03, self.pitch_bias + .08 * math.sin(frame / 23) - self.pitch))
            want.add('w')
            if foes and rng.random() < .5:
                want.add('j')
        if self.strafe_left <= 0:
            self.strafe = rng.choice([None, 'a', 'd', 'a', 'd'])
            self.strafe_left = rng.randint(10, 36)
        self.strafe_left -= 1
        if self.strafe and not self.escape_left:
            want.add(self.strafe)
        if rng.random() < .008:
            taps.append('shift')
        if rng.random() < .006:
            taps.append(' ')
        if ('c' in self.held and rng.random() < .97) or ('c' not in self.held and rng.random() < .004):
            want.add('c')
        if frame == 0 or rng.random() < .015:
            taps.append(rng.choice(['1', '2', '3']))
        if state['ammo'] <= 2 and not state['reloading'] and rng.random() < .2:
            taps.append('r')
        down = sorted((want | set(taps)) - self.held)
        up = sorted((self.held - want) | (self.held & TAPS))
        self.held = (self.held - set(up)) | set(down)
        return dict(down=down, up=up, yaw=self.yaw, pitch=self.pitch)


def make_jobs(map_id, blocks, points, clips, offset):
    jobs = []
    order = random.Random(offset).sample(range(len(blocks)), len(blocks))
    for i in range(clips):
        rng = random.Random(offset + i)
        if rng.random() < .2 and points:
            point = rng.choice(points)
            x, z, anchor = point['x'], point['z'], point['id']
            yaw = (point['yaw'] if point['yaw'] is not None else 0.0) + rng.uniform(-.6, .6)
        else:
            x, z = blocks[order[i % len(blocks)]]
            x, z, anchor = x + rng.uniform(-3, 3), z + rng.uniform(-3, 3), None
            yaw = ((i * 3 + i // len(blocks)) % 8) * math.pi / 4 + rng.uniform(-.3, .3)
        (allies, enemies) = rng.choices([t for t, _ in TEAMS], [w for _, w in TEAMS])[0]
        distance, side = pick(rng, DISTANCE), pick(rng, SIDES)
        engage, spread = DISTANCE[distance][0], DISTANCE[distance][1]
        bearing = yaw + SIDES[side][0]
        enemy_at = []
        for _ in range(enemies):
            r, lateral = rng.uniform(*spread), rng.uniform(-5, 5)
            enemy_at.append([x + math.sin(bearing) * r + math.cos(bearing) * lateral, z + math.cos(bearing) * r - math.sin(bearing) * lateral])
        ally_at = []
        for _ in range(allies):
            b, r = yaw + rng.choice([-1, 1]) * rng.uniform(.9, 1.9), rng.uniform(2.5, 6)
            ally_at.append([x + math.sin(b) * r, z + math.cos(b) * r])
        jobs.append(dict(map=map_id, clip=i, seed=offset + i, anchor=anchor, teams=[allies + 1, enemies], distance=distance, side=side,
                         pitch_bias=rng.choice([-.22, -.04, .18, .3]),
                         setup=dict(seed=offset + i + 1, width=W, height=H, engage=list(engage),
                                    player=dict(x=x, z=z, yaw=yaw, pitch=-.04), allies=ally_at, enemies=enemy_at)))
    return jobs


async def open_pages(browser, port, map_id):
    pages = []
    for look in LOOKS:
        page = await browser.new_page(viewport={'width': W, 'height': H}, device_scale_factor=1)
        await page.add_init_script('window.__CAPTURE=true')
        page.on('pageerror', lambda e, look=look: print(f'[{look}] pageerror {e}', flush=True))
        for _ in range(100):
            try:
                await page.goto(f'http://127.0.0.1:{port}/?map={map_id}&style={look}')
                break
            except Exception:
                await asyncio.sleep(.3)
        pages.append(page)
    for page in pages:
        await page.wait_for_function('window.__cap && window.__cap.ready()', timeout=300_000)
    return pages


def encoder(path):
    return subprocess.Popen(['ffmpeg', '-y', '-hide_banner', '-loglevel', 'error', '-f', 'image2pipe', '-c:v', 'png', '-r', str(FPS),
                             '-i', '-', '-c:v', 'libx264', '-preset', 'veryfast', '-crf', '12', '-pix_fmt', 'yuv444p',
                             '-threads', '2', str(path)], stdin=subprocess.PIPE)


def drift(states):
    ref, worst = states[0], 0.0
    for s in states[1:]:
        worst = max(worst, abs(s['player']['x'] - ref['player']['x']), abs(s['player']['z'] - ref['player']['z']),
                    abs(s['bullets'] - ref['bullets']))
        for a, b in zip(s['actors'], ref['actors']):
            worst = max(worst, abs(a['x'] - b['x']), abs(a['z'] - b['z']), abs(a['health'] - b['health']),
                        float(a['alive'] != b['alive']) * 100)
    return worst


async def capture(pages, job, out, renderer):
    tmp = out.with_name(out.name + '.part')
    if tmp.exists():
        for f in sorted(tmp.rglob('*'), reverse=True):
            f.unlink() if f.is_file() else f.rmdir()
    (tmp / 'cond').mkdir(parents=True)
    setups = [await page.evaluate('cfg => window.__cap.setup(cfg)', job['setup']) for page in pages]
    policy = CombatPolicy(job['seed'], job['map'], job['pitch_bias'], job['distance'])
    state, camera = setups[0], None
    encoders = {look: encoder(tmp / f'{look}.mp4') for look in LOOKS}
    cond = {0} | {4 * l - 3 for l in range(1, 1 + (FPS * SECONDS - 1) // 4)}
    log = (tmp / 'frames.jsonl').open('w')
    worst, started, bullet_frames, actor_frames = 0.0, time.perf_counter(), 0, []
    for frame in range(FPS * SECONDS):
        action = policy.step(state, camera, frame)
        gray = await pages[0].evaluate('a => window.__cap.step(a)', dict(action, capture=True))
        # with --styled-frames N the styled pages only run (and are only recorded) for the first N frames
        styled = await asyncio.gather(*[p.evaluate('a => window.__cap.step(a)', dict(action, capture=True, cam=gray['camera']))
                                        for p in pages[1:]]) if frame < STYLED_FRAMES else []
        for look, result in zip(LOOKS, [gray, *styled]):
            png = base64.b64decode(result['image'].split(',', 1)[1])
            encoders[look].stdin.write(png)
            if look == 'original' and frame in cond:
                (tmp / 'cond' / f'{sorted(cond).index(frame):06d}.png').write_bytes(png)
        state, camera = gray['state'], gray['camera']
        gap = drift([gray['state']] + [s['state'] for s in styled])
        worst = max(worst, gap)
        bullet_frames += state['bullets'] > 0
        actor_frames.append(sum(a['alive'] for a in state['actors']))
        log.write(json.dumps(dict(frame=frame, input=action, state=state, camera=camera, drift=round(gap, 6))) + '\n')
    log.close()
    for proc in encoders.values():
        proc.stdin.close()
        proc.wait()
    meta = dict(job, frames=FPS * SECONDS, fps=FPS, width=W, height=H, looks=LOOKS, styled_frames=STYLED_FRAMES, renderer=renderer,
                cond_frames=sorted(cond), max_drift=worst, bullet_frame_share=round(bullet_frames / (FPS * SECONDS), 3),
                mean_alive_bots=round(sum(actor_frames) / len(actor_frames), 2), shots=state['shots'], hits=state['hits'],
                kills=state['kills'], wall_s=round(time.perf_counter() - started, 1))
    (tmp / 'meta.json').write_text(json.dumps(meta, indent=1, ensure_ascii=False))
    tmp.rename(out)
    return meta


async def run(args):
    server = subprocess.Popen(['python3', str(args.game / 'start.py'), '--port', str(args.port)],
                              stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        async with async_playwright() as pw:
            browser = await pw.chromium.launch(executable_path=CHROME, headless=True, args=GPU_ARGS)
            if args.mode == 'plan':
                jobs = []
                for m, map_id in enumerate(MAPS):
                    pages = await open_pages(browser, args.port, map_id)
                    blocks = await pages[0].evaluate('s => window.__cap.blocks(s)', args.block)
                    points = await pages[0].evaluate('window.__cap.points()')
                    jobs += make_jobs(map_id, blocks, points, args.clips, args.seed_base * (m + 1))
                    print(json.dumps(dict(map=map_id, blocks=len(blocks), landmarks=len(points), clips=args.clips)), flush=True)
                    for page in pages:
                        await page.close()
                args.out.mkdir(parents=True, exist_ok=True)
                (args.out / 'jobs.jsonl').write_text(''.join(json.dumps(j) + '\n' for j in jobs))
                return
            jobs = [json.loads(l) for l in (args.out / 'jobs.jsonl').read_text().splitlines() if l.strip()]
            jobs = [j for j in jobs if j['map'] == args.map][args.shard::args.shards]
            pages = await open_pages(browser, args.port, args.map)
            renderer = await pages[0].evaluate("(() => {const g=document.createElement('canvas').getContext('webgl2');"
                                                "const e=g.getExtension('WEBGL_debug_renderer_info');return g.getParameter(e.UNMASKED_RENDERER_WEBGL);})()")
            print(json.dumps(dict(map=args.map, shard=args.shard, jobs=len(jobs), renderer=renderer)), flush=True)
            for job in jobs:
                out = args.out / job['map'] / f"clip_{job['clip']:06d}"
                if out.exists():
                    continue
                meta = await capture(pages, job, out, renderer)
                print(json.dumps(dict(clip=f"{job['map']}/{out.name}", teams=job['teams'], distance=job['distance'],
                                      drift=meta['max_drift'], bullets=meta['bullet_frame_share'], shots=meta['shots'],
                                      kills=meta['kills'], wall_s=meta['wall_s'])), flush=True)
            await browser.close()
    finally:
        server.terminate()


def main():
    global SECONDS, STYLED_FRAMES
    parser = argparse.ArgumentParser()
    parser.add_argument('mode', choices=('plan', 'work'))
    parser.add_argument('--game', type=Path, required=True)
    parser.add_argument('--out', type=Path, required=True)
    parser.add_argument('--map', choices=MAPS)
    parser.add_argument('--clips', type=int, default=300)
    parser.add_argument('--block', type=float, default=10)
    parser.add_argument('--seed-base', type=int, default=100_000, help='job seeds are seed_base * (map index + 1) + clip')
    parser.add_argument('--shard', type=int, default=0)
    parser.add_argument('--shards', type=int, default=1)
    parser.add_argument('--port', type=int, default=8131)
    parser.add_argument('--looks', default=','.join(LOOKS),
                        help="pages to render, the graybox first (the sakura package adds 'sakura')")
    parser.add_argument('--seconds', type=int, default=SECONDS, help='clip length (long-rollout clips: 120)')
    parser.add_argument('--styled-frames', type=int, default=None,
                        help='record the styled looks only for the first N frames (long graybox clips need just frame 0)')
    args = parser.parse_args()
    SECONDS = args.seconds
    STYLED_FRAMES = args.styled_frames or FPS * SECONDS
    globals()['LOOKS'] = tuple(args.looks.split(','))
    if LOOKS[0] != 'original':
        parser.error('--looks must start with original (the graybox page drives the others)')
    asyncio.run(run(args))


if __name__ == '__main__':
    main()