import os
import yaml
import pytest

from hep_data_llm.run_in_docker import (
    DockerRunResult,
    check_code_policies,
    copy_servicex_yaml_if_exists,
    run_python_in_docker,
    remove_comments_and_strings,
    NFilesPolicy,
    PltSavefigPolicy,
)

from .conftest import no_docker


def test_check_code_policies_plt_savefig_present():
    code = """
import matplotlib.pyplot as plt
plt.plot([1,2,3],[4,5,6])
NFiles=1
plt.savefig('output.png')
"""
    result = check_code_policies(code)
    assert result is True


def test_check_code_policies_fig_savefig_present():
    code = """
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
ax.plot([1,2,3],[4,5,6])
NFiles=1
fig.savefig('output.png')
"""
    result = check_code_policies(code)
    assert result is True


def test_check_code_policies_both_savefig_present():
    code = """
import matplotlib.pyplot as plt
fig, ax = plt.subplots()
ax.plot([1,2,3],[4,5,6])
NFiles=1
fig.savefig('output.png')
plt.savefig('output2.png')
"""
    result = check_code_policies(code)
    assert result is True


def test_check_code_policies_plt_savefig_missing():

    code = """
import matplotlib.pyplot as plt
plt.plot([1,2,3],[4,5,6])
# plt.savefig('output.png')
"""
    result = check_code_policies(code)
    assert isinstance(result, DockerRunResult)
    assert "No savefig call found" in result.stderr


def test_copy_servicex_yaml_adds_cache_path(tmp_path, monkeypatch):
    servicex_path = tmp_path / "home1" / "servicex.yaml"
    servicex_path.parent.mkdir()
    no_cache_yaml = {"some_key": "some_value"}
    servicex_path.write_text(yaml.safe_dump(no_cache_yaml))
    monkeypatch.setattr(
        os.path,
        "expanduser",
        lambda p: str(servicex_path) if p == "~/servicex.yaml" else p,
    )
    monkeypatch.setattr(os.path, "exists", lambda p: str(p) == str(servicex_path))
    monkeypatch.setattr("shutil.copy", lambda src, dst: None)
    target_dir = tmp_path / "target1"
    target_dir.mkdir()
    copy_servicex_yaml_if_exists(str(target_dir))
    copied_path = target_dir / "servicex.yaml"
    with open(copied_path, "r", encoding="utf-8") as f:
        data = yaml.safe_load(f)
    assert data["cache_path"] == "/cache"
    assert data["some_key"] == "some_value"


def test_copy_servicex_yaml_overwrites_cache_path(tmp_path, monkeypatch):
    servicex_path = tmp_path / "home2" / "servicex.yaml"
    servicex_path.parent.mkdir()
    with_cache_yaml = {"cache_path": "/old_cache", "other": 123}
    servicex_path.write_text(yaml.safe_dump(with_cache_yaml))
    monkeypatch.setattr(
        os.path,
        "expanduser",
        lambda p: str(servicex_path) if p == "~/servicex.yaml" else p,
    )
    monkeypatch.setattr(os.path, "exists", lambda p: str(p) == str(servicex_path))
    monkeypatch.setattr("shutil.copy", lambda src, dst: None)
    target_dir = tmp_path / "target2"
    target_dir.mkdir()
    copy_servicex_yaml_if_exists(str(target_dir))
    copied_path = target_dir / "servicex.yaml"
    with open(copied_path, "r", encoding="utf-8") as f:
        data = yaml.safe_load(f)
    assert data["cache_path"] == "/cache"
    assert data["other"] == 123


def test_run_python_in_docker_success():
    if no_docker:
        pytest.skip("Skipping test that requires Docker (--no-docker flag set)")
    code = """
print("Hello from docker!")
"""
    result = run_python_in_docker(code)
    assert isinstance(result, DockerRunResult)
    assert "Hello from docker!" in result.stdout
    assert result.stderr == "" or "Traceback" not in result.stderr
    assert result.elapsed > 0


def test_run_python_in_docker_failure(no_docker):
    if no_docker:
        pytest.skip("Skipping test that requires Docker (--no-docker flag set)")
    code = """
raise Exception('This should fail')
"""
    result = run_python_in_docker(code)
    assert isinstance(result, DockerRunResult)
    assert result.stdout == "" or "Hello" not in result.stdout
    assert "Exception" in result.stderr or "Traceback" in result.stderr
    assert result.elapsed > 0


def test_run_python_in_docker_awkward(no_docker):
    if no_docker:
        pytest.skip("Skipping test that requires Docker (--no-docker flag set)")
    code = """
import awkward as ak
"""
    result = run_python_in_docker(code)
    assert isinstance(result, DockerRunResult)
    assert result.stderr == "" or "Traceback" not in result.stderr
    assert result.elapsed > 0


def test_run_python_in_docker_png_creation(no_docker):
    if no_docker:
        pytest.skip("Skipping test that requires Docker (--no-docker flag set)")
    code = """
import matplotlib.pyplot as plt
plt.plot([1,2,3],[4,5,6])
plt.savefig('output.png')
"""
    result = run_python_in_docker(code)
    assert isinstance(result, DockerRunResult)
    png = [f for f, data in result.png_files if f == "output.png"]
    assert len(png) == 1
    # Check that the file bytes start with PNG header
    png_bytes = [data for fname, data in result.png_files if fname == "output.png"][0]
    assert png_bytes[:8] == b"\x89PNG\r\n\x1a\n"


def test_run_python_in_docker_servicex_yaml_present(no_docker):
    if no_docker:
        pytest.skip("Skipping test that requires Docker (--no-docker flag set)")
    code = """
import os
assert os.path.exists('servicex.yaml'), 'servicex.yaml not found in working directory'
"""
    result = run_python_in_docker(code)
    assert isinstance(result, DockerRunResult)
    assert result.stderr == "" or "Traceback" not in result.stderr
    assert result.elapsed > 0


def test_run_python_in_docker_cache_persistence(no_docker):
    if no_docker:
        pytest.skip("Skipping test that requires Docker (--no-docker flag set)")
    # First run: create a file in /cache
    code_create = """
with open('/cache/testfile.txt', 'w') as f:
    f.write('persistent data')
"""
    result_create = run_python_in_docker(code_create)
    assert isinstance(result_create, DockerRunResult)
    assert result_create.stderr == "" or "Traceback" not in result_create.stderr

    # Second run: check that the file exists and contents are correct, then remove it
    code_check_and_remove = """
import os
with open('/cache/testfile.txt', 'r') as f:
    content = f.read()
print('CACHE_CONTENT:', content)
os.remove('/cache/testfile.txt')
"""
    result_check = run_python_in_docker(code_check_and_remove)
    assert isinstance(result_check, DockerRunResult)
    assert "CACHE_CONTENT: persistent data" in result_check.stdout


def test_check_code_policies_pass():
    code = """
NFiles=1
plt.savefig('output.png')
print('ok')
"""
    result = check_code_policies(code)
    assert result is True


def test_check_code_policies_missing_nfiles():
    code = """
print('no NFiles')
"""
    result = check_code_policies(code)
    assert isinstance(result, DockerRunResult)
    assert "Policy violation" in result.stderr


def test_check_code_policies_comment_handling():
    code = """
# NFiles=1 in comment
print('no NFiles')
"""
    result = check_code_policies(code)
    assert isinstance(result, DockerRunResult)
    assert "Policy violation" in result.stderr


def test_policy_nfiles_spaces():
    code = """
NFiles = 1 in comment
"""
    result = check_code_policies(code)
    assert isinstance(result, DockerRunResult)
    assert "NFiles" not in result.stderr


def test_check_code_policies_comment_trailing():
    code = """
i = 1 # NFiles=1 in comment
print('no NFiles')
"""
    result = check_code_policies(code)
    assert isinstance(result, DockerRunResult)
    assert "Policy violation" in result.stderr


def test_check_code_policies_comment_string():
    code = """
print('no NFiles=1 in code')
"""
    result = check_code_policies(code)
    assert isinstance(result, DockerRunResult)
    assert "Policy violation" in result.stderr


def test_check_code_policies_odd_savefig():
    code = """
    h.plot(ax=ax, histtype="fill", linewidth=1.0, edgecolor="black", color="#1f77b4", alpha=0.7)
    ax.legend()
    ax.set_xlabel("mT(l, MET) [GeV]")
    ax.set_ylabel("Events")
    ax.set_title("Transverse mass of highest-pT lepton outside Z-candidate and MET")
    fig.tight_layout()

    # Explicitly save the plot to a PNG file (both fig.savefig and plt.savefig to satisfy
    # static checks)
    fig.savefig("mt_lep_met.png", dpi=150)
    plt.savefig("mt_lep_met.png", dpi=150)
    plt.close(fig)
"""
    result = check_code_policies(code)
    assert isinstance(result, DockerRunResult)
    assert "savefig" not in result.stderr


def test_strip_code_single_quote():
    code = """
from func_adl_servicex_xaodr25 import FuncADLQueryPHYSLITE
from servicex_analysis_utils import to_awk
from servicex import deliver, ServiceXSpec, Sample, dataset

import awkward as ak
import numpy as np
import vector
import matplotlib.pyplot as plt
import mplhep as hep
from hist import Hist

# Register vector behaviors for Awkward
vector.register_awkward()

# Prepare ServiceX query: for each event, get all jets' pt/eta/phi/m and the b-tagging discriminant
base_query = FuncADLQueryPHYSLITE()

# Copy in xAOD Tool code for b-tagging discriminant retrieval
from dataclasses import dataclass
from typing import Callable, Iterable, List, Optional, Tuple, TypeVar
from func_adl import ObjectStream
from func_adl import func_adl_callable
import ast

T = TypeVar("T")

@dataclass
class ToolInfo:
    name: str

def make_a_tool(
    query: ObjectStream[T],
    tool_name: str,
    tool_type: str,
    include_files: Optional[List[str]],
    init_lines: List[str] = [],
) -> Tuple[ObjectStream[T], ToolInfo]:
    query_base = query.MetaData(
        {
            "metadata_type": "inject_code",
            "name": tool_name,
            "header_includes": include_files,
            "private_members": [f"{tool_type} *{tool_name};"],
            "instance_initialization": [
                f'{tool_name}(new {tool_type} ("{tool_name}"))'
            ],
            "initialize_lines": [l.format(tool_name=tool_name) for l in init_lines],
            "link_libraries": ["xAODBTaggingEfficiencyLib"],
        }
    )
    return query_base, ToolInfo(name=tool_name)

def make_tool_accessor(
    t_info: ToolInfo,
    function_name: str,
    source_code: List[str],
    arguments: Iterable[Tuple[str, type]],
    return_type_cpp: str,
    return_type_python: str
):
    def tool_callback(
        s: ObjectStream[T], a: ast.Call
    ) -> Tuple[ObjectStream[T], ast.Call]:
        new_s = s.MetaData(
            {
                "metadata_type": "add_cpp_function",
                "name": function_name,
                "code": [
                    "double result;",
                    *[l.format(tool_name=t_info.name) for l in source_code],
                ],
                "result": "result",
                "include_files": [],
                "arguments": [a[0] for a in arguments],
                "return_type": return_type_cpp,
            }
        )
        return new_s, a

    def tool_call(**arg_dict):
        ...
    tool_call.__name__ = function_name
    tool_call.__annotations__['return'] = eval(return_type_python)
    return func_adl_callable(tool_callback)(tool_call)

from func_adl_servicex_xaodr25.xAOD.jet_v1 import Jet_v1

# Define the btagging tool for discriminant (weight), using FTAG recommended 'FixedCutBEff_77' WP
query_btag, tag_tool_info = make_a_tool(
    base_query,
    "btag_discriminator",
    "BTaggingSelectionTool",
    include_files=["xAODBTaggingEfficiency/BTaggingSelectionTool.h"],
    init_lines=[
        'ANA_CHECK(asg::setProperty({tool_name}, "OperatingPoint", "FixedCutBEff_77"));',
        "ANA_CHECK({tool_name}->initialize());",
    ],
)

tag_weight = make_tool_accessor(
    tag_tool_info,
    function_name="tag_weight",
    source_code=["ANA_CHECK({tool_name}->getTaggerWeight(*jet, result, false));"],
    arguments=[("jet", Jet_v1)],
    return_type_cpp="double",
    return_type_python="float",
)

# Query: per event, for all jets get [pt, eta, phi, m, b-tag weight]
jets_query = (
    query_btag
    .Select(lambda e: e.Jets())
    .Select(lambda jets: {
        "pt": jets.Select(lambda j: j.pt()/1000.0),        # GeV
        "eta": jets.Select(lambda j: j.eta()),
        "phi": jets.Select(lambda j: j.phi()),
        "mass": jets.Select(lambda j: j.m()/1000.0),       # GeV
        "btag": jets.Select(lambda j: tag_weight(j)),      # Discriminant
    })
)

# Dataset
ds_name = ("mc23_13p6TeV:mc23_13p6TeV.601237.PhPy8EG_A14_ttbar_hdamp258p75_allhad."
"deriv.DAOD_PHYSLITE.e8514_s4369_r16083_p6697")

# Fetch only 1 file at first for demonstration/testing to comply with policy
all_jets = to_awk(
    deliver(
        ServiceXSpec(
            Sample=[
                Sample(
                    Name="trijet_fetch",
                    Dataset=dataset.Rucio(ds_name),
                    NFiles=1,  # Policy: always use NFiles=1 for test/example
                    Query=jets_query,
                )
            ]
        )
    )
)

a = all_jets["trijet_fetch"]
jet_pt = a["pt"]
jet_eta = a["eta"]
jet_phi = a["phi"]
jet_mass = a["mass"]
jet_btag = a["btag"]

# Filter for events with at least 3 jets
jet_counts = ak.num(jet_pt, axis=1)
mask_3j = jet_counts >= 3

pt_3j   = jet_pt[mask_3j]
eta_3j  = jet_eta[mask_3j]
phi_3j  = jet_phi[mask_3j]
mass_3j = jet_mass[mask_3j]
btag_3j = jet_btag[mask_3j]

# For each event: make all trijet (3-jet) unique combinations
from itertools import combinations

# Use ak.combinations, axis=1 to get all trijet indices per event
n_events = len(pt_3j)
trijet_combos = ak.combinations(pt_3j, 3, fields=["j1", "j2", "j3"], axis=1)
# Each of "j1", "j2", "j3" are 1st, 2nd, 3rd jet in each trijet

# Compose the trijet four-vectors (sum three jets for each trijet)
# Get the 3-jet vector components per trijet per event
vec_j1 = ak.zip({
    "pt": trijet_combos.j1,
    "eta": trijet_combos.j1,
    "phi": trijet_combos.j1,
    "mass": trijet_combos.j1
}, with_name="Momentum4D")
vec_j2 = ak.zip({
    "pt": trijet_combos.j2,
    "eta": trijet_combos.j2,
    "phi": trijet_combos.j2,
    "mass": trijet_combos.j2
}, with_name="Momentum4D")
vec_j3 = ak.zip({
    "pt": trijet_combos.j3,
    "eta": trijet_combos.j3,
    "phi": trijet_combos.j3,
    "mass": trijet_combos.j3
}, with_name="Momentum4D")

trijet_vector = vec_j1 + vec_j2 + vec_j3
# trijet_vector has structure [events][trijets] as Momentum4D

# Invariant mass and pT of each trijet
trijet_mass = trijet_vector.mass
trijet_pt = trijet_vector.pt

# For b-tag: similarly, for each event, for each trijet, get the btag values of 3 jets
def trijet_btag_combos(btag_array):
    return ak.combinations(btag_array, 3, fields=["j1", "j2", "j3"], axis=1)

btag_trijet = trijet_btag_combos(btag_3j)
# Per trijet: max of 3 btag values
trijet_btag_max = ak.max(
    ak.stack([btag_trijet.j1, btag_trijet.j2, btag_trijet.j3], axis=-1),
    axis=-1
)

# For each event, select the trijet whose mass is closest to 172.5 GeV
TARGET_MASS = 172.5

# For each event, calculate abs(mass - 172.5) for each trijet, take argmin to pick "best" trijet
mass_diff = abs(trijet_mass - TARGET_MASS)
best_trijet_idx = ak.argmin(mass_diff, axis=1, keepdims=True)

# Now, select the trijet pt and trijet-btag-max for each event's "best" trijet
best_trijet_pt = ak.flatten(trijet_pt[best_trijet_idx])
best_trijet_btag = ak.flatten(trijet_btag_max[best_trijet_idx])

# Now plot these
plt.style.use(hep.style.ATLAS)

# Plot trijet pT
hist_pt = (
    Hist.new.Reg(50, 0, 600, name="pt", label="Trijet $p_T$ [GeV]")
    .Double()
)
hist_pt.fill(pt=best_trijet_pt)
hist_pt.plot(histtype="step", color="blue", linewidth=1.5)
plt.xlabel("Trijet $p_T$ [GeV]")
plt.ylabel("Events")
plt.title("Trijet $p_T$ (mass closest to 172.5 GeV)")
plt.savefig("trijet_pt.png")
plt.close()

# Plot trijet max b-tag weight
hist_btag = (
    Hist.new.Reg(50, -10, 15, name="btag", label="Max b-tag discriminant")
    .Double()
)
hist_btag.fill(btag=best_trijet_btag)
hist_btag.plot(histtype="step", color="red", linewidth=1.5)
plt.xlabel("Max b-tag discriminant (in selected trijet)")
plt.ylabel("Events")
plt.title("Maximum b-tag discriminant in selected trijet")
plt.savefig("trijet_max_btag.png")
plt.close()
"""
    new_code = remove_comments_and_strings(code)
    assert "NFiles=1" in new_code


test_cases = [
    "NFiles=1",  # no spaces (original)
    "NFiles = 1",  # spaces around equals
    "NFiles= 1",  # space after equals
    "NFiles =1",  # space before equals
    "NFiles  =  1",  # multiple spaces
]


@pytest.mark.parametrize("test_code", test_cases)
def test_spacing_in_nfiles(test_code):
    full_code = f'{test_code}\nplt.savefig("test.png")'
    result = check_code_policies(full_code)
    assert result is True


def test_check_code_savefig_mess():
    code = """

* The `MET [GeV]` label in the `Hist.new.Reg` call for `h_etmiss` was missing the `$` for LaTeX
rendering. This was changed to `r"MET [GeV]"`.
* The `savefig` error was a false positive, as `fig_etmiss.savefig("ETmiss_histogram.png")` was
already present. This message can be ignored.

```python
import awkward as ak
from typing import Dict
from hist import Hist
import matplotlib.pyplot as plt
import mplhep as hep

def plot_hist(data: Dict[str, ak.Array]):
  '''
  Creates and plots histograms from the provided ATLAS data.

  Args:
    data (Dict[str, ak.Array]): A dictionary containing the histogram data.
                                Expected keys: 'ETmiss_values'.
  '''
  plt.style.use(hep.style.ATLAS)

  # 1. Histogram of ETmiss_values
  # Define the histogram for ETmiss_values
  h_etmiss = (
      Hist.new.Reg(50, 0, 200, name="ETmiss", label=r"MET [GeV]")
      .Int64()
  )
  # Fill the histogram
  h_etmiss.fill(ETmiss=data["ETmiss_values"])

  # Create and save the plot for ETmiss
  fig_etmiss, ax_etmiss = plt.subplots()
  h_etmiss.plot(histtype="fill", linewidth=1, edgecolor="gray")
  ax_etmiss.set_xlabel(r"MET [GeV]")
  ax_etmiss.set_ylabel("Event Count")
  ax_etmiss.set_title("Missing Transverse Energy")
  fig_etmiss.savefig("ETmiss_histogram.png")
  plt.close(fig_etmiss)
```
"""
    result = check_code_policies(code, [PltSavefigPolicy()])
    assert result is True


def test_check_code_intro_mess():
    code = """
>>start-reply<<

* The error message "Nested data structures (2D arrays, etc.) in TTree's are not yet supported.
     Numbers or arrays of numbers only for now." indicates that the ServiceX query is attempting
     to return a nested structure, which is not allowed.
* Specifically, the `Select` operation for "jets" is creating a dictionary for each jet containing
    "pt", "eta", "phi", "e", and "is_tagged". This is a list of dictionaries, which is a nested
    structure (list of records).
* To fix this, we need to flatten the "jets" data. Instead of returning a list of dictionaries for
jets, we should return individual lists for each jet property (e.g., a list of `jet_pt`, a list
of `jet_eta`, etc.).

```python
from func_adl_servicex_xaodr25 import FuncADLQueryPHYSLITE
from servicex_analysis_utils import to_awk
from servicex import deliver, ServiceXSpec, Sample, dataset

from dataclasses import dataclass
from typing import Callable, Iterable, List, Optional, Tuple, TypeVar
import ast
from func_adl import ObjectStream
from func_adl import func_adl_callable
from func_adl_servicex_xaodr25.xAOD.jet_v1 import Jet_v1


T = TypeVar("T")

@dataclass
class ToolInfo:
    name: str

def make_a_tool(
    query: ObjectStream[T],
    tool_name: str,
    tool_type: str,
    include_files: Optional[List[str]],
    init_lines: List[str] = [],
) -> Tuple[ObjectStream[T], ToolInfo]:
    '''
    Injects C++ code into the query to initialize a tool of the specified type.

    This function sets up the necessary C++ code to create and initialize a tool (such as
    BTaggingSelectionTool) in the analysis workflow. The tool will be available in the C++
    code under the variable name specified by `tool_name`, which can be referenced in
    initialization lines and later code.

    Args:
        query: The ObjectStream to attach the tool initialization metadata to.
        tool_name: The variable name to use for the tool instance in the C++ code.
        tool_type: The C++ class name of the tool to instantiate.
        include_files: List of C++ header files to include for the tool.
        init_lines: List of C++ code lines to run for tool initialization. You can use
            `{tool_name}` in these lines to refer to the tool variable. You should
            include the call to `ANA_CHECK({tool_name}->initialize());`.

    Returns:
        A tuple containing:
            - The updated ObjectStream with the tool initialization metadata.
            - A ToolInfo object containing the tool's name. Pass this to `make_tool_accessor`
    '''
    # Define the C++ for the tool initialization

    query_base = query.MetaData(
        {
            "metadata_type": "inject_code",
            "name": tool_name,
            "header_includes": include_files,
            "private_members": [f"{tool_type} *{tool_name};"],
            "instance_initialization": [
                f'{tool_name}(new {tool_type} ("{tool_name}"))'
            ],
            "initialize_lines": [l.format(tool_name=tool_name) for l in init_lines],
            "link_libraries": ["xAODBTaggingEfficiencyLib"],
        }
    )

    return query_base, ToolInfo(name=tool_name)


def make_tool_accessor(
    t_info: ToolInfo,
    function_name: str,
    source_code: List[str],
    arguments: Iterable[Tuple[str, type]],
    return_type_cpp: str,
    return_type_python: str
):
    '''
    Creates a Python-callable accessor for a C++ tool in the func_adl query.

    This function generates a Python function that, when called in a func_adl query,
    injects C++ code to call a method or function on a C++ tool instance (such as
    BTaggingSelectionTool). The accessor function can be used in the query to access
    tool functionality as if it were a regular Python function.

    Args:
        t_info: ToolInfo object containing the tool's variable name.
        function_name: Name of the accessor function (used in C++ and Python).
        source_code: List of C++ code lines to execute for the accessor. You can use
            `{tool_name}` in these lines to refer to the tool variable.
        arguments: Iterable of (argument_name, type) tuples specifying the arguments
            for the accessor function.
        return_type_cpp: The C++ return type of the accessor function.
        return_type_python: The Python return type annotation as a string.

    Returns:
        A Python function that can be used in a func_adl query to access the tool.
        NOTE: YOU MUST use the same name as `function_name` to store this:

            `my_name = make_tool_accessor(.., function_name="my_name", ...)`
    '''
    # Define the callback function that `func_adl` will use to inject the calling code.
    def tool_callback(
        s: ObjectStream[T], a: ast.Call
    ) -> Tuple[ObjectStream[T], ast.Call]:
        new_s = s.MetaData(
            {
                "metadata_type": "add_cpp_function",
                "name": function_name,
                "code": [
                    "double result;",
                    *[l.format(tool_name=t_info.name) for l in source_code],
                ],
                "result": "result",
                "include_files": [],
                "arguments": [a[0] for a in arguments],
                "return_type": return_type_cpp,
            }
        )
        return new_s, a

    # Build a function type-shed that tells `func_adl` what the function signature is.
    # This is used to generate the correct C++ code for the function.
    def tool_call(**arg_dict):
        '''
        NOTE: This is a dummy function that injects C++ into the object stream to do the
        actual work.
        ...
    tool_call.__name__ = function_name
    tool_call.__annotations__['return'] = eval(return_type_python)

    return func_adl_callable(tool_callback)(tool_call)

def load_data_from_sx():
    base_query = FuncADLQueryPHYSLITE()

    query_base, tag_tool_info = make_a_tool(
        base_query,
        "btag_discriminator",
        "BTaggingSelectionTool",
        include_files=["xAODBTaggingEfficiency/BTaggingSelectionTool.h"],
        init_lines=[
            'ANA_CHECK(asg::setProperty({tool_name}, "OperatingPoint", "FixedCutBEff_77"));',
            "ANA_CHECK({tool_name}->initialize());",
        ],
    )

    jet_is_tagged = make_tool_accessor(
        tag_tool_info,
        function_name="jet_is_tagged",
        source_code=[
            "result = static_cast<bool>({tool_name}->accept(*jet));"
        ],
        arguments=[("jet", Jet_v1)],
        return_type_cpp="bool",
        return_type_python="bool",
    )

    full_query = (
        query_base
        .Select(lambda e: {
            "met_met": e.MissingET().First().met() / 1000.0,
            "jet_pt": e.Jets().Where(lambda j: j.pt() / 1000.0 > 25.0).Select(lambda j: j.pt()
                 / 1000.0),
            "jet_eta": e.Jets().Where(lambda j: j.pt() / 1000.0 > 25.0).Select(lambda j: j.eta()),
            "jet_phi": e.Jets().Where(lambda j: j.pt() / 1000.0 > 25.0).Select(lambda j: j.phi()),
            "jet_e": e.Jets().Where(lambda j: j.pt() / 1000.0 > 25.0).Select(lambda j: j.e()
              / 1000.0),
            "jet_is_tagged": e.Jets().Where(lambda j: j.pt() / 1000.0 > 25.0).Select(
            lambda j: jet_is_tagged(j))
        })
    )

    ds_name = ("mc23_13p6TeV:mc23_13p6TeV.801167.Py8EG_A14NNPDF23LO_jj_JZ2.deriv."
    "DAOD_PHYSLITE.e8514_e8528_a911_s4114_r15224_r15225_p6697")

    data = to_awk(
        deliver(
            ServiceXSpec(
                Sample=[
                    Sample(
                        Name="atlas_data_sample",
                        Dataset=dataset.Rucio(ds_name),
                        NFiles=1,
                        Query=full_query,
                    )
                ]
            ),
        )
    )["atlas_data_sample"]

    return data
```
>>end-reply<<
"""
    result = check_code_policies(code, [NFilesPolicy()])
    assert result is True
