"""Build a self-contained v3 dataset from completed annotations, one episode at a time."""

import json
from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
from fractions import Fraction
from pathlib import Path

import av
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.parquet as pq
from studio import relative_path, validate_subset_features

from lerobot.datasets.compute_stats import aggregate_stats, get_feature_stats
from lerobot.datasets.io_utils import load_tasks, write_stats, write_tasks

DATA_PATH = "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet"
VIDEO_PATH = "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4"
EPISODES_PATH = "meta/episodes/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet"


def trim_video(source, destination, offset, length, fps, *, lossless=False):
    """Seek to an episode and encode exactly its frames, with timestamps starting at zero."""
    destination.parent.mkdir(parents=True, exist_ok=True)
    rate = Fraction(str(fps)).limit_denominator(100000)
    with av.open(str(source)) as src, av.open(str(destination), "w") as dst:
        stream = src.streams.video[0]
        stream.codec_context.thread_count = 2
        output = dst.add_stream(
            "libx264rgb" if lossless else "libx264",
            rate=rate,
            options={"crf": "0" if lossless else "18", "preset": "fast", "g": "30"},
        )
        output.width, output.height = stream.codec_context.width, stream.codec_context.height
        output.pix_fmt = "rgb24" if lossless else "yuv420p"
        output.codec_context.thread_count = 2
        output.time_base = 1 / rate
        src.seek(int(offset / stream.time_base), stream=stream, backward=True)
        count = 0
        for frame in src.decode(stream):
            if frame.pts is None:
                continue
            timestamp = float(frame.pts * stream.time_base)
            if timestamp < offset - 0.5 / fps:
                continue
            if abs(timestamp - (offset + count / fps)) > 0.51 / fps:
                raise ValueError("Source video does not match the episode's frame timestamps.")
            encoded = frame.reformat(format=output.pix_fmt)
            encoded.pts, encoded.time_base = count, 1 / rate
            for packet in output.encode(encoded):
                dst.mux(packet)
            count += 1
            if count == length:
                break
        if count != length:
            raise ValueError(f"Source video ended early: expected {length} frames, found {count}.")
        for packet in output.encode():
            dst.mux(packet)
    # Decode metadata alone is insufficient: explicitly verify the encoded frame count.
    with av.open(str(destination)) as check:
        if sum(1 for _ in check.decode(video=0)) != length:
            raise ValueError("Trimmed video frame count does not match the episode.")


class EpisodeSubset:
    """Only completed episodes enter this writer. Upload each returned file set atomically."""

    def __init__(self, source_root, root, source_ids, camera_workers=1):
        if camera_workers not in (1, 2, 3):
            raise ValueError("Export supports one to three camera workers")
        self.camera_workers = camera_workers
        self.source_root, self.root = Path(source_root), Path(root)
        self.root.mkdir(parents=True, exist_ok=True)
        self.info = json.loads((self.source_root / "meta/info.json").read_text())
        validate_subset_features(self.info)
        self.source_info = deepcopy(self.info)
        self.cameras = [k for k, f in self.info["features"].items() if f["dtype"] == "video"]
        self.source_tasks = load_tasks(self.source_root)
        metadata = pq.read_table(
            self.source_root / "meta/episodes", filters=[("episode_index", "in", list(source_ids))]
        )
        self.source_rows = {int(r["episode_index"]): r for r in metadata.to_pylist()}
        if set(self.source_rows) != set(source_ids):
            raise ValueError("Episode subset metadata does not match the selected IDs.")
        self.mapping, self.tasks, self.task_mapping = [], [], {}
        self.total_frames, self.stats = 0, {}
        self.info.update(data_path=DATA_PATH, video_path=VIDEO_PATH, chunks_size=1000)
        for camera in self.cameras:
            video_info = self.info["features"][camera].setdefault("info", {})
            video_info.update(
                {
                    "video.codec": "h264",
                    "video.pix_fmt": "yuv420p",
                    "video.fps": self.info["fps"],
                    "video.g": 30,
                    "video.crf": 18,
                    "video.preset": "fast",
                    "video.extra_options": {},
                    "has_audio": False,
                }
            )

    def add(self, source_id, table, *, hand_features=None):
        """Materialize an annotated episode, retaining all its cameras and existing other styles."""
        if any(m["source_episode_index"] == source_id for m in self.mapping):
            raise ValueError("Episode was already added to the subset.")
        row = self.source_rows[source_id]
        table = table.filter(pc.equal(table["episode_index"], source_id))
        length = int(row["length"])
        if len(table) != length or table["frame_index"].to_pylist() != list(range(length)):
            raise ValueError("Episode data length or frame order does not match its metadata.")
        if hand_features:
            description = {
                "coordinates": "fixed segment-start FLU (X forward, Y left, Z up)",
                "wrist_position_units": "inferred metres; monocular metric scale is model-dependent",
                "rotation_units": "axis-angle radians",
                "validity_feature": "annotation.hand_action_valid",
                "observed_training_feature": "annotation.hand_action_observed_valid",
                "motion_infilling": False,
                "reference_feature": "annotation.hand_reference_frame",
                "invalid_action": "zero placeholder; must be masked during training",
                "action_statistics": "computed over valid predictions only; absent when none are available",
                "minimum_segment_frames": 30,
            }
            for name, feature in hand_features.items():
                self.info["features"][name] = deepcopy(feature)
            self.info["features"]["action"]["info"] = description
            if "action" in self.source_info["features"]:
                self.info["features"]["annotation.original_action"] = deepcopy(
                    self.source_info["features"]["action"]
                )
        new_id = len(self.mapping)
        chunk, file = divmod(new_id, 1000)
        changes = {
            "episode_index": [new_id] * length,
            "index": list(range(self.total_frames, self.total_frames + length)),
        }
        source_task_names = {int(v): str(k) for k, v in self.source_tasks["task_index"].items()}
        task_indices = table["task_index"].to_pylist()
        for old in sorted(set(task_indices)):
            name = source_task_names[int(old)]
            if name not in self.tasks:
                self.tasks.append(name)
            self.task_mapping[int(old)] = self.tasks.index(name)
        changes["task_index"] = [self.task_mapping[int(old)] for old in task_indices]
        for name, values in changes.items():
            table = table.set_column(
                table.schema.get_field_index(name),
                table.schema.field(name),
                pa.array(values, type=table[name].type),
            )
        paths = []
        episode = {
            "episode_index": new_id,
            "tasks": list(dict.fromkeys(source_task_names[int(i)] for i in task_indices)),
            "length": length,
            "dataset_from_index": self.total_frames,
            "dataset_to_index": self.total_frames + length,
            "data/chunk_index": chunk,
            "data/file_index": file,
            "meta/episodes/chunk_index": chunk,
            "meta/episodes/file_index": file,
        }

        def trim_camera(camera):
            source_video = self.source_root / relative_path(
                self.source_info["video_path"].format(
                    video_key=camera,
                    chunk_index=int(row[f"videos/{camera}/chunk_index"]),
                    file_index=int(row[f"videos/{camera}/file_index"]),
                    episode_index=row.get("episode_index"),
                )
            )
            target = self.root / relative_path(
                VIDEO_PATH.format(video_key=camera, chunk_index=chunk, file_index=file)
            )
            trim_video(
                source_video,
                target,
                float(row[f"videos/{camera}/from_timestamp"]),
                length,
                float(self.info["fps"]),
            )
            return target

        if self.camera_workers > 1:
            with ThreadPoolExecutor(max_workers=min(self.camera_workers, len(self.cameras))) as pool:
                targets = list(pool.map(trim_camera, self.cameras))
        else:
            targets = [trim_camera(camera) for camera in self.cameras]
        for camera, target in zip(self.cameras, targets, strict=True):
            paths.append(target)
            episode.update(
                {
                    f"videos/{camera}/chunk_index": chunk,
                    f"videos/{camera}/file_index": file,
                    f"videos/{camera}/from_timestamp": 0.0,
                    f"videos/{camera}/to_timestamp": length / self.info["fps"],
                }
            )
        data_path = self.root / DATA_PATH.format(chunk_index=chunk, file_index=file)
        data_path.parent.mkdir(parents=True, exist_ok=True)
        pq.write_table(table, data_path, row_group_size=length)
        paths.append(data_path)
        # Numeric/video stats describe these same frames; only remapped indices need recomputing.
        episode_stats = {}
        for key, value in row.items():
            if key.startswith("stats/") and value is not None:
                feature, stat = key[len("stats/") :].rsplit("/", 1)
                if feature in self.info["features"]:
                    episode_stats.setdefault(feature, {})[stat] = np.atleast_1d(np.asarray(value))
        for feature in changes:
            episode_stats[feature] = get_feature_stats(np.asarray(changes[feature]), axis=0, keepdims=True)
        if hand_features:
            for feature in [*hand_features, "annotation.original_action"]:
                if feature in table.column_names:
                    values = np.asarray(table[feature].to_pylist())
                    if feature == "action":
                        values = values[
                            np.asarray(table["annotation.hand_action_valid"].to_pylist(), dtype=bool)
                        ]
                        if not len(values):
                            episode_stats.pop(feature, None)
                            continue
                    episode_stats[feature] = get_feature_stats(values, axis=0, keepdims=values.ndim == 1)
        for feature, stats in episode_stats.items():
            for stat, value in stats.items():
                episode[f"stats/{feature}/{stat}"] = value.tolist()
        episode_path = self.root / EPISODES_PATH.format(chunk_index=chunk, file_index=file)
        episode_path.parent.mkdir(parents=True, exist_ok=True)
        episode_table = pa.Table.from_pylist([episode])
        if hand_features:
            # A first unobserved episode must declare the same nullable stats
            # columns as later observed/infilled episodes. Null payloads carry
            # no measured action statistics and never enter global aggregation.
            for stat in ("min", "max", "mean", "std", "count", "q01", "q10", "q50", "q90", "q99"):
                name = f"stats/action/{stat}"
                values = episode.get(name, [0] if stat == "count" else None)
                dtype = pa.list_(pa.int64() if stat == "count" else pa.float64())
                array = pa.array([values], type=dtype)
                index = episode_table.schema.get_field_index(name)
                episode_table = (
                    episode_table.set_column(index, name, array)
                    if index >= 0
                    else episode_table.append_column(name, array)
                )
        pq.write_table(episode_table, episode_path)
        paths.append(episode_path)
        self.stats = aggregate_stats([self.stats, episode_stats])
        self.total_frames += length
        self.mapping.append({"episode_index": new_id, "source_episode_index": source_id})
        self.info.update(
            total_episodes=len(self.mapping),
            total_frames=self.total_frames,
            total_tasks=len(self.tasks),
            splits={"train": f"0:{len(self.mapping)}"},
        )
        if "total_videos" in self.info:
            self.info["total_videos"] = len(self.mapping) * len(self.cameras)
        if "total_chunks" in self.info:
            self.info["total_chunks"] = chunk + 1
        # The annotation writer may have added language schemas and tool definitions.
        annotated_info = json.loads((self.source_root / "meta/info.json").read_text())
        for name, feature in annotated_info["features"].items():
            if feature["dtype"] == "language":
                self.info["features"][name] = feature
        if "tools" in annotated_info:
            self.info["tools"] = annotated_info["tools"]
        info_path = self.root / "meta/info.json"
        info_path.write_text(json.dumps(self.info, indent=2) + "\n")
        write_stats(self.stats, self.root)
        write_tasks(
            pd.DataFrame({"task_index": range(len(self.tasks))}, index=pd.Index(self.tasks, name="task")),
            self.root,
        )
        map_path = self.root / "annotation_studio/episode_mapping.json"
        map_path.parent.mkdir(parents=True, exist_ok=True)
        map_path.write_text(json.dumps(self.mapping, indent=2) + "\n")
        paths.extend([info_path, self.root / "meta/stats.json", self.root / "meta/tasks.parquet", map_path])
        return new_id, paths