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

from __future__ import annotations

import json
import logging
import re
from collections import defaultdict
from dataclasses import dataclass
from typing import DefaultDict

from pants.backend.shell.lint.shellcheck.subsystem import Shellcheck
from pants.backend.shell.shell_setup import ShellSetup
from pants.backend.shell.target_types import ShellSourceField
from pants.base.specs import AddressSpecs, DescendantAddresses
from pants.core.util_rules.external_tool import DownloadedExternalTool, ExternalToolRequest
from pants.engine.addresses import Address
from pants.engine.collection import DeduplicatedCollection
from pants.engine.fs import Digest, DigestSubset, MergeDigests, PathGlobs
from pants.engine.platform import Platform
from pants.engine.process import FallibleProcessResult, Process, ProcessCacheScope
from pants.engine.rules import Get, MultiGet, collect_rules, rule
from pants.engine.target import (
    Dependencies,
    DependenciesRequest,
    ExplicitlyProvidedDependencies,
    HydratedSources,
    HydrateSourcesRequest,
    InferDependenciesRequest,
    InferredDependencies,
    SourcesPaths,
    SourcesPathsRequest,
    Targets,
    WrappedTarget,
)
from pants.engine.unions import UnionRule
from pants.util.frozendict import FrozenDict
from pants.util.logging import LogLevel
from pants.util.ordered_set import OrderedSet

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class ShellMapping:
    """A mapping of Shell file names to their owning file address."""

    mapping: FrozenDict[str, Address]
    ambiguous_modules: FrozenDict[str, tuple[Address, ...]]


@rule(desc="Creating map of Shell file names to Shell targets", level=LogLevel.DEBUG)
async def map_shell_files() -> ShellMapping:
    all_targets = await Get(Targets, AddressSpecs([DescendantAddresses("")]))
    shell_tgts = tuple(tgt for tgt in all_targets if tgt.has_field(ShellSourceField))
    sources_per_target = await MultiGet(
        Get(SourcesPaths, SourcesPathsRequest(tgt[ShellSourceField])) for tgt in shell_tgts
    )

    files_to_addresses: dict[str, Address] = {}
    files_with_multiple_owners: DefaultDict[str, set[Address]] = defaultdict(set)
    for tgt, sources in zip(shell_tgts, sources_per_target):
        for f in sources.files:
            if f in files_to_addresses:
                files_with_multiple_owners[f].update({files_to_addresses[f], tgt.address})
            else:
                files_to_addresses[f] = tgt.address

    # Remove files with ambiguous owners.
    for ambiguous_f in files_with_multiple_owners:
        files_to_addresses.pop(ambiguous_f)

    return ShellMapping(
        mapping=FrozenDict(sorted(files_to_addresses.items())),
        ambiguous_modules=FrozenDict(
            (k, tuple(sorted(v))) for k, v in sorted(files_with_multiple_owners.items())
        ),
    )


class ParsedShellImports(DeduplicatedCollection):
    sort_input = True


@dataclass(frozen=True)
class ParseShellImportsRequest:
    # NB: We parse per-file, rather than per-target. This is necessary so that we can have each
    # file in complete isolation without its sibling files present so that Shellcheck errors when
    # trying to source a sibling file, which then allows us to extract that path.
    digest: Digest
    fp: str


PATH_FROM_SHELLCHECK_ERROR = re.compile(r"Not following: (.+) was not specified as input")


@rule
async def parse_shell_imports(
    request: ParseShellImportsRequest, shellcheck: Shellcheck