#!/usr/bin/env python
###############################################################################
# (c) Copyright 2024 CERN for the benefit of the LHCb Collaboration           #
#                                                                             #
# This software is distributed under the terms of the GNU General Public      #
# Licence version 3 (GPL Version 3), copied verbatim in the file "COPYING".  #
#                                                                             #
# In applying this licence, CERN does not waive the privileges and immunities #
# granted to it by virtue of its status as an Intergovernmental Organization  #
# or submit itself to any jurisdiction.                                       #
###############################################################################
"""Re-slice an existing local DST file to a smaller event range.

Use this after capture-fixture to narrow down the DST to only the events
that make the CWL test workflows produce output.

On first run the original file is backed up to <file>.bak and the sliced
output replaces it.  On subsequent runs the .bak (unsliced original) is
used as the source so you can refine first_evt / evt_max without losing
the full file.

Prerequisites:
    - CVMFS available (for lb-run)

Usage:
    pixi run reslice-fixture <dst-file> --first-evt 50 --evt-max 5

Example:
    pixi run reslice-fixture tests/test-input-dst/00070793_00000685_7.AllStreams_slice.dst --first-evt 15 --evt-max 20
"""

from __future__ import annotations

import argparse
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

import yaml


def main():
    parser = argparse.ArgumentParser(
        description="Re-slice an existing local DST to a smaller event range",
    )
    parser.add_argument(
        "dst_file", type=Path, help="Path to the local DST file to re-slice"
    )
    parser.add_argument(
        "--first-evt",
        type=int,
        default=0,
        help="First event index to keep (default: 0)",
    )
    parser.add_argument(
        "--evt-max",
        type=int,
        default=10,
        help="Number of events to keep (default: 10)",
    )
    args = parser.parse_args()

    dst_file = args.dst_file.resolve()
    bak_file = dst_file.with_suffix(dst_file.suffix + ".bak")

    # Determine the source file to slice from:
    # - If a .bak exists, always slice from it (the unsliced original)
    # - Otherwise slice from the dst_file itself (first run)
    if bak_file.exists():
        source_file = bak_file
        print(f"Backup:    {bak_file} (using existing backup as source)")
    elif dst_file.exists():
        source_file = dst_file
    else:
        print(f"ERROR: DST file not found: {dst_file}", file=sys.stderr)
        sys.exit(1)

    original_size = source_file.stat().st_size
    print(f"Input:     {source_file} ({original_size / 1024:.0f} KB)")
    print(f"Range:     first_evt={args.first_evt}, evt_max={args.evt_max}")

    with tempfile.TemporaryDirectory() as tmpdir:
        tmpdir = Path(tmpdir)
        output_file = tmpdir / dst_file.name

        slice_config = {
            "input_files": [str(source_file)],
            "first_evt": args.first_evt,
            "evt_max": args.evt_max,
            "output_file": str(output_file),
            "simulation": True,
            "input_process": "Hlt2",
            "input_type": "ROOT",
            "input_raw_format": 0.5,
            "data_type": "Upgrade",
            "geometry_version": "run3/2024.Q1.2-v00.00",
            "conditions_version": "master",
            "output_type": "ROOT",
            "compression": "ZSTD:6",
            "root_ioalg_name": "RootIOAlgExt",
        }

        config_path = tmpdir / "reslice_config.yaml"
        with open(config_path, "w") as f:
            yaml.dump(slice_config, f, default_flow_style=False, sort_keys=False)

        print(
            f"Running:   lb-run LHCb/v58r2 lbexec GaudiConf.mergeDST:dst {config_path}"
        )
        result = subprocess.run(
            [
                "lb-run",
                "LHCb/v58r2",
                "lbexec",
                "GaudiConf.mergeDST:dst",
                str(config_path),
            ],
        )
        if result.returncode != 0:
            print(
                f"ERROR: Slice failed with exit code {result.returncode}",
                file=sys.stderr,
            )
            sys.exit(1)

        if not output_file.exists():
            print(f"ERROR: Output file not created: {output_file}", file=sys.stderr)
            sys.exit(1)

        new_size = output_file.stat().st_size

        # Back up the original on first run (before any slicing)
        if not bak_file.exists():
            shutil.move(str(dst_file), str(bak_file))
            print(f"Backup:    {dst_file} -> {bak_file}")

        shutil.move(str(output_file), str(dst_file))

    print(f"Output:    {dst_file} ({new_size / 1024:.0f} KB)")
    print(
        f"Reduction: {original_size / 1024:.0f} KB -> {new_size / 1024:.0f} KB "
        f"({100 * (1 - new_size / original_size):.0f}% smaller)"
    )


if __name__ == "__main__":
    main()
