#!/usr/bin/env python3
#
# Minimal, self-contained reproducer for a babeltrace2 CTF-2 sink bug.
#
# A single trace that contains two data stream classes sharing ONE clock class
# is written to CTF 2 by sink.ctf.fs. The sink emits one `clock-class` metadata
# fragment PER stream class (both with the same id, here "the_clock") instead
# of emitting the shared clock class only once. Reading the resulting trace
# back then fails with:
#
#   Duplicate clock class fragment with ID `the_clock`.
#
# This depends only on babeltrace2 and its bundled `ctf`/`utils` plugins; no
# external trace files or third-party plugins are needed.
#
# Usage:
#   python3 repro.py <output-dir>
#
# The script writes the trace in a child process (so sink.ctf.fs flushes the
# metadata file on component teardown), then reads it back in the parent.
import os
import subprocess
import sys

import bt2


class TheIter(bt2._UserMessageIterator):
    def __init__(self, config, port):
        stream = port.user_data
        self._msgs = [
            self._create_stream_beginning_message(stream),
            self._create_stream_end_message(stream),
        ]
        self._i = 0

    def __next__(self):
        if self._i >= len(self._msgs):
            raise StopIteration
        msg = self._msgs[self._i]
        self._i += 1
        return msg


@bt2.plugin_component_class
class TheSource(bt2._UserSourceComponent, message_iterator_class=TheIter):
    """One trace, two stream classes, ONE shared clock class."""

    def __init__(self, config, params, obj):
        tc = self._create_trace_class(assigns_automatic_stream_class_id=False)
        cc = self._create_clock_class(frequency=1000000000, name="the_clock")
        sc0 = tc.create_stream_class(
            id=0,
            default_clock_class=cc,
            supports_packets=True,
            assigns_automatic_stream_id=False,
        )
        sc1 = tc.create_stream_class(
            id=1,
            default_clock_class=cc,
            supports_packets=True,
            assigns_automatic_stream_id=False,
        )
        trace = tc()
        self._add_output_port("out0", trace.create_stream(sc0, id=0))
        self._add_output_port("out1", trace.create_stream(sc1, id=1))


def write_trace(out_dir):
    ctf = bt2.find_plugin("ctf")
    utils = bt2.find_plugin("utils")
    sink_cls = ctf.sink_component_classes["fs"]

    # CTF 2 requires MIP 1. The ctf.fs sink has a single input port, so mux
    # the two source ports into it.
    graph = bt2.Graph(mip_version=1)
    src = graph.add_component(TheSource, "src")
    muxer = graph.add_component(utils.filter_component_classes["muxer"], "mux")
    sink = graph.add_component(
        sink_cls, "sink", params={"path": out_dir, "ctf-version": "2"}
    )
    out_ports = list(src.output_ports.values())
    graph.connect_ports(out_ports[0], list(muxer.input_ports.values())[0])
    graph.connect_ports(out_ports[1], list(muxer.input_ports.values())[1])
    graph.connect_ports(
        list(muxer.output_ports.values())[0],
        list(sink.input_ports.values())[0],
    )
    graph.run()


def read_trace(trace_dir):
    for _ in bt2.TraceCollectionMessageIterator(trace_dir):
        pass


def main():
    if len(sys.argv) < 2:
        print("usage: repro.py <output-dir>", file=sys.stderr)
        return 2

    # Internal entry points, invoked as child processes.
    if sys.argv[1] == "--write":
        write_trace(sys.argv[2])
        return 0
    if sys.argv[1] == "--read":
        read_trace(sys.argv[2])
        return 0

    out_dir = sys.argv[1]
    trace_dir = os.path.join(out_dir, "trace")

    # Write in a child process so the sink flushes the metadata on teardown.
    subprocess.run([sys.executable, __file__, "--write", out_dir], check=True)
    print("wrote CTF-2 trace to", trace_dir)

    with open(os.path.join(trace_dir, "metadata"), "rb") as f:
        n_clock = f.read().count(b'"type":"clock-class"')
    print("clock-class fragments written to metadata:", n_clock)

    # Read back with the CLI for a clear, unwrapped error message.
    proc = subprocess.run(
        ["babeltrace2", trace_dir],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
        text=True,
    )
    if proc.returncode != 0:
        print("\nREAD-BACK FAILED (bug reproduced):", file=sys.stderr)
        for line in proc.stderr.splitlines():
            if "Duplicate clock class" in line or "Invalid fragment" in line:
                print("  " + line.split("] ", 1)[-1], file=sys.stderr)
        return 1

    print("read-back OK (bug NOT reproduced)")
    return 0


if __name__ == "__main__":
    sys.exit(main())