# Copyright 2026 Toyota Research Institute.  All rights reserved.

import io
import json
import os
import subprocess
from glob import glob
from typing import Dict

import cv2
import numpy as np
import torch
import zarr
import pickle as pkl
from PIL import Image

from zarr.storage import ZipStore

from anydata.utils.misc import frame_name, join_dict, stack_sample
from anydata.geometry.depth import encode_dense

try:
    from torchcodec.encoders import VideoEncoder
except Exception:  # NOTE(GAIA 06-30): was `except ImportError` -- too narrow. Broken torchcodec (missing
    pass           # FFmpeg/CUDA libs) raises OSError, not ImportError, which leaked and crashed startup at
                   # module load. Match the decord-fallback guards in read.py/autodecode.py.


NT = np.ndarray | torch.Tensor
BIT_COMPRESS_LABELS = [
    'depth','semantic',
    'optflow_fwd','optflow_bwd',
    'mask_rgb','mask_depth',
]


# NOTE(bvh): storage keys were renamed 'frame'->'frames', 'sequence'->'videos'.
# Accept legacy names so older on-disk metadata.json / external callers still work.
_LEGACY_STORAGE = {'frame': 'frames', 'sequence': 'videos'}


def _norm_storage(storage):
    return _LEGACY_STORAGE.get(storage, storage)


def rgb_ext(storage):
    return {'frames':'jpg', 'videos':'mp4'}[_norm_storage(storage)]


def data_ext(storage):
    return {'frames':'npz', 'videos':'zarr'}[_norm_storage(storage)]


def create_folder(filename):
    """Create a new folder (if needed) to host a file"""
    if '/' not in filename: return
    os.makedirs(os.path.dirname(filename), exist_ok=True)


def write_pickle(filename, data):
    create_folder(filename)
    if not filename.endswith('.pkl'):
        filename = filename + '.pkl'
    with open(filename, 'wb') as f:
        pkl.dump(data, f)


def write_image(filename: str, image: Image.Image | np.ndarray, quality=90):
    """Writes a jpg image to disk, with a given quality"""
    create_folder(filename)
    if isinstance(image, np.ndarray):
        if len(image.shape) == 3:
            image = image[:, :, ::-1]
        if filename.endswith('.jpg'):  # Save as jpg with desired quality
            cv2.imwrite(filename, image, [int(cv2.IMWRITE_JPEG_QUALITY), quality])
        else:  # Save image with proper extension
            cv2.imwrite(filename, image)
    else:
        image.save(filename)


def write_npz(filename: str, data, bits=None):
    """Writes a npz file to disk"""
    for label in BIT_COMPRESS_LABELS:
        if label in data:
            data[label], encode_info = encode_dense(data[label], bits=bits)
            data['encode_info'] = encode_info
    filename = filename.removesuffix('.npz') + '.npz'
    create_folder(filename)
    np.savez_compressed(filename, **data)
    return filename


def write_json(filename, data):
    """Writes a json file to disk"""
    create_folder(filename)
    with open(filename, 'w') as f:
        json.dump(data, f, indent=4)


def write_png32(filename, data: np.ndarray):
    """Write 1 float32 image with 4 uint8 channels (RGBA)"""
    create_folder(filename)
    data = data.astype(np.float32)
    cv2.imwrite(filename, data.view((np.uint8, 4)))


def write_png8(filename, data: np.ndarray):
    """Write float32 as"""
    create_folder(filename)
    data = (data * 1000).astype(np.uint16)
    data = np.clip(data, 0, 2**16 - 1)
    cv2.imwrite(filename, data)


def write_image_seq(folder: str, rgbs: NT, ext: str = "jpg", frame_name_fn=frame_name, **kwargs):
    """Writes a sequence of (H, W, 3) images to disk, one per timestep"""
    os.makedirs(folder, exist_ok=True)
    if isinstance(rgbs, torch.Tensor):
        rgbs = rgbs.cpu().numpy()
    leftover_h = rgbs.shape[1] % 4
    if leftover_h > 0:
        rgbs = rgbs[:, :-leftover_h, :]
    leftover_w = rgbs.shape[2] % 4
    if leftover_w > 0:
        rgbs = rgbs[:, :, :-leftover_w]
    kwargs = {k:v for k,v in kwargs.items() if k in ['quality']}
    for t in range(rgbs.shape[0]):
        write_image(f"{folder}/{frame_name_fn(t)}.{ext}", rgbs[t], **kwargs)
    return folder


def write_video_seq(dest: str | None, rgbs: NT, ext: str = "mp4", fps: int = 10, crf=18, **kwargs) -> str | bytes:
    """
    Encodes a sequence of images and compiles into a video and saves to disk if dest is
    provided otherwise returns the encoded video bytes. 
    Expects input shape (T, H, W, 3) with values in [0, 255].
    """
    if isinstance(rgbs, np.ndarray):
        rgbs = torch.from_numpy(rgbs)
    leftover_h = rgbs.shape[1] % 4
    if leftover_h > 0:
        rgbs = rgbs[:, :-leftover_h, :]
    leftover_w = rgbs.shape[2] % 4
    if leftover_w > 0:
        rgbs = rgbs[:, :, :-leftover_w]
    assert rgbs.ndim == 4 and rgbs.shape[-1] == 3, "Expecting input to have shape (T, H, W, 3)"
    try:  # Use torchcodec as default
        encoder = VideoEncoder(frames=rgbs.permute(0, 3, 1, 2), frame_rate=float(fps))
    except Exception:  # Fallback to PyAV if needed
        encoder = PyAVEncoder(frames=rgbs, frame_rate=float(fps))
    if isinstance(dest, str):
        dest = dest.removesuffix(f'.{ext}') + f".{ext}"
        create_folder(dest)
        encoder.to_file(dest, crf=crf, extra_options={'g': '1'}, **kwargs)
    else:
        stream = io.BytesIO()
        encoder.to_file_like(stream, ext, crf=crf, extra_options={'g': '1'}, **kwargs)
        dest = stream.getvalue()
    del rgbs
    return dest


def write_npzs_seq(folder: str, data_dict: Dict[str, np.ndarray], frame_name_fn=frame_name):
    """Writes a sequence of npz files to disk, one per timestep"""
    os.makedirs(folder, exist_ok=True)
    encode_info = data_dict.pop('encode_info', None)  # remove info dict
    flat_dict = join_dict(data_dict)
    seq_lens = [len(flat_dict[key]) for key in flat_dict]
    assert all(l == seq_lens[0] for l in seq_lens), "All sequences must have the same length"
    for t in range(seq_lens[0]):
        frame_dict = {k: flat_dict[k][t] for k in flat_dict}
        if encode_info is not None:
            frame_dict['encode_info'] = encode_info  # Add info dict
        np.savez_compressed(f"{folder}/{frame_name_fn(t)}.npz", **frame_dict)
    del data_dict
    return folder


def write_zarr_seq(filename: str, data_dict: Dict[str, np.ndarray | dict]):
    """Writes a sequence of zarr arrays to disk. Supports nested dicts, which are stored as groups."""
    filename = filename.removesuffix(f'.zarr') + f".zarr"
    create_folder(filename)
    store = ZipStore(filename, mode='w')
    root = zarr.group(store)
    _write_zarr_group(root, data_dict)
    store.close()
    del data_dict
    return filename


def _write_zarr_group(group: zarr.Group, data_dict: Dict[str, np.ndarray | dict]):
    """Recursively writes a nested dict of numpy arrays into a zarr group."""
    for key, val in data_dict.items():
        if isinstance(val, dict):
            _write_zarr_group(group.require_group(key), val)
        elif isinstance(val, np.ndarray) and (
                np.issubdtype(val.dtype, np.number) or
                np.issubdtype(val.dtype, bool)
            ):
            group.create_array(key, data=val, chunks=(1, *val.shape[1:]) if val.ndim > 0 else "auto")
        else:
            print(f"Zarr is unreliable for non-numeric dtypes. Skipping '{key}'")


def write_seq(path: str, dict_or_array: Dict[str, np.ndarray] | NT, ext=None, sto=None, bits=None, **kwargs):
    """
    Writes a sequence of data to disk, one per timestep. Automatically determines whether to write 
    images or npz files based on the input type.
    """
    assert ext is not None or sto is not None, "Please provide extension or storage"

    # Stack frames together if it's a list
    if isinstance(dict_or_array, list):
        dict_or_array = np.stack(dict_or_array, 0)

    # Get extension automatically from storage + label
    if ext is None: 
        label = path.split('/')[-2]
        if label == 'rgb': 
            ext = rgb_ext(sto)
        elif label in BIT_COMPRESS_LABELS or label.startswith('mask_'): 
            ext = data_ext(sto)
            if not isinstance(dict_or_array, dict):
                if len(dict_or_array.shape) == 3:  # Expand to 1-channel
                    dict_or_array = np.expand_dims(dict_or_array, 1)
                else:  # Permute so channel is second dimension
                    dict_or_array = np.transpose(dict_or_array, (0, 3, 1, 2)).copy()
                dict_or_array = {label: dict_or_array}  # Convert to dictionary for storage
        else: 
            raise ValueError(f'Invalid label {label} and storage {sto}')

    # Zarr-specific logic (to avoid repetition in every converter)
    if isinstance(dict_or_array, dict):    
        for label, val in list(dict_or_array.items()):
            if label in BIT_COMPRESS_LABELS or label.startswith('mask_'):
                if label in BIT_COMPRESS_LABELS:  # dense labels to compress
                    data, encode_info = encode_dense(dict_or_array[label], bits=bits)
                else:  # dense labels to not compress (masks)
                    data, encode_info = dict_or_array[label], None
                dict_or_array[label] = data  # Store encoded value
                if encode_info is not None:  # Include info as well if available
                    dict_or_array['encode_info'] = encode_info

    sto = _norm_storage(sto)  # normalize legacy 'frame'/'sequence' to 'frames'/'videos'
    if ext in ['jpg', 'png'] and sto == 'frames':
        write_image_seq(path, dict_or_array, ext=ext, **kwargs)
    elif ext == 'mp4' and sto == 'videos':
        write_video_seq(path, dict_or_array, ext=ext, **kwargs)
    elif ext == 'zarr' and sto == 'videos':
        write_zarr_seq(path, dict_or_array, **kwargs)
    elif ext == 'npz' and sto in ['frames', 'videos']:
        (write_npz if sto == 'videos' else write_npzs_seq)(path, dict_or_array, **kwargs)
    else:
        raise ValueError(f"Unsupported combination of storage {sto} and extension {ext}")


# TODO: extract_frames_from_mp4() also exists, maybe merge to avoid code duplication
def write_jpg_from_mp4(folder, video, start_frame=0, end_frame=1e12):
    os.makedirs(folder, exist_ok=True)
    cmd = [
        "ffmpeg",
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-i",
        video,
        "-vf",
        f"select='between(n,{start_frame},{end_frame - 1})'",
        "-vsync",
        "vfr",
        "-frame_pts",
        "1",
        "-q:v",
        "2",
        f'{folder}/%010d.jpg',
    ]
    subprocess.run(cmd, capture_output=False)
    rgb_filenames = glob(f'{folder}/*.jpg')
    rgb = Image.open(rgb_filenames[0])
    return rgb.size[::-1], len(rgb_filenames)


def write_video(filename, video_array, fps=30):
    """
    Write video with PyAV compatibility handling for different versions.

    Args:
        filename: Output video file path
        video_array: Tensor or array of shape (T, H, W, 3) with values in [0, 255]
        fps: Frames per second
    """
    import torch
    import torchvision

    create_folder(filename)

    try:
        # Try using torchvision's write_video directly
        if isinstance(video_array, np.ndarray):
            video_array = torch.from_numpy(video_array)
        torchvision.io.write_video(filename, video_array, fps=fps)
    except TypeError as e:
        print(f"TypeError encountered during write_video: {e}")
        raise

def make_stacked_numpy(data):
    """Convert stacked lists to numpy arrays"""
    if isinstance(data, dict):
        for key, val in data.items():
            data[key] = make_stacked_numpy(val)
        return data
    elif isinstance(data, list):
        try:
            return np.array(data)
        except ValueError:
            # Ragged per-frame data (e.g. language prompt lists of differing
            # length when a frame falls on an action-segment boundary) can't be
            # stacked into a homogeneous array; keep it as an object array.
            arr = np.empty(len(data), dtype=object)
            arr[:] = data
            return arr
    else:
        return data


def write_lowdim(args, dst, labels, num_frames, lowdim: Dict[str, Dict[str, np.ndarray]]):
    """Write lowdim data to disk"""
    removed = [] # Log removed labels to not add again
    for key, val in lowdim.items():
        cam = key.split('/')[-2]
        for label, value in val.items():
            # Label has been removed, skip
            if label in removed:
                continue
            # Add label to labels
            if label not in labels:
                labels.append(label)
            # Start num_frames cound
            if label not in num_frames[cam]:
                num_frames[cam][label] = 0
            num_frames[cam][label] += 1
            # Label-specific checks
            if label in ['intrinsics','extrinsics']:
                ### Valid checks
                check_empty = np.absolute(value).sum() != 0 
                check_nan = not np.isnan(value.sum())
                check_inf = not np.isinf(value.sum())
                ### Remove label if it failed checks
                if not (check_empty and check_nan and check_inf):
                    if label in labels:
                        labels.remove(label)
                    removed.append(label)

    ### Video saving
    if args.storage == 'videos':
        cams = set([key.split('/')[-2] for key in lowdim.keys()])
        keys = sorted(list(lowdim.keys()))
        for cam in cams:
            lowdim_cam = stack_sample([lowdim[key] for key in keys if key.split('/')[-2] == cam])
            lowdim_cam = make_stacked_numpy(lowdim_cam)
            filename_lowdim_out = f'{dst}/lowdim/{cam}'
            write_npz(filename_lowdim_out, lowdim_cam)
    ### Frame saving
    else:
        for filename, val in lowdim.items():
            write_npz(filename, val)


def write_label(dst, cam, storage, label, data, labels, resolution, num_frames, bits=None):
    """Write a dense label to disk and store relevant information"""
    if len(data) > 0:
        ### Valid checks
        if isinstance(data, dict):
            vals = [data[key] for key in sorted(list(data.keys()))]
        else:
            vals = data
        check_shape = all([vals[0].shape == vals[i].shape for i in range(len(vals))])
        check_nan = all([not np.isnan(val.sum()) for val in vals])
        check_inf = all([not np.isinf(val.sum()) for val in vals])
        ### Remove label if it failed checks
        if not (check_shape and check_nan and check_inf):
            if label in labels:
                labels.remove(label)
            return
        ### Add label if not there
        if label not in labels:
            labels.append(label)
        ### Default bit values
        if bits is None:
            if label.startswith('mask'):
                bits = 1
            elif label.startswith('semantic'):
                bits = 8
            elif label.startswith('optflow'):
                bits = 24
        ### Write data (and delete from memory) and update dictionaries
        write_seq(f'{dst}/{label}/{cam}', vals, sto=storage, bits=bits)
        resolution[cam][label] = vals[0].shape[:2]
        num_frames[cam][label] = len(data)
        del data


def write_labels(dst, cam, storage, dense, labels, resolution, num_frames):
    """Write each label from a dict to disk"""
    for key, val in dense.items():
        write_label(dst, cam, storage, key, val, labels, resolution, num_frames)


class PyAVEncoder:
    """PyAV-backed drop-in replacement for torchcodec's VideoEncoder."""

    def __init__(self, frames: torch.Tensor, frame_rate=30):
        frames = frames.detach().cpu().numpy()
        assert frames.ndim == 4 and frames.shape[3] == 3, "Expecting frames to have shape (T, H, W, 3)"
        if np.issubdtype(frames.dtype, np.floating) and frames.max() <= 1.0:
            frames = frames * 255.0
        self.frames = np.clip(frames, 0, 255).astype(np.uint8, copy=False)
        self.frame_rate = int(frame_rate)

    def to_file(self, dest: str, **kwargs):
        with open(dest, 'wb') as f:
            self.to_file_like(f, ext=dest.split('.')[-1], **kwargs)

    def to_file_like(self, fileobj, ext: str ='mp4', *, codec=None, crf=None, extra_options=None):
        """
        codec: str, optional: The codec to use for encoding (e.g., 'libx264', 'h264'). 
            If not specified, the default codec will be used.
        crf: int or float, optional: The CRF (Constant Rate Factor) value for controlling video quality. 
            Lower values result in higher quality and larger file sizes, while higher values result in 
            lower quality and smaller file sizes. Typical values are between 18 (visually lossless) and 
            28 (more compressed). If not specified, the default CRF value for the codec will be used.
        extra_options: dict, optional: Additional encoding options to pass to the encoder.
        """
        import av
        from av.video.stream import VideoStream
        with av.open(fileobj, mode='w', format=ext) as container:
            stream: VideoStream = container.add_stream(codec or 'libx264', rate=self.frame_rate)
            stream.width = int(self.frames.shape[2])
            stream.height = int(self.frames.shape[1])
            stream.pix_fmt = 'yuv420p'
            stream.options = {str(k): str(v) for k,v in (extra_options or {}).items()}
            if crf is not None:
                stream.options['crf'] = str(crf)
            for frame in self.frames:
                video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
                for packet in stream.encode(video_frame):
                    container.mux(packet)
            for packet in stream.encode():
                container.mux(packet)