// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

using AwesomeAssertions;

using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.Generators;
using Microsoft.VisualStudio.TestTools.UnitTesting;

namespace Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.UnitTests;

[TestClass]
public sealed class ReflectionMetadataGeneratorTests
{
    private const string MinimalMSTestStub = """
        namespace Microsoft.VisualStudio.TestTools.UnitTesting
        {
            [System.AttributeUsage(System.AttributeTargets.Class)]
            public sealed class TestClassAttribute : System.Attribute {}

            [System.AttributeUsage(System.AttributeTargets.Method)]
            public class TestMethodAttribute : System.Attribute {}
        }
        """;

    private const string RuntimeHookStub = """
        namespace Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration
        {
            public static class ReflectionMetadataHook
            {
                public static void Register(System.Reflection.Assembly assembly, System.Type[] types, System.Collections.Generic.IReadOnlyDictionary<System.Type, System.Reflection.MethodInfo[]> testMethods) { }
            }
        }
        """;

    [TestMethod]
    public void Generator_DiscoversTestClassesAndMethods_AndEmitsModuleInitializer()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                public class MyTests
                {
                    [TestMethod]
                    public void Test1() {}

                    [TestMethod]
                    public void Test2() {}

                    public void NotATest() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        result.GeneratedSources.Should().HaveCount(1);

        // Full-snapshot match so any change to the emitter shape is caught by this test
        // (rather than only the small set of substrings other tests probe for).
        string generated = NormalizeNewlines(result.GeneratedSources[0].SourceText.ToString());
        const string expected = """
            // <auto-generated />
            // This file was generated by the MSTest source generator.
            // Do not edit it manually; changes will be lost on the next build.
            #nullable enable

            using System;
            using System.Collections.Generic;
            using System.Diagnostics.CodeAnalysis;
            using System.Reflection;
            using System.Runtime.CompilerServices;

            namespace Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.Generated
            {
                /// <summary>Source-generated MSTest reflection metadata hook for this test assembly.</summary>
                internal static class MSTestSourceGeneratedReflectionMetadata
                {
                    [ModuleInitializer]
                    [DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(global::Sample.MyTests))]
                    internal static void Initialize()
                    {
                        var assembly = typeof(MSTestSourceGeneratedReflectionMetadata).Assembly;
                        var types = new Type[]
                        {
                            typeof(global::Sample.MyTests),
                        };
                        var testMethods = new Dictionary<Type, MethodInfo[]>
                        {
                            [typeof(global::Sample.MyTests)] = new MethodInfo[]
                            {
                                ResolveMethod(typeof(global::Sample.MyTests), "Test1", Type.EmptyTypes),
                                ResolveMethod(typeof(global::Sample.MyTests), "Test2", Type.EmptyTypes),
                            },
                        };
                        global::Microsoft.VisualStudio.TestPlatform.MSTestAdapter.PlatformServices.SourceGeneration.ReflectionMetadataHook.Register(assembly, types, testMethods);
                    }

                    private static MethodInfo ResolveMethod([DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicMethods | DynamicallyAccessedMemberTypes.NonPublicMethods)] Type type, string name, Type[] parameterTypes)
                    {
                        const BindingFlags flags = BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Instance | BindingFlags.Static;
                        foreach (MethodInfo candidate in type.GetMethods(flags))
                        {
                            if (candidate.Name != name)
                            {
                                continue;
                            }

                            ParameterInfo[] candidateParameters = candidate.GetParameters();
                            if (candidateParameters.Length != parameterTypes.Length)
                            {
                                continue;
                            }

                            bool match = true;
                            for (int i = 0; i < candidateParameters.Length; i++)
                            {
                                if (candidateParameters[i].ParameterType != parameterTypes[i])
                                {
                                    match = false;
                                    break;
                                }
                            }

                            if (match)
                            {
                                return candidate;
                            }
                        }

                        throw new MissingMethodException(type.FullName, name);
                    }
                }
            }

            """;

        generated.Should().Be(NormalizeNewlines(expected));
    }

    [TestMethod]
    public void Generator_GeneratedFile_HasExpectedHintName()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                public class MyTests
                {
                    [TestMethod]
                    public void Test1() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        result.GeneratedSources.Should().HaveCount(1);

        // The hint name drives the on-disk path of the emitted file (when EmitCompilerGeneratedFiles
        // is enabled) and is used by downstream tooling to identify the generator output. It must
        // be deterministic: <AssemblyName><Constants.GeneratedFileSuffix>.
        result.GeneratedSources[0].HintName.Should().Be("TestSample.MSTestReflectionMetadata.g.cs");
    }

    [TestMethod]
    public void Generator_EmitsOverloadAwareMethodResolution()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                public class OverloadTests
                {
                    [TestMethod]
                    public void Run() {}

                    [TestMethod]
                    public void Run(int x) {}

                    [TestMethod]
                    public void Run(string s, int x) {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = NormalizeNewlines(result.GeneratedSources[0].SourceText.ToString());

        // Match the full testMethods block as a contiguous snippet so order, the new-array vs
        // Type.EmptyTypes choice, and the FullName-qualified parameter types are all locked in.
        const string expectedTypeMethodsBlock = """
                    var testMethods = new Dictionary<Type, MethodInfo[]>
                    {
                        [typeof(global::Sample.OverloadTests)] = new MethodInfo[]
                        {
                            ResolveMethod(typeof(global::Sample.OverloadTests), "Run", Type.EmptyTypes),
                            ResolveMethod(typeof(global::Sample.OverloadTests), "Run", new Type[] { typeof(global::System.Int32) }),
                            ResolveMethod(typeof(global::Sample.OverloadTests), "Run", new Type[] { typeof(global::System.String), typeof(global::System.Int32) }),
                        },
                    };
        """;

        generated.Should().Contain(NormalizeNewlines(expectedTypeMethodsBlock));
    }

    [TestMethod]
    public void Generator_EmittedSourceCompilesCleanly()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                public class MyTests
                {
                    [TestMethod]
                    public void Test1() {}

                    [TestMethod]
                    public void Test2(int x, string s) {}

                    [TestMethod]
                    internal void NonPublicTest() {}
                }
            }
            """;

        Compilation outputCompilation = RunGeneratorAndGetCompilation(MinimalMSTestStub, RuntimeHookStub, userCode);

        Diagnostic[] errors = outputCompilation.GetDiagnostics()
            .Where(d => d.Severity == DiagnosticSeverity.Error)
            .ToArray();

        errors.Should().BeEmpty("the generator output should compile without errors. Diagnostics: " + string.Join("\n", errors.Select(d => d.ToString())));
    }

    [TestMethod]
    public void Generator_SkipsStaticAndAbstractTestClasses()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                public static class StaticTests
                {
                    [TestMethod]
                    public static void Test1() {}
                }

                [TestClass]
                public abstract class AbstractTests
                {
                    [TestMethod]
                    public void Test1() {}
                }

                [TestClass]
                public class Concrete : AbstractTests
                {
                    [TestMethod]
                    public void Test2() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // Neither the static nor the abstract test class is a runnable test class, so they must
        // not be emitted into the `types[]` array or the `testMethods{}` dictionary. The
        // discovery-side rules have not changed.
        generated.Should().NotContain("typeof(global::Sample.StaticTests)");
        generated.Should().NotContain("[typeof(global::Sample.AbstractTests)]");

        // The concrete derived class is the runnable test class — it must be emitted.
        generated.Should().Contain("typeof(global::Sample.Concrete)");

        // Concrete must surface BOTH its own Test2 AND the inherited Test1 from AbstractTests.
        generated.Should().Contain("ResolveMethod(typeof(global::Sample.Concrete), \"Test1\", Type.EmptyTypes)");
        generated.Should().Contain("ResolveMethod(typeof(global::Sample.Concrete), \"Test2\", Type.EmptyTypes)");

        // The abstract base class must still be referenced by a [DynamicDependency] on the
        // generated module initializer so that members declared on it ([ClassInitialize],
        // [ClassCleanup], [AssemblyInitialize], [AssemblyCleanup], TestContext setter) are
        // preserved by the IL trimmer / Native AOT. Without this hint those members live only
        // on the abstract base and are trimmed because [DynamicDependency(All, typeof(Concrete))]
        // does not preserve base-type members.
        generated.Should().Contain("[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(global::Sample.AbstractTests))]");
    }

    [TestMethod]
    public void Generator_EmitsNestedTestClass_WithTypeofReference()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                public class Outer
                {
                    [TestClass]
                    public class Nested
                    {
                        [TestMethod]
                        public void Test1() {}
                    }
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // For nested types we must emit `typeof(Outer.Nested)` (dot-qualified C# name) — never
        // the runtime FullName form `Outer+Nested` which would not compile. The builder computes
        // TypesByName at runtime from typeof(T).FullName, so the emitted source must only carry
        // the C# typeof reference.
        generated.Should().Contain("typeof(global::Sample.Outer.Nested)");
        generated.Should().NotContain("Sample.Outer+Nested");
    }

    [TestMethod]
    public void Generator_HandlesEmptyAssembly()
    {
        const string userCode = """
            namespace Sample
            {
                public class NoTests {}
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();

        // When no [TestClass] types are discovered, the generator must not emit any source
        // (avoids polluting the compilation with an empty module initializer).
        result.GeneratedSources.Should().BeEmpty();
    }

    [TestMethod]
    public void Generator_SkipsGenericTestMethods()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                public class GenericMethodTests
                {
                    [TestMethod]
                    public void GenericTest<T>(T value) {}

                    [TestMethod]
                    public void GenericTestNoParams<T>() {}

                    [TestMethod]
                    public void NormalTest(int x) {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // Generic test methods cannot be emitted by the non-generic module initializer:
        // parameter types would include method-level type parameters (typeof(T)) which do not
        // bind at module-initializer scope, breaking compilation. They are silently skipped so
        // opting into the generator never introduces compile errors that reflection mode would
        // tolerate. Plain test methods on the same class must still be emitted.
        generated.Should().NotContain("\"GenericTest\"");
        generated.Should().NotContain("\"GenericTestNoParams\"");
        generated.Should().Contain("ResolveMethod(typeof(global::Sample.GenericMethodTests), \"NormalTest\", new Type[] { typeof(global::System.Int32) })");
    }

    [TestMethod]
    public void Generator_SkipsOpenGenericTestClasses()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                public class Generic<T>
                {
                    [TestMethod]
                    public void Test1() {}
                }

                [TestClass]
                public class Concrete
                {
                    [TestMethod]
                    public void Test2() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // Open generic test classes cannot be emitted as `typeof(Generic<T>)` at module-initializer
        // scope, so they are silently skipped. The non-generic sibling must still be emitted.
        generated.Should().NotContain("typeof(global::Sample.Generic");
        generated.Should().Contain("typeof(global::Sample.Concrete)");
    }

    [TestMethod]
    public void Generator_SkipsMethodsWithByRefParameters()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                public class ByRefTests
                {
                    [TestMethod]
                    public void RefTest(ref int x) {}

                    [TestMethod]
                    public void OutTest(out int x) { x = 0; }

                    [TestMethod]
                    public void InTest(in int x) {}

                    [TestMethod]
                    public void RefReadonlyTest(ref readonly int x) {}

                    [TestMethod]
                    public void NormalTest(int x) {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // By-ref signatures (ref/out/in/ref readonly) cannot round-trip through ResolveMethod's
        // typeof(T) == ParameterType check, so the generator omits these methods entirely.
        // The plain-parameter overload must still be emitted.
        generated.Should().NotContain("\"RefTest\"");
        generated.Should().NotContain("\"OutTest\"");
        generated.Should().NotContain("\"InTest\"");
        generated.Should().NotContain("\"RefReadonlyTest\"");
        generated.Should().Contain("ResolveMethod(typeof(global::Sample.ByRefTests), \"NormalTest\", new Type[] { typeof(global::System.Int32) })");
    }

    [TestMethod]
    public void Generator_SkipsInaccessibleNestedTestClasses()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                public class Outer
                {
                    [TestClass]
                    private class PrivateNested
                    {
                        [TestMethod]
                        public void Test1() {}
                    }

                    [TestClass]
                    public class PublicNested
                    {
                        [TestMethod]
                        public void Test2() {}
                    }
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // A private (or otherwise inaccessible) nested [TestClass] cannot be referenced from
        // the generated `internal` module initializer; it would compile as CS0122. The
        // generator skips such types but still emits siblings that are reachable.
        generated.Should().NotContain("PrivateNested");
        generated.Should().Contain("typeof(global::Sample.Outer.PublicNested)");
    }

    [TestMethod]
    public void Generator_SkipsFileLocalTestClasses()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                file class FileLocalTests
                {
                    [TestMethod]
                    public void HiddenTest() {}
                }

                [TestClass]
                public class VisibleTests
                {
                    [TestMethod]
                    public void OpenTest() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // `file`-scoped types are only addressable within their own source file. The generated
        // module initializer lives in a different file and would fail with CS9051. Skip them.
        generated.Should().NotContain("FileLocalTests");
        generated.Should().NotContain("\"HiddenTest\"");
        generated.Should().Contain("typeof(global::Sample.VisibleTests)");
    }

    [TestMethod]
    public void Generator_DedupesOverriddenMethodsAcrossInheritance()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                public abstract class BaseTests
                {
                    [TestMethod]
                    public virtual void TestA() {}
                }

                [TestClass]
                public class Derived : BaseTests
                {
                    [TestMethod]
                    public override void TestA() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // The base + override must collapse to a single emit. Counting occurrences of the
        // exact ResolveMethod call site catches double-emission regressions in the inheritance
        // walk's dedupe-by-signature logic.
        int count = CountOccurrences(generated, "ResolveMethod(typeof(global::Sample.Derived), \"TestA\", Type.EmptyTypes)");
        count.Should().Be(1);
    }

    [TestMethod]
    public void Generator_EmitsMultipleTestClasses_InOrder()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                public class AlphaTests
                {
                    [TestMethod]
                    public void TestA() {}
                }

                [TestClass]
                public class BetaTests
                {
                    [TestMethod]
                    public void TestB() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = NormalizeNewlines(result.GeneratedSources[0].SourceText.ToString());

        // Two [TestClass]es must produce two entries in the types/testMethods locals and
        // two [DynamicDependency] attributes on Initialize. We assert each block as a
        // contiguous snippet so the per-class layout is locked in.
        const string expectedDynamicDependencies = """
                [DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(global::Sample.AlphaTests))]
                [DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(global::Sample.BetaTests))]
        """;
        generated.Should().Contain(NormalizeNewlines(expectedDynamicDependencies));

        const string expectedTypes = """
                    var types = new Type[]
                    {
                        typeof(global::Sample.AlphaTests),
                        typeof(global::Sample.BetaTests),
                    };
        """;
        generated.Should().Contain(NormalizeNewlines(expectedTypes));

        const string expectedTypeMethods = """
                    var testMethods = new Dictionary<Type, MethodInfo[]>
                    {
                        [typeof(global::Sample.AlphaTests)] = new MethodInfo[]
                        {
                            ResolveMethod(typeof(global::Sample.AlphaTests), "TestA", Type.EmptyTypes),
                        },
                        [typeof(global::Sample.BetaTests)] = new MethodInfo[]
                        {
                            ResolveMethod(typeof(global::Sample.BetaTests), "TestB", Type.EmptyTypes),
                        },
                    };
        """;
        generated.Should().Contain(NormalizeNewlines(expectedTypeMethods));
    }

    [TestMethod]
    public void Generator_EmitsTestClass_InGlobalNamespace()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            [TestClass]
            public class GlobalNamespaceTests
            {
                [TestMethod]
                public void Test1() {}
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // A test class declared in the global namespace must be referenced with `global::TypeName`
        // (no namespace segment between `global::` and the type), and still produce both the
        // types-array entry and the testMethods-dictionary entry.
        generated.Should().Contain("typeof(global::GlobalNamespaceTests)");
        generated.Should().Contain("ResolveMethod(typeof(global::GlobalNamespaceTests), \"Test1\", Type.EmptyTypes)");
    }

    [TestMethod]
    public void Generator_DiscoversMethodsViaCustomTestMethodAttributeSubclass()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                public sealed class CustomTestMethodAttribute : TestMethodAttribute {}

                [TestClass]
                public class CustomAttributeTests
                {
                    [CustomTestMethod]
                    public void DerivedAttributeTest() {}

                    [TestMethod]
                    public void PlainTest() {}

                    public void NotATest() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // The generator's HasTestMethodAttribute walks the attribute class's base type chain,
        // so any subclass of TestMethodAttribute (e.g. DataTestMethod, custom attributes) must
        // be recognised the same way as TestMethod itself.
        generated.Should().Contain("ResolveMethod(typeof(global::Sample.CustomAttributeTests), \"DerivedAttributeTest\", Type.EmptyTypes)");
        generated.Should().Contain("ResolveMethod(typeof(global::Sample.CustomAttributeTests), \"PlainTest\", Type.EmptyTypes)");
        generated.Should().NotContain("\"NotATest\"");
    }

    [TestMethod]
    public void Generator_EmitsInternalTestClasses()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                internal class InternalTests
                {
                    [TestMethod]
                    public void Test1() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // Internal [TestClass]es are reachable from the generated `internal` module
        // initializer (same assembly) and so must be emitted just like public ones.
        generated.Should().Contain("typeof(global::Sample.InternalTests)");
        generated.Should().Contain("ResolveMethod(typeof(global::Sample.InternalTests), \"Test1\", Type.EmptyTypes)");
    }

    [TestMethod]
    public void Generator_EmitsDynamicDependencyForFullBaseTypeChain()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                public abstract class GrandparentTests
                {
                    [AssemblyInitialize]
                    public static void AssemblyInit(TestContext context) {}
                }

                public abstract class ParentTests : GrandparentTests
                {
                    [ClassInitialize]
                    public static void ClassInit(TestContext context) {}
                }

                [TestClass]
                public class Derived : ParentTests
                {
                    [TestMethod]
                    public void Test1() {}
                }
            }
            """;

        const string testContextStub = """
            namespace Microsoft.VisualStudio.TestTools.UnitTesting
            {
                public class TestContext {}

                [System.AttributeUsage(System.AttributeTargets.Method)]
                public sealed class AssemblyInitializeAttribute : System.Attribute {}

                [System.AttributeUsage(System.AttributeTargets.Method)]
                public sealed class ClassInitializeAttribute : System.Attribute {}
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, testContextStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // Every base type up the chain must be referenced by a [DynamicDependency] so the
        // trimmer keeps members like [AssemblyInitialize] (declared on GrandparentTests) and
        // [ClassInitialize] (declared on ParentTests). The concrete [TestClass] still owns the
        // discovery entry, but the trimmer roots are emitted per type.
        generated.Should().Contain("[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(global::Sample.Derived))]");
        generated.Should().Contain("[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(global::Sample.ParentTests))]");
        generated.Should().Contain("[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(global::Sample.GrandparentTests))]");

        // Bases must NOT appear in the types[] array or the testMethods{} dictionary — they are
        // not runnable test classes; discovery still goes through the concrete [TestClass].
        generated.Should().NotContain("typeof(global::Sample.GrandparentTests),");
        generated.Should().NotContain("[typeof(global::Sample.GrandparentTests)]");
        generated.Should().NotContain("typeof(global::Sample.ParentTests),");
        generated.Should().NotContain("[typeof(global::Sample.ParentTests)]");
    }

    [TestMethod]
    public void Generator_DedupesDynamicDependencyForBaseShared()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                public abstract class SharedBase
                {
                    [TestMethod]
                    public void Inherited() {}
                }

                [TestClass]
                public class DerivedA : SharedBase
                {
                    [TestMethod]
                    public void TestA() {}
                }

                [TestClass]
                public class DerivedB : SharedBase
                {
                    [TestMethod]
                    public void TestB() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // The shared abstract base produces only ONE [DynamicDependency] regardless of how
        // many derived test classes reference it. Without deduping, every shared base would
        // be emitted N times for N derived classes, bloating the generated source.
        int sharedBaseCount = CountOccurrences(
            generated,
            "[DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(global::Sample.SharedBase))]");
        sharedBaseCount.Should().Be(1);
    }

    [TestMethod]
    public void Generator_DoesNotEmitDynamicDependencyForObjectBase()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                [TestClass]
                public class StandaloneTests
                {
                    [TestMethod]
                    public void Test1() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // The base-type walk stops at System.Object — we never emit a [DynamicDependency] for
        // it (would be useless noise and would increase the surface of types we root with the trimmer).
        generated.Should().NotContain("typeof(global::System.Object)");
        generated.Should().NotContain("typeof(object)");
    }

    [TestMethod]
    public void Generator_DoesNotEmitDynamicDependencyForInaccessibleBase()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                public class Outer
                {
                    // A private nested base cannot be referenced from the generated `internal`
                    // module initializer (CS0122). Skip it from [DynamicDependency] emission to
                    // keep the generated source compilable.
                    private class PrivateBase
                    {
                        [TestMethod]
                        public void InheritedFromPrivate() {}
                    }

                    [TestClass]
                    public class PublicDerived : PrivateBase
                    {
                        [TestMethod]
                        public void TestB() {}
                    }
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // The private base must not surface in any reference; the derived class still must.
        generated.Should().NotContain("PrivateBase");
        generated.Should().Contain("typeof(global::Sample.Outer.PublicDerived)");
    }

    [TestMethod]
    public void Generator_DoesNotEmitDynamicDependencyForGenericBase()
    {
        const string userCode = """
            using Microsoft.VisualStudio.TestTools.UnitTesting;

            namespace Sample
            {
                public abstract class GenericBase<T>
                {
                    [TestMethod]
                    public void InheritedTest() {}
                }

                [TestClass]
                public class Derived : GenericBase<int>
                {
                    [TestMethod]
                    public void Test1() {}
                }
            }
            """;

        GeneratorRunResult result = RunGenerator(MinimalMSTestStub, userCode);

        result.Diagnostics.Should().BeEmpty();
        string generated = result.GeneratedSources[0].SourceText.ToString();

        // For this conservative iteration we skip generic base types from [DynamicDependency]
        // emission — emitting `typeof(GenericBase<int>)` is technically valid but adds edge
        // cases (open vs. closed, type-parameter capture) that the rest of the generator does
        // not exercise. The derived class itself is still emitted, and the inherited test
        // method is still surfaced through the base-walk in the methods collection.
        generated.Should().NotContain("typeof(global::Sample.GenericBase");
        generated.Should().Contain("typeof(global::Sample.Derived)");
        generated.Should().Contain("ResolveMethod(typeof(global::Sample.Derived), \"InheritedTest\", Type.EmptyTypes)");
    }

    private static string NormalizeNewlines(string value)
        => value.Replace("\r\n", "\n");

    private static int CountOccurrences(string source, string needle)
    {
        int count = 0;
        int index = 0;
        while ((index = source.IndexOf(needle, index, StringComparison.Ordinal)) >= 0)
        {
            count++;
            index += needle.Length;
        }

        return count;
    }

    private static GeneratorRunResult RunGenerator(params string[] sources)
    {
        IEnumerable<SyntaxTree> trees = sources.Select(s => CSharpSyntaxTree.ParseText(s));
        var references = new MetadataReference[]
        {
            MetadataReference.CreateFromFile(typeof(object).Assembly.Location),
            MetadataReference.CreateFromFile(typeof(System.Runtime.CompilerServices.ModuleInitializerAttribute).Assembly.Location),
        };

        var compilation = CSharpCompilation.Create(
            "TestSample",
            trees,
            references,
            new CSharpCompilationOptions(OutputKind.DynamicallyLinkedLibrary));

        GeneratorDriver driver = CSharpGeneratorDriver.Create(new ReflectionMetadataGenerator());
        driver = driver.RunGeneratorsAndUpdateCompilation(compilation, out _, out _);

        return driver.GetRunResult().Results[0];
    }

    private static Compilation RunGeneratorAndGetCompilation(params string[] sources)
    {
        IEnumerable<SyntaxTree> trees = sources.Select(s => CSharpSyntaxTree.ParseText(s));
        var references = new MetadataReference[]
        {
            MetadataReference.CreateFromFile(typeof(object).Assembly.Location),
            MetadataReference.CreateFromFile(typeof(System.Runtime.CompilerServices.ModuleInitializerAttribute).Assembly.Location),
            MetadataReference.CreateFromFile(typeof(System.Reflection.Assembly).Assembly.Location),
            MetadataReference.CreateFromFile(typeof(System.Collections.Generic.Dictionary<,>).Assembly.Location),
            MetadataReference.CreateFromFile(typeof(System.Reflection.MethodInfo).Assembly.Location),
            MetadataReference.CreateFromFile(typeof(System.Reflection.BindingFlags).Assembly.Location),
            MetadataReference.CreateFromFile(typeof(System.MissingMethodException).Assembly.Location),
            MetadataReference.CreateFromFile(typeof(System.Linq.Enumerable).Assembly.Location),
        };

        var compilation = CSharpCompilation.Create(
            "TestSample",
            trees,
            references,
            new CSharpCompilationOptions(OutputKind.DynamicallyLinkedLibrary));

        GeneratorDriver driver = CSharpGeneratorDriver.Create(new ReflectionMetadataGenerator());
        driver.RunGeneratorsAndUpdateCompilation(compilation, out Compilation outputCompilation, out _);

        return outputCompilation;
    }
}
