# trivy_plugin/register.py
from pants.engine.internals.selectors import Get
from pants.engine.rules import collect_rules, goal_rule
from pants.engine.goal import Goal, GoalSubsystem
from pants.engine.process import Process, ProcessResult
from pants.engine.unions import UnionRule
from pants.backend.python.util_rules.pex import PexProcess, PexRequest
from pants.engine.fs import Digest, MergeDigests, Snapshot
from pants.backend.python.goals.package_pex_binary import PexBinaryFieldSet


class TrivyScanSubsystem(GoalSubsystem):
    name = "trivyscan"
    help = "Run Trivy security scan on PEX binaries."


class TrivyScanGoal(Goal):
    subsystem_cls = TrivyScanSubsystem


@goal_rule
async def run_trivy_scan(field_set: PexBinaryFieldSet) -> TrivyScanGoal:
    # Create a PEX with a loose layout
    pex_request = field_set.to_pex_request(layout="loose")
    pex_result = await Get(ProcessResult, PexProcess, pex_request)

    # Assuming trivy is installed and your path is correctly set,
    # and that trivy can access the directory structure as needed:
    trivy_process = Process(
        argv=["trivy", "rootfs", pex_result.output_directory],
        input_digest=pex_result.output_digest,
        description="Running Trivy scan.",
        output_files=["trivy_report.txt"],  # Assuming Trivy outputs to a report file
        output_directories=[pex_result.output_directory]
        # Include output directory to ensure all files are available
    )
    trivy_result = await Get(ProcessResult, Process, trivy_process)

    return TrivyScanGoal(exit_code=trivy_result.exit_code)


def rules():
    return [
        *collect_rules(),
    ]