"""Game capture without a browser: the Node engine steps the game, the GPU draws it.

``NativeSceneRenderer`` has the interface ``GameCaptureWorker`` expects from a
renderer (the same one ``NativeMatrixRenderer`` implements), so the session,
its capture thread and the model loop are unchanged. Condition frames stay on
the GPU as (H, W, 3) uint8 tensors and go to the proxy encoder as they are;
the game-window frames come back to the host for the paired video, as before.

``install()`` is called from the deployment adapter. Switches:

  EDEN_NATIVE_CAPTURE=1        use it for the slugs in EDEN_NATIVE_SLUGS
  EDEN_NATIVE_SLUGS=eiffel-assault
  EDEN_NATIVE_AUX=full         render every frame for the game-window video;
                               'condition' renders only the model's frames and
                               repeats the latest one in that video. Defaults
                               to 'full' only when that video is sent to the
                               player (FLASHDREAMS_PAIRED_VIDEO=1).
  EDEN_LATE_CAPTURE=1          run each chunk as one burst at the end of its
                               paced window, so the model sees fresh input
  EDEN_CHUNK_PERIOD_MS=437     pace a 12-frame chunk to this wall time. Faster
                               capture would otherwise speed the game up, since
                               presentation follows production; 437 ms is the
                               live median before this change.
  EDEN_NATIVE_FEATURES=1       encode each condition frame with the proxy encoder
                               as soon as it is rendered, on the capture stream,
                               and hand the model the finished features; the
                               model loop then skips its own proxy pass.

In a burst the chunk is one engine command: the engine streams each condition
frame's scene back the moment it is simulated, the GPU renders and encodes it
while the engine keeps stepping, and the chunk is ready one render and one
single-frame encode after the last condition frame.
"""
import json
import os
import threading
import time
from collections import deque
from dataclasses import replace
from pathlib import Path

import torch
from loguru import logger

from engine_client import Engine, standalone_config
from native_render import NativeRenderer

# Proxy features by chunk id, from the capture thread to RenderControlStream.take().
FEATURES = {}
_FEATURES_KEPT = 4


class ProxyFeatures:
    """Model input that already went through the proxy encoder."""
    __slots__ = ('proxy_features',)

    def __init__(self, proxy_features):
        self.proxy_features = proxy_features


def enabled_for(slug):
    if os.environ.get('EDEN_NATIVE_CAPTURE', '0') != '1':
        return False
    return slug in os.environ.get('EDEN_NATIVE_SLUGS', 'eiffel-assault').split(',')


def _late_capture():
    return os.environ.get('EDEN_LATE_CAPTURE', '1') == '1'


def _features_enabled():
    return os.environ.get('EDEN_NATIVE_FEATURES', '1') == '1'


def _proxy_encoder():
    """A second copy of the live proxy encoder, built from the same profile as the runner's."""
    from lightx2v.models.input_encoders.edenlive.proxy_dino import ProxyDINO
    root = Path(os.environ.get('EDEN_DEPLOY_ROOT', '/root/autodl-tmp/edenlive'))
    profile = json.loads(Path(os.environ.get('EDEN_5090_PROFILE', root / 'runtime-5090.json')).read_text())
    # As EdenLiveRunner.load_proxy_encoder() chooses for the cached Sage kernels.
    attention = 'sage_attn2' if profile.get('self_attn_1_type') in ('sage_attn2_sm90_fused', 'sage_attn2_sm120_cached') else 'torch_sdpa'
    return ProxyDINO(str(root / 'models/EdenLive-x2v'), profile.get('proxy_resize_mode', 'bilinear'),
                     use_cuda_graph=profile.get('use_cuda_graph', False), attention_type=attention,
                     quantized_ckpt=profile.get('proxy_quantized_ckpt'))


def _aux_full():
    default = 'full' if os.environ.get('FLASHDREAMS_PAIRED_VIDEO') == '1' else 'condition'
    return os.environ.get('EDEN_NATIVE_AUX', default) == 'full'


def chunk_period_s(frame_count):
    return float(os.environ.get('EDEN_CHUNK_PERIOD_MS', '437')) / 1000 * frame_count / 12


def split_release(events):
    """Hold back from the first release of a key pressed earlier in the same batch.

    Delivered together, a press and its release would land before one
    simulation step and a game that polls held keys would never see the tap.
    The remainder is delivered with the next frame.
    """
    pressed = set()
    for index, event in enumerate(events):
        action = event.get('action')
        if action == 'pointer':
            continue
        if event.get('phase') == 'down':
            pressed.add(action)
        elif event.get('phase') == 'up' and action in pressed:
            return events[:index], events[index:]
    return events, []


class NativeSceneRenderer:
    def __init__(self, slug, *, width=960, height=544, games_root=None):
        root = Path(games_root or os.environ.get('EDEN_GAMES_ROOT', '/root/autodl-tmp/edenlive/EdenLive/realtime/games'))
        entry = next(g for g in json.loads((root / 'registry.json').read_text())['games'] if g['slug'] == slug)
        self.slug, self.width, self.height = slug, width, height
        self.engine = Engine(slug, width, height, games_root=root, config=standalone_config(entry))
        # Its own high-priority stream: a render never queues behind model kernels.
        self.stream = torch.cuda.Stream(priority=-1)
        with torch.cuda.stream(self.stream):
            self.renderer = NativeRenderer(self.engine.scene, self.engine.scene_arrays)
        self.webgl_renderer = f'native nvdiffrast on {torch.cuda.get_device_name()}'
        self.last_evaluate_s = self.last_decode_s = 0.0
        self.last_burst_s = 0.0
        self._bursts = deque(maxlen=8)
        self._target_end = None
        self._carry = []
        self._last_host = None
        self._last_device = None
        self.proxy = None
        self.measure_frames = False
        self.frame_to_features_s = []
        if _features_enabled():
            with torch.cuda.stream(self.stream):
                self.proxy = _proxy_encoder()
        logger.info('Game2V capture: {} runs browser-free ({} draws)', slug, len(self.engine.scene['draws']))

    def render_chunk(self, chunk, *, duration_s=0.0, drain_events=None):
        started = time.perf_counter()
        late = _late_capture() and duration_s > 0
        if late:
            # Chunks end on a fixed schedule, each one period after the last
            # one's scheduled end, so finishing a burst early never compounds
            # into a faster game. After a pause the schedule restarts.
            previous = self._target_end
            target = started + duration_s if previous is None or previous < started - duration_s else previous + duration_s
            self._target_end = target
            if self._bursts:
                time.sleep(max(0.0, target - (max(self._bursts) + 0.003) - time.perf_counter()))
        else:
            self._target_end = None
        burst_started = time.perf_counter()
        if late or duration_s <= 0:
            return self._burst(chunk, drain_events, burst_started)
        aux_full = _aux_full()
        frames, images = [], []
        engine_s = 0.0
        packets = chunk.frames
        draws = []

        def send(index):
            # Input is read when a step is queued, as the browser reads it per frame.
            events = self._carry + (drain_events() if drain_events is not None else [])
            events, self._carry = split_release(events)
            draw = packets[index].capture_condition or aux_full or (self._last_device is None and not draws)
            draws.append(draw)
            self.engine.send_advance(packets[index].simulation_tick, events, capture=draw)

        with torch.cuda.stream(self.stream):
            # Paced without a burst (EDEN_LATE_CAPTURE=0): one step at a time,
            # input read before each frame, spread over the chunk's window.
            for index, packet in enumerate(packets):
                send(index)
                step_started = time.perf_counter()
                header, scene = self.engine.receive()
                engine_s += time.perf_counter() - step_started
                if draws[index]:
                    self._last_device = self.renderer.render(scene)
                    images.append(self._last_device)
                frames.append((packet, header.get('state'), self._last_device, len(images) - 1))
                remaining = started + duration_s * (index + 1) / len(packets) - time.perf_counter()
                if remaining > 0:
                    time.sleep(remaining)
            # Game-window frames for the paired video, in one copy.
            host = torch.stack(images).cpu().numpy() if images else None
        # Everything the model will read from this chunk is finished on the
        # render stream before the chunk is handed over.
        self.stream.synchronize()
        out = []
        for packet, state, device, slot in frames:
            ui = host[slot] if slot >= 0 else self._last_host
            out.append(replace(packet, ui_frame=ui, ui_state=state,
                               condition_frame=device if packet.capture_condition else None))
        if host is not None:
            self._last_host = host[-1]
        burst = time.perf_counter() - burst_started
        self.last_burst_s = burst
        self._bursts.append(burst)
        self.last_evaluate_s = engine_s
        self.last_decode_s = burst - engine_s
        return replace(chunk, frames=tuple(out))

    def _burst(self, chunk, drain_events, burst_started):
        """The whole chunk as one engine command, condition frames rendered and encoded as they arrive."""
        packets = chunk.frames
        aux_full = _aux_full()
        # Input that piled up before the burst opens it; a release that would
        # cancel a press in the same frame moves to the next frame.
        pending = self._carry + (drain_events() if drain_events is not None else [])
        plan = []
        for index, packet in enumerate(packets):
            now, pending = split_release(pending)
            draw = packet.capture_condition or aux_full or (index == 0 and self._last_device is None)
            plan.append(dict(targetTick=packet.simulation_tick, events=now, capture=draw))
        self._carry = pending
        self.engine.send_chunk(plan)
        rendered, features = {}, []
        encode = self.proxy is not None
        with torch.cuda.stream(self.stream):
            while True:
                header, scene = self.engine.receive()
                if header['type'] != 'scene':
                    break
                index = header['index']
                arrived = time.perf_counter()
                image = self.renderer.render(scene)
                rendered[index] = image
                if encode and packets[index].capture_condition:
                    # The encoder's CUDA graph reuses its output buffer.
                    features.append(self.proxy.infer(image.permute(2, 0, 1)[None].contiguous()).clone())
                    if self.measure_frames:
                        # Measurement only: wait for this frame's features here.
                        torch.cuda.current_stream().synchronize()
                        self.frame_to_features_s.append(time.perf_counter() - arrived)
            states = header['states']
            engine_done = time.perf_counter()
            order = sorted(rendered)
            host = torch.stack([rendered[i] for i in order]).cpu().numpy() if order else None
        self.stream.synchronize()
        ready = time.perf_counter()
        # Engine-to-features for the chunk's last rendered frame, and for the
        # whole chunk from the moment the engine command went out.
        self.last_frame_to_features_s = ready - arrived if rendered else 0.0
        self.last_engine_s = engine_done - burst_started
        out, slot = [], -1
        for index, packet in enumerate(packets):
            if index in rendered:
                self._last_device = rendered[index]
                slot = order.index(index)
            ui = host[slot] if slot >= 0 else self._last_host
            out.append(replace(packet, ui_frame=ui, ui_state=states[index],
                               condition_frame=rendered.get(index) if packet.capture_condition else None))
        if host is not None:
            self._last_host = host[-1]
        if encode and len(features) == sum(p.capture_condition for p in packets):
            FEATURES[chunk.chunk_id] = torch.cat(features, dim=2)
            for stale in sorted(FEATURES)[:-_FEATURES_KEPT]:
                FEATURES.pop(stale, None)
        burst = time.perf_counter() - burst_started
        self.last_burst_s = burst
        self._bursts.append(burst)
        self.last_evaluate_s = header['timing']['chunk_ms'] / 1000
        self.last_decode_s = burst - self.last_evaluate_s
        return replace(chunk, frames=tuple(out))

    def probe(self, *, drain_events=None):
        events = self._carry + (drain_events() if drain_events is not None else [])
        self._carry = []
        return {'state': self.engine.probe(events).get('state')}

    def close(self):
        engine, self.engine = self.engine, None
        if engine is not None:
            engine.close()
        self.renderer = self.proxy = None
        self._last_device = self._last_host = None
        FEATURES.clear()


def install(game_renderer, session):
    """Route enabled slugs to NativeSceneRenderer and pace their chunks to a fixed period."""
    if getattr(game_renderer, '_native_capture_installed', False):
        return
    original_type = game_renderer._renderer_type
    def renderer_type(slug, games_root=None):
        if enabled_for(slug):
            return NativeSceneRenderer
        return original_type(slug, games_root)
    game_renderer._renderer_type = renderer_type

    original_pace = session._pace_duration_s
    def pace_duration_s(state, step_index):
        if not enabled_for(state.config.game_slug):
            return original_pace(state, step_index)
        return chunk_period_s(int(state.pipeline.get_num_output_frames(step_index)))
    session._pace_duration_s = pace_duration_s
    game_renderer._native_capture_installed = True