package click.kamil.springstatemachineexporter.ast.app;

import click.kamil.springstatemachineexporter.model.State;
import click.kamil.springstatemachineexporter.model.Transition;
import click.kamil.springstatemachineexporter.ast.common.CodebaseContext;
import org.eclipse.jdt.core.dom.TypeDeclaration;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.List;

import static org.assertj.core.api.Assertions.assertThat;

class StateMachineAggregatorTest {

    @TempDir
    Path tempDir;

    private CodebaseContext context;
    private StateMachineAggregator aggregator;

    @BeforeEach
    void setUp() {
        context = new CodebaseContext();
        aggregator = new StateMachineAggregator(context);
    }

    @Test
    void shouldAggregateTransitionsFromInheritance() throws IOException {
        String baseSource = """
                package com.example;
                public abstract class BaseConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        transitions.withExternal().source("BASE_S1").target("BASE_S2").event("BASE_E1");
                    }
                }
                """;
        String childSource = """
                package com.example;
                public class ChildConfig extends BaseConfig {
                    @Override
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        super.configure(transitions);
                        transitions.withExternal().source("CHILD_S1").target("CHILD_S2").event("CHILD_E1");
                    }
                }
                """;

        writeClass("BaseConfig", baseSource);
        writeClass("ChildConfig", childSource);
        context.scan(tempDir);

        TypeDeclaration childTd = context.getTypeDeclaration("ChildConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(childTd);

        assertThat(transitions).hasSize(2);
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("BASE_E1"));
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("CHILD_E1"));
    }

    @Test
    void shouldFollowMethodCallsAcrossHierarchy() throws IOException {
        String baseSource = """
                package com.example;
                public abstract class BaseConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        addBaseTransitions(transitions);
                    }
                    protected void addBaseTransitions(StateMachineTransitionConfigurer transitions) {
                        transitions.withExternal().source("BASE_S1").target("BASE_S2").event("BASE_E1");
                    }
                }
                """;
        String childSource = """
                package com.example;
                public class ChildConfig extends BaseConfig {
                    @Override
                    protected void addBaseTransitions(StateMachineTransitionConfigurer transitions) {
                        transitions.withExternal().source("OVERRIDDEN_S1").target("OVERRIDDEN_S2").event("OVERRIDDEN_E1");
                    }
                }
                """;

        writeClass("BaseConfig", baseSource);
        writeClass("ChildConfig", childSource);
        context.scan(tempDir);

        TypeDeclaration childTd = context.getTypeDeclaration("ChildConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(childTd);

        assertThat(transitions).hasSize(1);
        assertThat(transitions.get(0).getEvent().toString()).isEqualTo("OVERRIDDEN_E1");
    }

    @Test
    void shouldResolveParametersInMethodCalls() throws IOException {
        String source = """
                package com.example;
                public class ParamConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        addCustomTransition(transitions, "START", "END", "GO");
                    }
                    private void addCustomTransition(StateMachineTransitionConfigurer t, String src, String tgt, String evt) {
                        t.withExternal().source(src).target(tgt).event(evt);
                    }
                }
                """;

        writeClass("ParamConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("ParamConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(td);

        assertThat(transitions).hasSize(1);
        assertThat(transitions.get(0).getSourceStates()).extracting(State::toString).containsExactly("START");
        assertThat(transitions.get(0).getTargetStates()).extracting(State::toString).containsExactly("END");
        assertThat(transitions.get(0).getEvent().toString()).isEqualTo("GO");
    }

    @Test
    void shouldParseStreamOfForEachTransitions() throws IOException {
        String source = """
                package com.example;
                import java.util.stream.Stream;
                public class StreamConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        Stream.of("E1", "E2").forEach(e -> {
                            transitions.withExternal().source("S1").target("S2").event(e);
                        });
                    }
                }
                """;

        writeClass("StreamConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("StreamConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(td);

        assertThat(transitions).hasSize(2);
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("E1"));
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("E2"));
    }

    @Test
    void shouldAggregateStatesAcrossInheritance() throws IOException {
        String baseSource = """
                package com.example;
                public abstract class BaseConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineStateConfigurer states) throws Exception {
                        states.withStates().initial("BASE_START");
                    }
                }
                """;
        String childSource = """
                package com.example;
                public class ChildConfig extends BaseConfig {
                    public void configure(StateMachineStateConfigurer states) throws Exception {
                        super.configure(states);
                        states.withStates().state("CHILD_STATE").end("CHILD_END");
                    }
                }
                """;

        writeClass("BaseConfig", baseSource);
        writeClass("ChildConfig", childSource);
        context.scan(tempDir);

        TypeDeclaration childTd = context.getTypeDeclaration("ChildConfig");
        aggregator.aggregateStates(childTd);

        assertThat(aggregator.getInitialStates()).containsExactly("BASE_START");
        assertThat(aggregator.getEndStates()).containsExactly("CHILD_END");
    }

    @Test
    void shouldParseStreamOfForEachStates() throws IOException {
        String source = """
                package com.example;
                import java.util.stream.Stream;
                public class StreamStateConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineStateConfigurer states) throws Exception {
                        Stream.of("S1", "S2").forEach(s -> {
                            states.withStates().initial(s);
                        });
                    }
                }
                """;

        writeClass("StreamStateConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("StreamStateConfig");
        aggregator.aggregateStates(td);

        assertThat(aggregator.getInitialStates()).containsExactlyInAnyOrder("S1", "S2");
    }

    @Test
    void shouldAggregateTransitionsFromDeepInheritance() throws IOException {
        String grandparentSource = """
                package com.example;
                public abstract class GrandparentConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        transitions.withExternal().source("GP_S1").target("GP_S2").event("GP_E1");
                    }
                }
                """;
        String parentSource = """
                package com.example;
                public abstract class ParentConfig extends GrandparentConfig {
                    @Override
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        super.configure(transitions);
                        transitions.withExternal().source("P_S1").target("P_S2").event("P_E1");
                    }
                }
                """;
        String childSource = """
                package com.example;
                public class ChildConfig extends ParentConfig {
                    @Override
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        super.configure(transitions);
                        transitions.withExternal().source("C_S1").target("C_S2").event("C_E1");
                    }
                }
                """;

        writeClass("GrandparentConfig", grandparentSource);
        writeClass("ParentConfig", parentSource);
        writeClass("ChildConfig", childSource);
        context.scan(tempDir);

        TypeDeclaration childTd = context.getTypeDeclaration("ChildConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(childTd);

        assertThat(transitions).hasSize(3);
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("GP_E1"));
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("P_E1"));
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("C_E1"));
    }

    @Test
    void shouldParseEnhancedForLoopTransitions() throws IOException {
        String source = """
                package com.example;
                import java.util.List;
                public class LoopConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        for (String e : List.of("E1", "E2")) {
                            transitions.withExternal().source("S1").target("S2").event(e);
                        }
                    }
                }
                """;

        writeClass("LoopConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("LoopConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(td);

        assertThat(transitions).hasSize(2);
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("E1"));
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("E2"));
    }

    @Test
    void shouldIgnoreStatesInsideFork() throws IOException {
        String source = """
                package com.example;
                import org.springframework.statemachine.config.builders.StateMachineStateConfigurer;
                public class ForkStateConfig extends StateMachineConfigurerAdapter {
                    @Override
                    public void configure(StateMachineStateConfigurer states) throws Exception {
                        states
                            .withStates()
                            .initial("CREATED")
                            .fork("PROCESSING")
                            .and()
                            .withStates()
                            .parent("PROCESSING")
                            .initial("SHIPPING")
                            .initial("BILLING")
                            .end("COMPLETED");
                    }
                }
                """;

        writeClass("ForkStateConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("ForkStateConfig");
        aggregator.aggregateStates(td);

        assertThat(aggregator.getInitialStates()).containsExactly("CREATED");
        assertThat(aggregator.getEndStates()).containsExactly("COMPLETED");
    }

    @Test
    void shouldHandleNestedForkAndRegion() throws IOException {
        String source = """
                package com.example;
                import org.springframework.statemachine.config.builders.StateMachineStateConfigurer;
                public class ComplexStateConfig extends StateMachineConfigurerAdapter {
                    @Override
                    public void configure(StateMachineStateConfigurer states) throws Exception {
                        states
                            .withStates()
                            .initial("CREATED")
                            .fork("PROCESSING")
                            .and()
                            .withStates()
                            .parent("PROCESSING")
                            .region("SHIPPING")
                            .and()
                            .withStates()
                            .parent("SHIPPING")
                            .initial("PREPARING")
                            .end("SHIPPED")
                            .and()
                            .withStates()
                            .parent("PROCESSING")
                            .region("BILLING")
                            .and()
                            .withStates()
                            .parent("BILLING")
                            .initial("CALCULATING")
                            .end("BILLED")
                            .and()
                            .withStates()
                            .parent("PROCESSING")
                            .end("COMPLETED");
                    }
                }
                """;

        writeClass("ComplexStateConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("ComplexStateConfig");
        aggregator.aggregateStates(td);

        assertThat(aggregator.getInitialStates()).containsExactly("CREATED");
        assertThat(aggregator.getEndStates()).containsExactlyInAnyOrder("COMPLETED", "SHIPPED", "BILLED");
    }

    @Test
    void shouldFollowMethodCallsInOtherClasses() throws IOException {
        String otherClassSource = """
                package com.example;
                public class ExternalTransitions {
                    public static void addCommon(StateMachineTransitionConfigurer transitions) {
                        transitions.withExternal().source("S1").target("S2").event("EXT_E1");
                    }
                }
                """;
        String mainSource = """
                package com.example;
                public class MainConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        ExternalTransitions.addCommon(transitions);
                    }
                }
                """;

        writeClass("ExternalTransitions", otherClassSource);
        writeClass("MainConfig", mainSource);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("MainConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(td);

        // Currently StateMachineAggregator only searches within the hierarchy of the searchStartClass.
        // It might not find ExternalTransitions.addCommon unless it's static or we handle imports.
        // Let's see if it works or if we need to improve the aggregator.
        assertThat(transitions).hasSize(1);
        assertThat(transitions.get(0).getEvent().toString()).isEqualTo("EXT_E1");
    }

    @Test
    void shouldHandleRecursionSafety() throws IOException {
        String source = """
                package com.example;
                public class RecursiveConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        transitions.withExternal().source("S1").target("S2").event("E1");
                        recursiveMethod(transitions);
                    }
                    private void recursiveMethod(StateMachineTransitionConfigurer transitions) {
                        transitions.withExternal().source("S2").target("S3").event("E2");
                        recursiveMethod(transitions);
                    }
                }
                """;

        writeClass("RecursiveConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("RecursiveConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(td);

        // Should not crash and should find both transitions
        assertThat(transitions).hasSize(2);
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("E1"));
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("E2"));
    }

    @Test
    void shouldParseTransitionsInsideTryBlock() throws IOException {
        String source = """
                package com.example;
                public class TryConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        try {
                            transitions.withExternal().source("S1").target("S2").event("TRY_E1");
                        } catch (Exception e) {
                            // ignore
                        }
                    }
                }
                """;

        writeClass("TryConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("TryConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(td);

        assertThat(transitions).hasSize(1);
        assertThat(transitions.get(0).getEvent().toString()).isEqualTo("TRY_E1");
    }

    @Test
    void shouldHandleComplexStateConfigurationWithMultipleBlocks() throws IOException {
        String source = """
                package com.example;
                public class MultiBlockConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineStateConfigurer states) throws Exception {
                        states
                            .withStates()
                            .initial("START")
                            .state("S1")
                            .end("END_S1")
                            .and()
                            .withStates()
                            .parent("S1")
                            .initial("SUB_START")
                            .end("SUB_END");
                    }
                }
                """;

        writeClass("MultiBlockConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("MultiBlockConfig");
        aggregator.aggregateStates(td);

        assertThat(aggregator.getInitialStates()).containsExactlyInAnyOrder("START", "SUB_START");
        assertThat(aggregator.getEndStates()).containsExactlyInAnyOrder("END_S1", "SUB_END");
    }

    @Test
    void shouldIgnoreStatesInsideRegion() throws IOException {
        String source = """
                package com.example;
                public class RegionStateConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineStateConfigurer states) throws Exception {
                        states
                            .withStates()
                            .initial("CREATED")
                            .region("PROCESSING")
                            .and()
                            .withStates()
                            .parent("PROCESSING")
                            .initial("SHIPPING")
                            .end("COMPLETED");
                    }
                }
                """;

        writeClass("RegionStateConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("RegionStateConfig");
        aggregator.aggregateStates(td);

        assertThat(aggregator.getInitialStates()).containsExactly("CREATED");
        assertThat(aggregator.getEndStates()).containsExactly("COMPLETED");
    }

    @Test
    void shouldHandleEmptyConfigureMethod() throws IOException {
        String source = """
                package com.example;
                public class EmptyConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineStateConfigurer states) throws Exception {}
                }
                """;

        writeClass("EmptyConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("EmptyConfig");
        aggregator.aggregateStates(td);

        assertThat(aggregator.getInitialStates()).isEmpty();
    }

    @Test
    void shouldHandleMultipleMethodsHandlingTransitions() throws IOException {
        String source = """
                package com.example;
                public class MultiMethodConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        transitions.withExternal().source("S1").target("S2").event("E1");
                        otherTransitions(transitions);
                    }
                    private void otherTransitions(StateMachineTransitionConfigurer transitions) {
                        transitions.withExternal().source("S2").target("S3").event("E2");
                    }
                }
                """;

        writeClass("MultiMethodConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("MultiMethodConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(td);

        assertThat(transitions).hasSize(2);
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("E1"));
        assertThat(transitions).anyMatch(t -> t.getEvent().toString().equals("E2"));
    }

    @Test
    void shouldHandleMethodCallsWithEnumArguments() throws IOException {
        String source = """
                package com.example;
                public class EnumConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        addTransition(transitions, States.S1, States.S2);
                    }
                    private void addTransition(StateMachineTransitionConfigurer t, States s, States e) {
                        t.withExternal().source(s).target(e).event("ENUM_E1");
                    }
                }
                """;

        writeClass("EnumConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("EnumConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(td);

        assertThat(transitions).hasSize(1);
        assertThat(transitions.get(0).getSourceStates()).extracting(State::toString).containsExactly("States.S1");
    }

    @Test
    void shouldIgnoreStaticMethodsThatAreNotCalled() throws IOException {
        String source = """
                package com.example;
                public class UncalledStaticConfig extends StateMachineConfigurerAdapter {
                    public void configure(StateMachineTransitionConfigurer transitions) throws Exception {
                        transitions.withExternal().source("S1").target("S2").event("E1");
                    }
                    public static void uncalled(StateMachineTransitionConfigurer transitions) {
                        transitions.withExternal().source("S2").target("S3").event("E2");
                    }
                }
                """;

        writeClass("UncalledStaticConfig", source);
        context.scan(tempDir);

        TypeDeclaration td = context.getTypeDeclaration("UncalledStaticConfig");
        List<Transition> transitions = aggregator.aggregateTransitions(td);

        assertThat(transitions).hasSize(1);
        assertThat(transitions.get(0).getEvent().toString()).isEqualTo("E1");
    }

    private void writeClass(String name, String source) throws IOException {
        Files.writeString(tempDir.resolve(name + ".java"), source);
    }
}
