package org.implab.gradle.variants.artifacts.internal;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Stream;

import org.eclipse.jdt.annotation.NonNullByDefault;
import org.gradle.api.Action;
import org.gradle.api.InvalidUserDataException;
import org.gradle.api.Task;
import org.gradle.api.file.ConfigurableFileCollection;
import org.gradle.api.file.Directory;
import org.gradle.api.file.DirectoryProperty;
import org.gradle.api.file.FileCollection;
import org.gradle.api.file.FileSystemLocation;
import org.gradle.api.model.ObjectFactory;
import org.gradle.api.provider.Provider;
import org.gradle.api.tasks.Sync;
import org.gradle.api.tasks.TaskContainer;
import org.gradle.api.tasks.TaskProvider;
import org.gradle.language.base.plugins.LifecycleBasePlugin;
import org.implab.gradle.common.core.lang.FilePaths;
import org.implab.gradle.common.core.lang.Strings;
import org.implab.gradle.variants.artifacts.ArtifactAssembly;
import org.implab.gradle.variants.artifacts.ArtifactAssemblySpec;
import org.implab.gradle.variants.artifacts.ArtifactSlot;
import org.implab.gradle.variants.artifacts.OutputSelectionSpec;
import org.implab.gradle.variants.core.Layer;
import org.implab.gradle.variants.core.Role;
import org.implab.gradle.variants.sources.CompileUnit;
import org.implab.gradle.variants.sources.CompileUnitsView;
import org.implab.gradle.variants.sources.RoleProjectionsView;
import org.implab.gradle.variants.sources.SourceSetMaterializer;

/**
 * Adapts slot contribution declarations to materialized
 * {@link ArtifactAssembly}
 * handles.
 *
 * <p>
 * Contribution-based assemblies create one {@link Sync} task per
 * {@link ArtifactSlot}. The task copies all collected slot inputs into a single
 * output directory. That output directory is then registered in
 * {@link ArtifactAssemblyRegistry} as the published artifact for the slot.
 *
 * <p>
 * Task-produced assemblies bypass the managed copy task. The producer task is
 * registered directly in {@link ArtifactAssemblyRegistry}, and its mapped output
 * file or directory becomes the published slot artifact.
 *
 * <p>
 * Input collection uses {@link SlotContributionVisitor}. Each contribution is
 * converted to a {@link SlotInputKey}; duplicate keys are ignored so that
 * repeated
 * topology-based selections do not add the same input twice.
 */
@NonNullByDefault
public class ArtifactAssemblyHandler {
    private final ObjectFactory objects;

    private final ArtifactAssemblyRegistry assemblyRegistry;

    private final DirectoryProperty assembliesDirectory;

    private final TaskContainer tasks;

    private final CompileUnitsView compileUnitsView;

    private final RoleProjectionsView roleProjectionsView;

    private final SourceSetMaterializer sourceSetMaterializer;

    private final Map<ArtifactSlot, SlotAssembly> slotInputs = new HashMap<>();

    private final Map<ArtifactSlot, AssemblyMode> assemblyModes = new HashMap<>();

    public ArtifactAssemblyHandler(
            ObjectFactory objects,
            TaskContainer tasks,
            ArtifactAssemblyRegistry assemblyRegistry,
            CompileUnitsView compileUnitsView,
            RoleProjectionsView roleProjectionsView,
            SourceSetMaterializer sourceSetMaterializer) {
        this.objects = objects;
        this.tasks = tasks;
        this.assemblyRegistry = assemblyRegistry;
        this.compileUnitsView = compileUnitsView;
        this.roleProjectionsView = roleProjectionsView;
        this.sourceSetMaterializer = sourceSetMaterializer;

        assembliesDirectory = objects.directoryProperty();
    }

    public DirectoryProperty getAssembliesDirectory() {
        return assembliesDirectory;
    }

    public void configureAssembly(ArtifactSlot artifactSlot, Action<? super ArtifactAssemblySpec> action) {
        var spec = new DefaultArtifactAssemblySpec(artifactSlot);
        action.execute(spec);
    }

    private void useAssemblyMode(ArtifactSlot artifactSlot, AssemblyMode mode) {
        var previous = assemblyModes.putIfAbsent(artifactSlot, mode);
        if (previous != null && previous != mode) {
            throw new InvalidUserDataException("Artifact slot '" + artifactSlot
                    + "' cannot mix task-produced artifact and contribution-based assembly");
        }
    }

    /**
     * Creates the assembly task for the given slot and registers its output
     * artifact.
     */
    private SlotAssembly createSlotAssembly(ArtifactSlot artifactSlot) {
        var assembly = new SlotAssembly();
        var fileCollection = assembly.inputs();

        var outputDirectory = outputDirectory(artifactSlot);

        var task = tasks.register(assembleTaskName(artifactSlot), Sync.class, copy -> {
            copy.setGroup(LifecycleBasePlugin.BUILD_GROUP);
            copy.into(outputDirectory);
            copy.from(fileCollection);
        });

        assemblyRegistry.register(artifactSlot, task, t -> outputDirectory);

        return assembly;
    }

    private String assembleTaskName(ArtifactSlot artifactSlot) {
        var variantName = artifactSlot.variant().getName();
        var slotName = artifactSlot.slot().getName();

        return "assembleVariantArtifactSlot"
                + "_v" + variantName.length() + "_" + variantName
                + "_s" + slotName.length() + "_" + slotName;
    }

    private Provider<Directory> outputDirectory(ArtifactSlot artifactSlot) {
        return assembliesDirectory.dir(
                FilePaths.cat(
                        artifactSlot.variant().getName(),
                        artifactSlot.slot().getName()));
    }

    /**
     * Collects slot contributions into one {@link ConfigurableFileCollection}.
     */
    private class ContributionVisitor implements SlotContributionVisitor {
        // artifact slot for this assembly
        private final ArtifactSlot artifactSlot;

        // seen inputs, used for deduplication
        private final SlotAssembly assembly;

        ContributionVisitor(ArtifactSlot artifactSlot, SlotAssembly assembly) {
            this.artifactSlot = artifactSlot;
            this.assembly = assembly;
        }

        @Override
        public void visit(DirectContribution contribution) {
            contribute(
                    SlotInputKey.newUniqueKey("Direct input for " + artifactSlot),
                    contribution.input());
        }

        @Override
        public void visit(VariantOutputsContribution contribution) {
            var units = compileUnitsView.getUnitsForVariant(artifactSlot.variant());
            contributeCompileUnits(units, contribution.outputs());
        }

        @Override
        public void visit(RoleOutputsContribution contribution) {
            var roleProjection = roleProjectionsView.requireProjection(artifactSlot.variant(),
                    contribution.role());
            var units = roleProjectionsView.getUnits(roleProjection);

            contributeCompileUnits(units, contribution.outputs());

        }

        @Override
        public void visit(LayerOutputsContribution contribution) {
            var unit = compileUnitsView.requireUnit(artifactSlot.variant(), contribution.layer());
            contributeCompileUnits(Set.of(unit), contribution.outputs());
        }

        private void contributeCompileUnits(Set<CompileUnit> units, Set<String> outputs) {
            units.stream()
                    // expand variant compile units, make (compileUnit, outputName) pairs
                    .flatMap(unit -> outputs.stream()
                            .map(output -> new CompileUnitOutputKey(unit, output)))
                    .forEach(key -> contribute(
                            key,
                            sourceSetMaterializer.getSourceSet(key.unit())
                                    .map(s -> s.output(key.outputName()))));
        }

        private void contribute(SlotInputKey key, Object resolvedInput) {
            assembly.addSlotInput(key, resolvedInput);
        }
    }

    /** Mutable input state for one slot assembly. */
    class SlotAssembly {
        private final ConfigurableFileCollection inputs = objects.fileCollection();
        private final Set<SlotInputKey> seen = new HashSet<>();

        public void addSlotInput(SlotInputKey key, Object input) {
            if (!seen.add(key))
                return;
            inputs.from(input);
        }

        public FileCollection inputs() {
            return inputs;
        }
    }

    private enum AssemblyMode {
        CONTRIBUTIONS,
        TASK_PRODUCER
    }

    /**
     * Default DSL facade for collecting {@link SlotContribution} declarations.
     *
     * <p>
     * The spec does not validate topology references immediately. It translates DSL
     * calls to contribution objects and passes them to the supplied consumer;
     * semantic
     * validation happens later when the assembly handler resolves contributions
     * against the finalized source model.
     */
    class DefaultArtifactAssemblySpec implements ArtifactAssemblySpec {

        private final ArtifactSlot artifactSlot;

        DefaultArtifactAssemblySpec(ArtifactSlot artifactSlot) {
            this.artifactSlot = artifactSlot;
        }

        @Override
        public void from(Object artifact) {
            contribute(new DirectContribution(artifact));
        }

        @Override
        public <T extends Task> void producedBy(
                TaskProvider<T> task,
                Function<? super T, ? extends Provider<? extends FileSystemLocation>> artifact) {
            registerProducedArtifact(task, artifact);
        }

        @Override
        public void fromVariant(Action<? super OutputSelectionSpec> action) {
            contribute(new VariantOutputsContribution(outputs(action)));
        }

        @Override
        public void fromRole(String roleName, Action<? super OutputSelectionSpec> action) {

            contribute(new RoleOutputsContribution(
                    objects.named(Role.class, roleName),
                    outputs(action)));
        }

        @Override
        public void fromLayer(String layerName, Action<? super OutputSelectionSpec> action) {
            contribute(new LayerOutputsContribution(
                    objects.named(Layer.class, layerName),
                    outputs(action)));
        }

        private static Set<String> outputs(Action<? super OutputSelectionSpec> action) {
            var spec = new OutputsSetSpec();
            action.execute(spec);
            return spec.outputs();
        }

        void contribute(SlotContribution contribution) {
            useAssemblyMode(artifactSlot, AssemblyMode.CONTRIBUTIONS);
            var assembly = slotInputs.computeIfAbsent(artifactSlot, ArtifactAssemblyHandler.this::createSlotAssembly);
            var contributionVisitor = new ContributionVisitor(artifactSlot, assembly);
            contribution.accept(contributionVisitor);
        }

        <T extends Task> void registerProducedArtifact(
                TaskProvider<T> task,
                Function<? super T, ? extends Provider<? extends FileSystemLocation>> artifact) {
            useAssemblyMode(artifactSlot, AssemblyMode.TASK_PRODUCER);
            assemblyRegistry.register(artifactSlot, task, artifact);
        }

    }

    /** Simple implementation of {@link OutputSelectionSpec}. */
    static class OutputsSetSpec implements OutputSelectionSpec {
        private final Set<String> outputs = new HashSet<>();

        @Override
        public void output(String name, String... extra) {
            Stream.concat(Stream.of(name), Stream.of(extra))
                    .map(Strings::requireNonBlank)
                    .forEach(outputs::add);
        }

        Set<String> outputs() {
            return Set.copyOf(outputs);
        }

    }

}
