# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

import itertools
import logging
import shlex
from itertools import chain
from pathlib import Path

from pants.core.target_types import ResourcesFieldSet, ResourcesGeneratorFieldSet
from pants.core.util_rules import stripped_source_files
from pants.core.util_rules.source_files import SourceFilesRequest
from pants.core.util_rules.stripped_source_files import strip_source_roots
from pants.core.util_rules.system_binaries import BashBinary, TouchBinary, ZipBinary
from pants.engine.fs import MergeDigests
from pants.engine.internals.selectors import concurrently
from pants.engine.intrinsics import merge_digests
from pants.engine.process import Process, execute_process_or_raise
from pants.engine.rules import collect_rules, implicitly, rule
from pants.engine.target import SourcesField
from pants.engine.unions import UnionRule
from pants.jvm import compile
from pants.jvm.compile import (
    ClasspathDependenciesRequest,
    ClasspathEntry,
    ClasspathEntryRequest,
    ClasspathEntryRequests,
    CompileResult,
    FallibleClasspathEntries,
    FallibleClasspathEntry,
    compile_classpath_entries,
)
from pants.jvm.strip_jar.strip_jar import StripJarRequest, strip_jar
from pants.jvm.subsystems import JvmSubsystem
from pants.util.logging import LogLevel

logger = logging.getLogger(__name__)


class JvmResourcesRequest(ClasspathEntryRequest):
    field_sets = (
        ResourcesFieldSet,
        ResourcesGeneratorFieldSet,
    )


@rule(desc="Assemble resources")
async def assemble_resources_jar(
    zip: ZipBinary,
    bash: BashBinary,
    touch: TouchBinary,
    jvm: JvmSubsystem,
    request: JvmResourcesRequest,
) -> FallibleClasspathEntry:
    # Request the component's direct dependency classpath, and additionally any prerequisite.
    # Filter out any dependencies that are generated by our current target so that each resource
    # only appears in a single input JAR.
    # NOTE: Generated dependencies will have the same dependencies as the current target, so we
    # don't need to inspect those dependencies.
    optional_prereq_request = [*((request.prerequisite,) if request.prerequisite else ())]
    fallibles = await concurrently(
        compile_classpath_entries(ClasspathEntryRequests(optional_prereq_request)),
        compile_classpath_entries(
            **implicitly(ClasspathDependenciesRequest(request, ignore_generated=True))
        ),
    )
    direct_dependency_classpath_entries = FallibleClasspathEntries(
        itertools.chain(*fallibles)
    ).if_all_succeeded()

    if direct_dependency_classpath_entries is None:
        return FallibleClasspathEntry(
            description=str(request.component),
            result=CompileResult.DEPENDENCY_FAILED,
            output=None,
            exit_code=1,
        )

    source_files = await strip_source_roots(
        **implicitly(
            SourceFilesRequest([tgt.get(SourcesField) for tgt in request.component.members])
        )
    )

    output_filename = f"{request.component.representative.address.path_safe_spec}.resources.jar"
    output_files = [output_filename]

    # #16231: Valid JAR files need the directories of each resource file as well as the files
    # themselves.

    paths = {Path(filename) for filename in source_files.snapshot.files}
    directories = {parent for path in paths for parent in path.parents}
    input_files = {str(path) for path in chain(paths, directories)}

    resources_jar_input_digest = source_files.snapshot.digest

    input_filenames = shlex.join(sorted(input_files))

    resources_jar_result = await execute_process_or_raise(
        **implicitly(
            Process(
                argv=[
                    bash.path,
                    "-c",
                    " ".join(
                        [
                            "TZ=UTC",
                            touch.path,
                            "-t 198001010000.00",
                            input_filenames,
                            "&&",
                            "TZ=UTC",
                            zip.path,
                            "-oX",
                            output_filename,
                            input_filenames,
                        ]
                    ),
                ],
                description=f"Build resources JAR for {request.component}",
                input_digest=resources_jar_input_digest,
                output_files=output_files,
                level=LogLevel.DEBUG,
            )
        )
    )

    output_digest = resources_jar_result.output_digest
    if jvm.reproducible_jars:
        output_digest = await strip_jar(
            **implicitly(StripJarRequest(output_digest, tuple(output_files)))
        )

    cpe = ClasspathEntry(output_digest, output_files, [])

    merged_cpe_digest = await merge_digests(
        MergeDigests(chain((cpe.digest,), (i.digest for i in direct_dependency_classpath_entries)))
    )

    merged_cpe = ClasspathEntry.merge(
        digest=merged_cpe_digest, entries=[cpe, *direct_dependency_classpath_entries]
    )

    return FallibleClasspathEntry(output_filename, CompileResult.SUCCEEDED, merged_cpe, 0)


def rules():
    return [
        *collect_rules(),
        *compile.rules(),
        *stripped_source_files.rules(),
        UnionRule(ClasspathEntryRequest, JvmResourcesRequest),
    ]
