#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Nov  1 17:02:10 2025

@author: blissadm
"""

import os

import imageio
import numpy as np
from ewokscore import Task
from silx.io import h5py_utils

DEFAULTS = {"detector_name": None, "target_width": 200, "max_size": 200 * 1024}


class AnimatedWebP(
    Task,
    input_names=[
        "nx_file",
    ],
    optional_input_names=["detector_name"],
    output_names=["nx_file", "detector_name"],
):
    """Export an animated WebP file from a nx-fluotomo"""

    def bin_stack(self, img, bin_y=2, bin_x=2):
        n, h, w = img.shape

        h2 = h // bin_y * bin_y
        w2 = w // bin_x * bin_x
        img = img[:, :h2, :w2]

        return img.reshape(n, h2 // bin_y, bin_y, w2 // bin_x, bin_x).mean(axis=(2, 4))

    def run(self):

        pars = {
            **DEFAULTS,
            **self.get_input_values(),
        }

        nx_file = pars["nx_file"]
        detector_name = pars["detector_name"]
        t_width = pars["target_width"]
        max_size = pars["max_size"]

        processed_data_folder = os.path.dirname(nx_file)
        galleryfolder = os.path.join(processed_data_folder, "gallery")
        os.makedirs(galleryfolder, exist_ok=True)
        os.chmod(
            galleryfolder, 0o770
        )  # nosec B103: group write required for shared pipeline

        with h5py_utils.open_item(nx_file, "/") as fd:
            for e in fd:

                if detector_name is not None and not e.startswith(detector_name):
                    continue

                data = fd[e]["data/data"][()]
                data[np.isnan(data)] = 0

                width = data.shape[2]

                binning = int(np.round(width / t_width))

                if binning < 0:
                    binning = 1

                data = self.bin_stack(data, binning, binning)

                mn = np.min(data)
                Mx = np.max(data)

                converted = np.array((data - mn) / (Mx - mn) * 255, dtype=np.uint8)

                #                imageio.v3.imwrite(
                #                    os.path.join(galleryfolder, f"{e}.webp"), converted,  quality=40, method=4,
                #                )

                encoded = imageio.v3.imwrite(
                    "<bytes>",
                    converted,
                    fps=5,
                    lossless=False,
                    extension=".webp",
                    quality=70,
                    method=4,
                )

                size = len(encoded)

                if size < max_size:
                    with open(os.path.join(galleryfolder, f"{e}.webp"), "wb") as f:
                        f.write(encoded)

        self.outputs.nx_file = nx_file
        self.outputs.detector_name = detector_name
