#!/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.                                       #
###############################################################################
"""Capture test fixtures for CWL-based tests.

Fully automated: captures DIRAC metadata, picks an LFN, resolves its
access URL, generates a slice config, runs the slice, and saves all
fixture files.  Tests auto-discover new fixtures — no test code changes
needed.

Prerequisites:
    - Valid DIRAC proxy  (lb-dirac-proxy --init)
    - CVMFS available   (for lb-run)
    - Run from inside an AP data-package repo

Usage:
    pixi run capture-fixture <production> <job>

Example:
    pixi run capture-fixture executabletests 2016_MagDown_PromptMC_D02KK
"""

from __future__ import annotations

import argparse
import importlib.resources
import json
import subprocess
import sys
import tempfile
from pathlib import Path

import yaml

TESTS_DIR = Path(__file__).parent
FIXTURES_DIR = TESTS_DIR / "fixtures"
TEST_DST_DIR = TESTS_DIR / "test-input-dst"
DATA_DIR = TESTS_DIR / "data"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def run_checked(
    cmd: list[str], *, capture: bool = True, **kwargs
) -> subprocess.CompletedProcess:
    """Run a command, printing it first, and raising on failure."""
    print(f"  $ {' '.join(str(c) for c in cmd)}")
    result = subprocess.run(
        cmd,
        capture_output=capture,
        text=True,
        **kwargs,
    )
    if result.returncode != 0:
        stderr = result.stderr if capture else ""
        raise RuntimeError(
            f"Command failed (exit {result.returncode}):\n"
            f"  {' '.join(str(c) for c in cmd)}\n{stderr}"
        )
    return result


# ---------------------------------------------------------------------------
# Step 1 — query_dirac  (via lb-dirac, same mechanism as LbAPCommon)
# ---------------------------------------------------------------------------


def run_query_dirac(input_query: dict, turbo: bool) -> dict:
    """Call LbAPCommon's query_dirac script via lb-dirac and return the result."""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Copy the query_dirac.py script (same as LbAPCommon.__main__ does)
        query_script = Path(tmpdir) / "query_dirac.py"
        query_script.write_text(
            (
                importlib.resources.files("LbAPCommon.data") / "query_dirac.py"
            ).read_text()
        )
        query_script.chmod(0o700)

        output_path = Path(tmpdir) / "query_result.json"
        cmd = [
            "lb-dirac",
            str(query_script),
            json.dumps(input_query),
            f"--output={output_path}",
        ]
        if turbo:
            cmd.append("--turbo")

        run_checked(cmd)
        return json.loads(output_path.read_text())


# ---------------------------------------------------------------------------
# Step 2 — BK query for LFNs  (via lb-dirac subprocess)
# ---------------------------------------------------------------------------


def query_bk_for_lfns(
    conditions_dict: dict,
    conditions_description: str,
    event_type: str,
    input_spec: dict,
) -> list[str]:
    """Query Bookkeeping for a single LFN matching the input dataset criteria."""
    # Self-contained script using only DIRAC/LHCbDIRAC imports (no LbAPLocal)
    script = """\
import json, sys, random
import DIRAC
DIRAC.initialize()
from DIRAC.Core.Utilities.ReturnValues import returnValueOrRaise
from DIRAC.DataManagementSystem.Client.DataManager import DataManager
from LHCbDIRAC.BookkeepingSystem.Client.BookkeepingClient import BookkeepingClient

spec = json.loads(sys.argv[1])
num_lfns = int(sys.argv[2])
cd = spec["conditions_dict"]

bk_query = {
    "FileType": cd["inFileType"],
    "EventType": spec["event_type"],
    "ConfigName": cd["configName"],
    "ConfigVersion": cd["configVersion"],
}
if spec.get("conditions_description"):
    bk_query["DataTakingConditions"] = spec["conditions_description"]
if cd.get("inProPass"):
    bk_query["ProcessingPass"] = cd["inProPass"]
dq = cd.get("inDataQualityFlag", "OK")
bk_query["DataQualityFlag"] = dq.split(",") if isinstance(dq, str) else dq
if cd.get("inProductionID") and cd["inProductionID"] != "ALL":
    bk_query["ProductionID"] = cd["inProductionID"]

result = returnValueOrRaise(
    BookkeepingClient().getFilesWithMetadata(
        bk_query | {"OnlyParameters": ["FileName", "FileSize"]}
    )
)
if result["TotalRecords"] == 0:
    raise ValueError(f"No files found for query: {json.dumps(bk_query, indent=2)}")

fi = result["ParameterNames"].index("FileName")
si = result["ParameterNames"].index("FileSize")
records = sorted(result["Records"], key=lambda r: r[si])
# Pick smallest files for faster slicing
records = records[:max(num_lfns * 3, 5)]
random.shuffle(records)

# Check for available replicas
dm = DataManager()
lfns = []
for rec in records:
    lfn = rec[fi]
    rep = returnValueOrRaise(dm.getReplicasForJobs([lfn], diskOnly=True))
    if rep["Successful"]:
        lfns.extend(rep["Successful"])
        if len(lfns) >= num_lfns:
            break

print(json.dumps(lfns[:num_lfns]))
"""

    dataset_spec = {
        "conditions_dict": conditions_dict,
        "conditions_description": conditions_description,
        "event_type": event_type,
    }

    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
        f.write(script)
        script_path = f.name

    try:
        result = run_checked(
            ["lb-dirac", "python", script_path, json.dumps(dataset_spec), "1"]
        )
        return json.loads(result.stdout.strip().split("\n")[-1])
    finally:
        Path(script_path).unlink(missing_ok=True)


# ---------------------------------------------------------------------------
# Step 3 — resolve access URL  (via lb-dirac dirac-dms-lfn-accessURL)
# ---------------------------------------------------------------------------


def resolve_access_url(lfn: str) -> str:
    """Resolve an LFN to a root:// access URL via DIRAC."""
    result = run_checked(["lb-dirac", "dirac-dms-lfn-accessURL", lfn])
    # Output contains lines like:
    #   /lhcb/...dst : root://host:port//path/...dst
    # Extract the first root:// URL found anywhere in the output.
    for line in result.stdout.strip().split("\n"):
        idx = line.find("root://")
        if idx != -1:
            return line[idx:].strip()
    raise RuntimeError(
        f"Could not extract root:// URL from dirac-dms-lfn-accessURL output:\n{result.stdout}"
    )


# ---------------------------------------------------------------------------
# Step 4 — generate slice config and run the slice
# ---------------------------------------------------------------------------


def generate_slice_config(
    access_url: str,
    output_filename: str,
    first_evt: int = 0,
    evt_max: int = 100,
) -> dict:
    """Generate a minimal slice YAML config for GaudiConf.mergeDST:dst."""
    return {
        "input_files": [access_url],
        "first_evt": first_evt,
        "evt_max": evt_max,
        "output_file": output_filename,
        "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",
    }


def run_slice(slice_yaml_path: Path) -> None:
    """Run the DST slicer via lb-run."""
    run_checked(
        [
            "lb-run",
            "LHCb/v58r2",
            "lbexec",
            "GaudiConf.mergeDST:dst",
            str(slice_yaml_path),
        ],
        capture=False,
    )


# ---------------------------------------------------------------------------
# Step 5 — update manifest.yaml
# ---------------------------------------------------------------------------


def update_manifest(
    fixture_dir: Path,
    production_name: str,
    ap_repo: str,
    job_name: str,
    dst_file_rel: str,
    lfn: str,
) -> None:
    """Create or update the manifest.yaml with the new job entry."""
    manifest_path = fixture_dir / "manifest.yaml"
    if manifest_path.exists():
        with open(manifest_path) as f:
            manifest = yaml.safe_load(f)
    else:
        manifest = {
            "production": production_name,
            "ap_repo": ap_repo,
            "jobs": {},
        }

    manifest["jobs"][job_name] = {
        "dst_file": dst_file_rel,
        "lfn": lfn,
    }

    with open(manifest_path, "w") as f:
        yaml.dump(manifest, f, default_flow_style=False, sort_keys=False)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main():
    parser = argparse.ArgumentParser(
        description="Capture CWL test fixtures from DIRAC + Bookkeeping",
    )
    parser.add_argument(
        "production", help="Production name (directory containing info.yaml)"
    )
    parser.add_argument("job", help="Job name within the production")
    parser.add_argument(
        "--ap-repo",
        default="data-pkg-repo",
        help="AP data-package repo name for manifest (default: data-pkg-repo)",
    )
    parser.add_argument(
        "--first-evt",
        type=int,
        default=0,
        help="First event to slice out (default: 0)",
    )

    parser.add_argument(
        "--evt-max",
        type=int,
        default=100,
        help="Number of events to keep in the sliced DST (default: 100)",
    )
    args = parser.parse_args()

    production_name = args.production
    job_name = args.job

    # ── Parse info.yaml ──────────────────────────────────────────────
    print(f"\n{'='*60}")
    print(f"Capturing fixture: {production_name}/{job_name}")
    print(f"{'='*60}\n")

    # Auto-discover info.yaml: check CWD first, then tests/data/*/
    info_yaml_path = Path(production_name) / "info.yaml"
    if not info_yaml_path.exists():
        # Search AP repos under tests/data/
        for repo_dir in sorted(DATA_DIR.iterdir()):
            candidate = repo_dir / production_name / "info.yaml"
            if candidate.exists():
                info_yaml_path = candidate
                args.ap_repo = repo_dir.name
                print(f"Found AP repo: {repo_dir.relative_to(TESTS_DIR)}")
                break
        else:
            print(
                f"ERROR: {production_name}/info.yaml not found in CWD or tests/data/*/.",
                file=sys.stderr,
            )
            sys.exit(1)

    from LbAPCommon import parse_yaml, render_yaml

    raw_yaml = info_yaml_path.read_text()
    rendered = render_yaml(raw_yaml)
    jobs_data = parse_yaml(rendered, production_name, None)

    if job_name not in jobs_data:
        available = ", ".join(jobs_data.keys())
        print(
            f"ERROR: Job '{job_name}' not found. Available: {available}",
            file=sys.stderr,
        )
        sys.exit(1)

    job_data = jobs_data[job_name]
    input_query = job_data["input"]
    turbo = job_data.get("turbo", False)

    # ── Step 1: query_dirac ──────────────────────────────────────────
    print("[1/6] Querying DIRAC for job metadata...")
    query_result = run_query_dirac(input_query, turbo)
    print(
        f"      Auto-conf: simulation={query_result['auto-conf-data']['simulation']}, "
        f"data_type={query_result['auto-conf-data']['data_type']}"
    )

    input_spec = query_result["input-spec"]
    conditions_dict = input_spec["conditions_dict"]
    conditions_description = input_spec["conditions_description"]
    event_type = input_spec["event_type"]

    # Save query_result.json
    job_fixture_dir = FIXTURES_DIR / production_name / job_name
    job_fixture_dir.mkdir(parents=True, exist_ok=True)
    query_result_path = job_fixture_dir / "query_result.json"
    query_result_path.write_text(json.dumps(query_result, indent=2) + "\n")
    print(f"      Saved: {query_result_path.relative_to(TESTS_DIR)}")

    # ── Step 2: BK query for LFN ────────────────────────────────────
    print("[2/6] Querying Bookkeeping for input LFN...")
    lfns = query_bk_for_lfns(
        conditions_dict, conditions_description, event_type, input_spec
    )
    if not lfns:
        print("ERROR: No LFNs found in Bookkeeping for this query.", file=sys.stderr)
        sys.exit(1)
    lfn = lfns[0]
    print(f"      LFN: {lfn}")

    # ── Step 3: resolve access URL ──────────────────────────────────
    print("[3/6] Resolving access URL...")
    access_url = resolve_access_url(lfn)
    print(f"      URL: {access_url}")

    # ── Step 4: generate slice config ───────────────────────────────
    # Derive output filename from LFN: keep original basename + _slice suffix
    lfn_basename = Path(lfn).stem  # e.g. "00070793_00000368_7.AllStreams"
    lfn_ext = Path(lfn).suffix  # e.g. ".dst"
    output_filename = f"{lfn_basename}_slice{lfn_ext}"

    TEST_DST_DIR.mkdir(parents=True, exist_ok=True)

    print(f"[4/6] Generating slice config (evt_max={args.evt_max})...")
    slice_config = generate_slice_config(
        access_url, output_filename, args.first_evt, args.evt_max
    )
    slice_yaml_path = TEST_DST_DIR / f"slice_{job_name}.yaml"
    with open(slice_yaml_path, "w") as f:
        yaml.dump(slice_config, f, default_flow_style=False, sort_keys=False)
    print(f"      Saved: {slice_yaml_path.relative_to(TESTS_DIR)}")

    # ── Step 5: run the slice ───────────────────────────────────────
    print(f"[5/6] Running DST slice ({output_filename})...")
    cwd_save = Path.cwd()
    try:
        # Run from the DST directory so the output file lands there
        import os

        os.chdir(TEST_DST_DIR)
        run_slice(slice_yaml_path)
    finally:
        os.chdir(cwd_save)

    dst_output_path = TEST_DST_DIR / output_filename
    if not dst_output_path.exists():
        print(f"ERROR: Sliced DST not created at {dst_output_path}", file=sys.stderr)
        sys.exit(1)
    print(
        f"      Output: {dst_output_path.relative_to(TESTS_DIR)} "
        f"({dst_output_path.stat().st_size / 1024:.0f} KB)"
    )

    # ── Step 6: update manifest ─────────────────────────────────────
    print("[6/6] Updating manifest.yaml...")
    dst_file_rel = output_filename
    fixture_dir = FIXTURES_DIR / production_name
    update_manifest(
        fixture_dir, production_name, args.ap_repo, job_name, dst_file_rel, lfn
    )
    print(f"      Updated: {(fixture_dir / 'manifest.yaml').relative_to(TESTS_DIR)}")

    # ── Done ────────────────────────────────────────────────────────
    print(f"\n{'='*60}")
    print("Fixture captured successfully!")
    print(f"{'='*60}")
    print("\nFiles created/updated:")
    print(f"  {query_result_path.relative_to(TESTS_DIR)}")
    print(f"  {dst_output_path.relative_to(TESTS_DIR)}")
    print(f"  {(fixture_dir / 'manifest.yaml').relative_to(TESTS_DIR)}")
    print(f"  {slice_yaml_path.relative_to(TESTS_DIR)}")
    print("\nTo commit:")
    print("  git add tests/fixtures/ tests/test-input-dst/")
    print(f"  git commit -m 'Add test fixture for {production_name}/{job_name}'")


if __name__ == "__main__":
    main()
