import json
import torch
import torch.nn as nn

from diffsynth.core.data.operators import ImageCropAndResize, DataProcessingOperator
from diffsynth.core.data.custom_operators import LoadVideoRange, TailPadFrames
from diffsynth.core.data.custom_units import WanVideoUnit_CutInjector
from diffsynth.core.data.custom_model_fn import model_fn_wan_video_with_cut


class BaseDataProfile:
    def __init__(self, args):
        self.args = args

    def load_and_transform(self, path):
        raise NotImplementedError

    def get_operator_map(self):
        return {}

    def get_data_keys(self):
        return []

    def get_extra_inputs(self):
        return []

    def configure_pipeline(self, pipe):
        pass


class PassThroughOp(DataProcessingOperator):
    def __call__(self, data):
        return data


class VidEventProfile(BaseDataProfile):
    """Data profile for ShotPlan multi-shot training samples.

    Expects a JSON list of records:
    {
        "file_path": "path/to/video.mp4",
        "start_frame": 102,
        "end_frame": 182,
        "cut_at": [26, 64],       # cut positions, frames relative to start_frame
        "type": "hardcut",
        "text": "Global caption ... Shot 1: ... Shot 2: ..."
    }
    """

    def load_and_transform(self, metadata_path):
        with open(metadata_path, 'r', encoding='utf-8') as f:
            raw_data = json.load(f)

        flattened_data = []
        for row in raw_data:
            path = row['file_path']
            start = row.get('start_frame', 0)
            end = row.get('end_frame', -1)

            cut_at = row.get('cut_at', [])
            item_type = row.get('type', 'hardcut')

            item = {
                "video": {
                    "path": path,
                    "start_frame": start,
                    "end_frame": end
                },
                "cut_info": {
                    "cuts": cut_at,
                    "type": item_type
                },
                "prompt": row.get('text', '')
            }
            flattened_data.append(item)

        return flattened_data

    def get_operator_map(self):
        resize_op = ImageCropAndResize(
            height=self.args.height, width=self.args.width,
            max_pixels=self.args.max_pixels,
            height_division_factor=16, width_division_factor=16
        )
        # TailPadFrames pads the payload (e.g. 80 frames) to 4k+1 (e.g. 81)
        # by repeating the last frame, so the loader never reads past the
        # window into the next shot of the source video.
        return {
            "video": LoadVideoRange(frame_processor=resize_op) >> TailPadFrames(target_len=81),
            "cut_info": PassThroughOp()
        }

    def get_data_keys(self):
        return ["video", "cut_info"]

    def get_extra_inputs(self):
        return ["cut_info"]

    def configure_pipeline(self, pipe):
        print("[Profile] Configuring pipeline for planning-token injection...")

        # A single learnable planning token, hardcut_embedding, is used.
        dim = pipe.dit.dim

        def register_token(name):
            param_name = f"{name}_embedding"
            if not hasattr(pipe.dit, param_name):
                print(f"  -> Registering token: {param_name}")
                token_tensor = torch.randn(1, 1, dim) * 0.02
                cut_param = nn.Parameter(token_tensor)
                pipe.dit.register_parameter(param_name, cut_param)
                getattr(pipe.dit, param_name).requires_grad = True

        register_token("hardcut")

        # Insert the injector unit right after the noise initializer.
        pipe.units = [u for u in pipe.units if not isinstance(u, WanVideoUnit_CutInjector)]

        insert_index = 0
        for i, u in enumerate(pipe.units):
            if u.__class__.__name__ == "WanVideoUnit_NoiseInitializer":
                insert_index = i + 1
                break

        pipe.units.insert(insert_index, WanVideoUnit_CutInjector())

        pipe.model_fn = model_fn_wan_video_with_cut

        if hasattr(pipe.dit, "require_vae_embedding"):
            pipe.dit.require_vae_embedding = True

        print("[Profile] Pipeline configured.")