"""Render the final decoded o-voxel directly, before UV or GLB extraction.

Uses TRELLIS.2's PbrMeshRenderer and HDR environments for material inspection.
"""
import json
import math
from pathlib import Path
from hf_space.timing import span


ENVIRONMENTS = ('studio', 'forest', 'sunset', 'courtyard')
INPUT_VIEW_LAYOUT = 'input_front_y_up_v1'


def input_camera(output):
    """Reuse the geometry projection camera, including its optimized scale."""
    output = Path(output)
    for path in (output.parent/'input-camera.json',
                 output.parent/(output.stem+'_work')/'input'/'camera.json'):
        if not path.is_file():
            continue
        camera = json.loads(path.read_text(encoding='utf-8'))
        fov = float(camera['camera_angle_x'])
        distance = float(camera['distance']) * float(camera.get('mesh_scale', 1.0))
        if not math.isfinite(fov) or not 0 < fov < math.pi or not math.isfinite(distance) or distance <= 0:
            raise ValueError('Invalid input projection camera')
        return {'camera_angle_x': fov, 'distance': distance, 'mesh_scale': 1.0}
    return None


def orbit_parameters(count, camera=None):
    # After the Z-up renderer's camera conversion below, yaw=pi/pitch=0
    # looks from +Z with +Y up, exactly as the input projection does.
    yaws = [math.pi + index * 2 * math.pi / count for index in range(count)]
    distance = camera['distance'] if camera else 2.0
    fov = math.degrees(camera['camera_angle_x']) if camera else 36.0
    return yaws, [0.0] * count, distance, fov


def load_environments():
    import os
    os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
    import cv2
    import torch
    from supratex.renderers import EnvMap
    folder = Path(__file__).resolve().parents[1]/'assets/preview-lighting'
    environments = {}
    for name in ENVIRONMENTS:
        image = cv2.imread(str(folder/f'{name}.exr'), cv2.IMREAD_UNCHANGED)
        if image is None:
            raise RuntimeError(f'Could not load HDR environment: {name}')
        environments[name] = EnvMap(torch.tensor(cv2.cvtColor(image, cv2.COLOR_BGR2RGB),
                                                dtype=torch.float32, device='cuda'))
    return environments


def save_and_render(output, representation):
    import numpy as np
    if representation.vertices.device.type != 'cuda':
        raise RuntimeError('Multiview rendering requires CUDA')
    output = Path(output)
    with span('ovoxel_save'):
        # Lossless float32 PBR data, original topology, no decimation or UV baking.
        np.savez(output, schema_version=np.int32(1), resolution=np.int32(1024),
                 vertices=representation.vertices.detach().cpu().numpy(),
                 faces=representation.faces.detach().cpu().numpy(),
                 coords=representation.coords.detach().cpu().numpy(),
                 attrs=representation.attrs.detach().cpu().numpy(),
                 origin=representation.origin.detach().cpu().numpy(),
                 voxel_size=np.float32(representation.voxel_size))
    render_views(output, representation)


def render_views(output, representation):
    import numpy as np
    import torch
    from PIL import Image
    output = Path(output)
    with span('preview_render'):
        original_faces = int(representation.faces.shape[0])
        preview = representation
        if original_faces > 16777216:
            with span('preview_simplify'):
                # TRELLIS.2's nvdiffrast face limit. Only simplify the view mesh;
                # the saved final o-voxel and future GLB retain original topology.
                preview = representation.to('cuda')
                preview.simplify(16777216)
        # Input geometry uses Y-up; TRELLIS's snapshot orbit uses Z-up.
        # Rotate only the render cameras, leaving final o-voxel coordinates intact.
        from supratex.utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics, render_frames
        count = 16
        camera = input_camera(output)
        if camera is not None:
            (output.parent/'input-camera.json').write_text(json.dumps(camera), encoding='utf-8')
        yaws, pitches, distance, fov = orbit_parameters(count, camera)
        extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitches, distance, fov)
        rotation = torch.tensor([[1.,0,0,0],[0,0,-1.,0],[0,1.,0,0],[0,0,0,1]],device='cuda')
        extrinsics = [matrix@rotation for matrix in extrinsics]
        with torch.no_grad():
            frames = render_frames(preview, extrinsics, intrinsics,
                options={'resolution':768,'ssaa':1,'peel_layers':2,
                         'near':min(.01, distance/100), 'far':max(10., distance+2)}, envmap=load_environments())
        folder = output.parent/'views'
        folder.mkdir(exist_ok=True)
        modes = ('normal','clay','base_color',*(f'shaded_{name}' for name in ENVIRONMENTS))
        for mode in modes:
            for index, pixels in enumerate(frames[mode]):
                mask = frames['mask'][index][..., :1].astype(np.float32)/255
                # Composite the same neutral background in every view.
                if mode.startswith('shaded_'):
                    pixels = np.clip(pixels.astype(np.float32)+(1-mask)*25,0,255).astype(np.uint8)
                else:
                    pixels = np.clip(pixels*mask+(1-mask)*25,0,255).astype(np.uint8)
                Image.fromarray(pixels).save(folder/f'{mode}-{index:02d}.webp',quality=90)
        (output.parent/'preview.json').write_text(json.dumps({
            'views':count,'modes':list(modes),'resolution':768,'device':'cuda',
            'renderer':'TRELLIS.2 PbrMeshRenderer','lighting':'HDR image-based lighting',
            'source_faces':original_faces,'preview_faces':int(preview.faces.shape[0]),
            'environments':list(ENVIRONMENTS),'default_mode':'shaded_forest',
            'view_layout':INPUT_VIEW_LAYOUT,'input_view':0,
            'input_camera':camera,
            'automatic_glb_export':False}))