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.EntryPoint;
import click.kamil.springstatemachineexporter.analysis.model.LifecycleTriggerMarkers;
import click.kamil.springstatemachineexporter.analysis.model.LinkResolution;
import click.kamil.springstatemachineexporter.analysis.model.MatchedTransition;
import click.kamil.springstatemachineexporter.analysis.model.TriggerPoint;
import click.kamil.springstatemachineexporter.analysis.service.CodebaseIntelligenceProvider;
import click.kamil.springstatemachineexporter.analysis.service.JsonExportContextFactory;
import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import click.kamil.springstatemachineexporter.model.State;
import click.kamil.springstatemachineexporter.model.Transition;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import click.kamil.springstatemachineexporter.analysis.enricher.matching.EventMatchingEngine;
import click.kamil.springstatemachineexporter.analysis.enricher.matching.StrictFqnMatchingEngine;
import click.kamil.springstatemachineexporter.analysis.enricher.routing.BeanResolutionEngine;
import click.kamil.springstatemachineexporter.analysis.enricher.routing.HeuristicBeanResolutionEngine;
import click.kamil.springstatemachineexporter.analysis.resolver.BooleanConstraintEvaluator;
import click.kamil.springstatemachineexporter.analysis.resolver.MachineEnumCanonicalizer;
import click.kamil.springstatemachineexporter.analysis.resolver.StateMachineTypeResolver;

public class TransitionLinkerEnricher implements AnalysisEnricher {

    private final BeanResolutionEngine routingEngine = new HeuristicBeanResolutionEngine();

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

        EventMatchingEngine matchingEngine = new StrictFqnMatchingEngine(context);

        List<CallChain> updatedChains = new ArrayList<>();
        List<Transition> stateMachineTransitions = result.getTransitions();
        StateMachineTypeResolver.MachineTypes machineTypes =
                JsonExportContextFactory.resolveMachineTypes(
                        result.getName(),
                        result.getStateTypeFqn(),
                        result.getEventTypeFqn(),
                        context);

        for (CallChain chain : result.getMetadata().getCallChains()) {
            TriggerPoint tp = chain.getTriggerPoint();
            if (tp == null || tp.getEvent() == null) {
                updatedChains.add(chain);
                continue;
            }

            tp = MachineEnumCanonicalizer.ensureCallChainPolymorphicEvents(
                    tp, machineTypes, context, stateMachineTransitions, true);
            tp = CallChainLinkPolicy.applyRestExternalPolicy(tp, chain, context);

            if (LifecycleTriggerMarkers.isLifecycle(tp.getEvent())) {
                updatedChains.add(chain);
                continue;
            }

            String triggerSource = tp.getSourceState() != null ? simplifySourceState(tp.getSourceState()) : null;
            if (CallChainLinkPolicy.shouldFailClosedOnAmbiguousCallGraphWiden(tp)) {
                updatedChains.add(chain.toBuilder()
                        .triggerPoint(tp)
                        .matchedTransitions(null)
                        .linkResolution(LinkResolution.AMBIGUOUS_WIDEN)
                        .build());
                continue;
            }

            boolean ambiguousSource = false;

            if (triggerSource == null) {
                Map<String, Set<String>> sourcesByEvent = new LinkedHashMap<>();
                Map<String, String> canonicalSourceBySimplified = new LinkedHashMap<>();
                for (Transition t : stateMachineTransitions) {
                    if (t.getEvent() != null && matchingEngine.matches(t.getEvent(), tp)
                            && isRoutedToCorrectMachine(chain, result.getName(), context, stateMachineTransitions)
                            && isConstraintCompatible(tp.getConstraint(), result.getName())) {
                        String smEventForLink = canonicalEvent(t.getEvent());
                        for (State smSourceState : t.getSourceStates()) {
                            String smSourceForLink = canonicalState(smSourceState);
                            String smSource = simplifySourceState(smSourceForLink);
                            sourcesByEvent.computeIfAbsent(smEventForLink, ignored -> new LinkedHashSet<>())
                                    .add(smSource);
                            canonicalSourceBySimplified.putIfAbsent(smSource, smSourceForLink);
                        }
                    }
                }
                if (sourcesByEvent.values().stream().anyMatch(sources -> sources.size() > 1)) {
                    ambiguousSource = true;
                } else {
                    Set<String> allDistinctSources = new LinkedHashSet<>();
                    sourcesByEvent.values().forEach(allDistinctSources::addAll);
                    if (allDistinctSources.size() == 1) {
                        String inferredSimplifiedSource = allDistinctSources.iterator().next();
                        String inferredCanonicalSource = canonicalSourceBySimplified.get(inferredSimplifiedSource);
                        tp = tp.toBuilder().sourceState(inferredCanonicalSource).build();
                        triggerSource = inferredSimplifiedSource;
                    }
                }
            }

            List<MatchedTransition> matched = new ArrayList<>();

            if (!ambiguousSource) {
            for (Transition t : stateMachineTransitions) {
                if (t.getEvent() != null && matchingEngine.matches(t.getEvent(), tp)) {
                    String smEventForLink = canonicalEvent(t.getEvent());
                    for (State smSourceState : t.getSourceStates()) {
                            String smSourceForLink = canonicalState(smSourceState);
                            String smSource = simplifySourceState(smSourceForLink);
                            if (triggerSource == null || triggerSource.equals(smSource)) {
                                if (t.getTargetStates() == null || t.getTargetStates().isEmpty()) {
                                    MatchedTransition mt = MatchedTransition.builder()
                                            .sourceState(smSourceForLink)
                                            .targetState(smSourceForLink)
                                            .event(smEventForLink)
                                            .build();
                                    if (isRoutedToCorrectMachine(chain, result.getName(), context, stateMachineTransitions)
                                            && isConstraintCompatible(tp.getConstraint(), result.getName())) {
                                        matched.add(mt);
                                    }
                                } else {
                                    for (State smTargetState : t.getTargetStates()) {
                                        String targetForLink = canonicalState(smTargetState);
                                        MatchedTransition mt = MatchedTransition.builder()
                                                .sourceState(smSourceForLink)
                                                .targetState(targetForLink)
                                                .event(smEventForLink)
                                                .build();
                                        if (isRoutedToCorrectMachine(chain, result.getName(), context, stateMachineTransitions)
                                                && isConstraintCompatible(tp.getConstraint(), result.getName())) {
                                            matched.add(mt);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }

            if (ambiguousSource) {
                tp = tp.toBuilder().ambiguous(true).build();
            }

            LinkResolution linkResolution = CallChainLinkPolicy.resolveLinkResolution(tp, matched, ambiguousSource);

            if (!matched.isEmpty()) {
                CallChain newChain = chain.toBuilder()
                        .triggerPoint(tp)
                        .matchedTransitions(matched)
                        .linkResolution(linkResolution)
                        .build();
                updatedChains.add(newChain);
            } else {
                updatedChains.add(chain.toBuilder()
                        .triggerPoint(tp)
                        .matchedTransitions(null)
                        .linkResolution(linkResolution)
                        .build());
            }
        }

        List<TriggerPoint> scopedTriggers = MachineScopeFilter.filterTriggersForMachine(
                result.getMetadata().getTriggers(), result.getName(), context, stateMachineTransitions);
        List<CallChain> scopedChains = MachineScopeFilter.filterCallChainsForMachine(
                updatedChains, result.getName(), context, stateMachineTransitions);
        List<EntryPoint> scopedEntryPoints = EntryPointScopeResolver.scopeFromCallChains(
                scopedChains, result.getMetadata().getEntryPoints());

        CodebaseMetadata updatedMetadata = CodebaseMetadata.builder()
                .triggers(scopedTriggers)
                .entryPoints(scopedEntryPoints)
                .callChains(scopedChains)
                .properties(result.getMetadata().getProperties())
                .build();

        result.setMetadata(updatedMetadata);
    }

    
    private boolean isRoutedToCorrectMachine(
            CallChain chain, String currentMachineName, CodebaseContext context, List<Transition> machineTransitions) {
        return routingEngine.hasProvenMachineAffinity(chain, currentMachineName, context, machineTransitions);
    }

    private boolean isConstraintCompatible(String constraint, String machineName) {
        if (constraint == null || machineName == null) {
            return true;
        }
        return BooleanConstraintEvaluator.isCompatibleWithMachineDomain(
                constraint, MachineDomainKeys.extractMachineDomainKey(machineName));
    }

    private final java.util.Map<String, String> simplifySourceCache = new java.util.HashMap<>();

    private String simplifySourceState(String name) {
        if (name == null) return null;
        if (simplifySourceCache.containsKey(name)) return simplifySourceCache.get(name);

        // For source states we must preserve enum type context to avoid collisions between
        // different state enums that share a constant name (e.g. OrderState.NEW vs InvoiceState.NEW).
        // Use EnumFormat.fn-like suffix (Type.CONST) when possible.
        String simplified = name;
        int lastDot = name.lastIndexOf('.');
        if (lastDot > 0) {
            int prevDot = name.lastIndexOf('.', lastDot - 1);
            if (prevDot > 0) {
                simplified = name.substring(prevDot + 1);
            }
        }

        simplifySourceCache.put(name, simplified);
        return simplified;
    }

    private String canonicalState(State state) {
        if (state == null) {
            return null;
        }
        return state.fullIdentifier() != null ? state.fullIdentifier() : state.rawName();
    }

    private String canonicalEvent(click.kamil.springstatemachineexporter.model.Event event) {
        if (event == null) {
            return null;
        }
        return event.fullIdentifier() != null ? event.fullIdentifier() : event.rawName();
    }

}