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

import json
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):
    """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("libx264", rate=rate, options={"crf": "18", "preset": "fast", "g": "30"})
        output.width, output.height = stream.codec_context.width, stream.codec_context.height
        output.pix_fmt = "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="yuv420p")
            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):
        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):
        """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.")
        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,
        }
        for camera in self.cameras:
            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"]),
                )
            )
            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"]),
            )
            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)
        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)
        pq.write_table(pa.Table.from_pylist([episode]), 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