from typing import Optional

from pants.backend.docker.target_types import DockerImageRepositoryField, DockerImageTarget
from pants.backend.python.goals.package_pex_binary import package_pex_binary, PexBinaryFieldSet, \
    built_pacakge_for_pex_from_targets_request
from pants.backend.python.target_types import PexBinary, PexBinaryDefaults
from pants.core.goals.package import OutputPathField
from pants.engine.console import Console
from pants.engine.fs import PathGlobs
from pants.engine.goal import Goal, GoalSubsystem
from pants.engine.internals.native_engine import Digest
from pants.engine.internals.selectors import Get
from pants.engine.process import ProcessResult, Process, FallibleProcessResult
from pants.engine.rules import goal_rule, collect_rules, rule
from pants.engine.target import Targets, TransitiveTargets, TransitiveTargetsRequest


@rule
async def get_docker_image_repository(target: PexBinary) -> str:
    transitive_targets = await Get[TransitiveTargets](TransitiveTargetsRequest([target.address]))
    for tgt in transitive_targets.dependencies:
        if isinstance(tgt, DockerImageTarget):
            return tgt.get(DockerImageRepositoryField).value
    return None


@rule
async def run_actual_scan(target: PexBinary) -> ProcessResult:
    digest = await Get(Digest, PathGlobs([target[OutputPathField].value]))
    repository = await Get[Optional[str]](get_docker_image_repository, target)
    if not repository:
        return ProcessResult(0, "No repository found for target {target.address}.")

    output_file = f"{repository}_trivy_report.txt"

    trivy_command = ["trivy", "rootfs", "--output", output_file, "/tmp/rootfs"]

    process = Process(
        argv=trivy_command,
        input_digest=digest,
        description=f"Running Trivy scan on {repository} PEX directory",
        output_files=(output_file,),
    )

    result = await Get[FallibleProcessResult](Process, process)
    return result


@rule
async def run_trivy_scan(targets: Targets) -> ProcessResult:
    pex_targets = [target for target in targets if isinstance(target, PexBinary)]
    print(pex_targets)
    if not pex_targets:
        return ProcessResult(0, "No PEX targets to scan.")

    for target in pex_targets:
        result = await run_actual_scan(target)
        if result.exit_code != 0:
            return result

    return result


class HelloWorldSubsystem(GoalSubsystem):
    name = "hello-world"
    help = "An example goal."


class HelloWorld(Goal):
    subsystem_cls = HelloWorldSubsystem
    environment_behavior = Goal.EnvironmentBehavior.LOCAL_ONLY


@goal_rule
async def hello_world(targets: Targets, console: Console) -> HelloWorld:
    console.print_stdout("Hello!")
    result = await run_trivy_scan(targets)
    return HelloWorld(exit_code=result.exit_code)


def rules():
    return collect_rules()