package click.kamil.springstatemachineexporter.analysis.enricher;

import click.kamil.springstatemachineexporter.analysis.model.AnalysisResult;
import click.kamil.springstatemachineexporter.analysis.model.CallChain;
import click.kamil.springstatemachineexporter.analysis.model.CodebaseMetadata;
import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
import click.kamil.springstatemachineexporter.analysis.resolver.MachineEnumCanonicalizer;
import click.kamil.springstatemachineexporter.analysis.resolver.StateMachineTypeResolver;
import click.kamil.springstatemachineexporter.analysis.service.CodebaseIntelligenceProvider;
import click.kamil.springstatemachineexporter.analysis.service.JsonExportContextFactory;
import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.stream.Collectors;

/**
 * Canonicalizes trigger-side enum identifiers ({@code event}, {@code polymorphicEvents},
 * {@code sourceState}) using the owning machine config's {@code <State, Event>} types.
 */
@Slf4j
public class TriggerCanonicalizationEnricher implements AnalysisEnricher {

    @Override
    public void enrich(AnalysisResult result, CodebaseContext context, CodebaseIntelligenceProvider intelligence) {
        if (result.getMetadata() == null) {
            return;
        }

        StateMachineTypeResolver.MachineTypes machineTypes = JsonExportContextFactory.resolveMachineTypes(
                result.getName(), result.getStateTypeFqn(), result.getEventTypeFqn(), context);
        if (machineTypes.stateTypeFqn() == null && machineTypes.eventTypeFqn() == null) {
            return;
        }

        List<TriggerPoint> triggers = result.getMetadata().getTriggers();
        List<TriggerPoint> canonicalTriggers = triggers == null ? null : triggers.stream()
                .map(trigger -> MachineEnumCanonicalizer.canonicalizeAndExpandTriggerPoint(
                        trigger, machineTypes, context))
                .collect(Collectors.toList());

        List<CallChain> callChains = result.getMetadata().getCallChains();
        List<CallChain> canonicalChains = callChains == null ? null : callChains.stream()
                .map(chain -> {
                    if (chain.getTriggerPoint() == null) {
                        return chain;
                    }
                    TriggerPoint expanded = MachineEnumCanonicalizer.prepareCallChainTrigger(
                            chain.getTriggerPoint(),
                            machineTypes,
                            context,
                            result.getTransitions(),
                            chain.getPathBindings());
                    TriggerPoint canonical = expanded.getEventTypeFqn() == null && expanded.getStateTypeFqn() == null
                            ? MachineEnumCanonicalizer.canonicalizeTriggerLabelsForLinking(
                                    expanded, machineTypes, context)
                            : MachineEnumCanonicalizer.canonicalizeTriggerPoint(
                                    expanded, machineTypes, context);
                    return chain.toBuilder().triggerPoint(canonical).build();
                })
                .collect(Collectors.toList());

        if (canonicalTriggers == null && canonicalChains == null) {
            return;
        }

        log.debug("Canonicalized triggers for {}", result.getName());

        result.setMetadata(CodebaseMetadata.builder()
                .triggers(canonicalTriggers != null ? canonicalTriggers : triggers)
                .entryPoints(result.getMetadata().getEntryPoints())
                .callChains(canonicalChains != null ? canonicalChains : callChains)
                .properties(result.getMetadata().getProperties())
                .build());
    }
}
