module MultidimensionalArrayTests

using JuliaFormatter
using JuliaFormatter.Internal: test_format, ALL_STYLES
using Test

@testset "array literal examples" begin
    # Don't really care how it formats, but make sure that it parses to the same AST. Some
    # of these are probably redundant; that's fine, we can just be safe.
    for T in ("", "T", "x = ")
        for s in (
            # These are taken from the docs:
            # https://docs.julialang.org/en/v1/manual/arrays/#man-array-literals
            "$(T)[1, 2, 3]",
            "$(T)[1, 2.3, 4//5]",
            "$(T)[1:2, 4:5]",
            "$(T)[1:2\n 4:5\n 6]",
            "$(T)[1:2  4:5  7:8]",
            "$(T)[[1,2]  [4,5]  [7,8]]",
            "$(T)[1 2 3]",
            "$(T)[1;; 2;; 3;; 4]",
            "$(T)[1 2\n 3 4]",
            "$(T)[zeros(Int, 2, 2) [1; 2]\n [3 4]            5]",
            "$(T)[[1 1]; 2 3; [4 4]]",
            "$(T)[zeros(Int, 2, 2) ; [3 4] ;; [1; 2] ; 5]",
            "$(T)[1:2; 4;; 1; 3:4]",
            "$(T)[1; 2;; 3; 4;; 5; 6;;;\n 7; 8;; 9; 10;; 11; 12]",
            "$(T)[1 3 5\n 2 4 6;;;\n 7 9 11\n 8 10 12]",
            "$(T)[1 2;;; 3 4;;;; 5 6;;; 7 8]",
            "$(T)[[1 2;;; 3 4];;;; [5 6];;; [7 8]]",
            "$(T)[1 2 ;;\n 3 4]",
            "$(T)[1;;]",
            "$(T)[2; 3;;;]",
            "$(T)[[1 2] [3 4]]",

            # These are supplements, not from the docs
            "$(T)[1\n2\n3\n4]",
            "$(T)[1\n2;;\n3\n4]",
            "$(T)[f(a)\n f(a);;\n f(a)\n f(a)]",
            "$(T)[\n1\n2;;\n3\n4\n]", # extra newlines
            "$(T)[]",
            "$(T)[;]",
            "$(T)[;;]",
            "$(T)[;;;]",
            "$(T)[1 2 ;;\n 3 4 ; 2 3 4 5]", # issue 1080
            "$(T)[1 2 ;;\n 3 4 ;;; 5 6 ;;\n 7 8]", # issue 1080
            "$(T)[1\n\n2\n\n3]", # issue 1242
        )
            for style in ALL_STYLES
                @testset let style = style
                    test_format(s, nothing, style; ast=true)
                end
            end
        end
    end
end

@testset "hcat and typed_hcat nodes" begin
    # hcat nodes are the ones without T; typed_hcat nodes are the ones with T. Both
    # should more or less be formatted the same way.
    for T in ("", "T")
        @testset "spaces only" begin
            # Additionally check that superfluous whitespace is removed.
            for s in ("x = $(T)[a b]", "x = $(T)[a      b     ]")
                for style in ALL_STYLES
                    test_format(s, "x = $(T)[a b]", style)
                end
                test_format(s, "x = $(T)[\n    a b\n]"; margin = 5+length(T))

                # TODO(penelopeysm): The indentation for other styles is all over the
                # place. But we can at least check for now that it parses to the same
                # AST (meaning that at worst, it's ugly, but not wrong), and that it's
                # idempotent.
                for style in ALL_STYLES,
                    kwargs in
                    ((;), (; margin = 5+length(T)), (; join_lines_based_on_source = false))

                    test_format(s, nothing, style; kwargs...)
                end
            end
        end

        @testset ";;\\n only" begin
            for s in ("x = $(T)[a;;\n b]", "x = $(T)[a    ;;\n   b   ]")
                # because the original matrix is already split across a newline,
                # the formatter will insert more newlines
                expected_default = "x = $(T)[\n    a;;\n    b\n]"
                test_format(s, expected_default)
                test_format(s, expected_default; margin = 5+length(T))

                ws = " " ^ (5 + length(T))
                expected_sciml = "x = $(T)[a;;\n$(ws)b]"
                test_format(s, expected_sciml, SciMLStyle())
                test_format(s, expected_sciml, SciMLStyle(); margin = 5+length(T))
                test_format(s, expected_sciml, YASStyle())
                test_format(s, expected_sciml, YASStyle(); margin = 5+length(T))
            end
        end

        @testset "mixture of space and ;;\\n" begin
            for s in ("x = $(T)[a b;;\n c d]", "x = $(T)[a    b   ;;\n   c    d   ]")
                # because the original matrix is already split across a newline,
                # the formatter will insert more newlines
                expected = "x = $(T)[\n    a b;;\n    c d\n]"
                test_format(s, expected)
                test_format(s, expected; margin = 5+length(T))

                # TODO(penelopeysm): The indentation for other styles is all over the
                # place. But we can at least check for now that it parses to the same
                # AST (meaning that at worst, it's ugly, but not wrong), and that it's
                # idempotent.
                kwargs_combos = (
                    (;),
                    (; margin = 5+length(T)),
                    (; join_lines_based_on_source = false)
                )
                for style in ALL_STYLES, kwargs in kwargs_combos
                    test_format(s, nothing, style; kwargs...)
                end
            end
        end

        @testset "comments are not lost" begin
            s = "[a b;; # comment\nc d]"
            for st in (DefaultStyle(), BlueStyle(), MinimalStyle())
                target = "[a b;; # comment\n    c d]"
                test_format(s, target, st; ast=true)
            end
            target = "[a b;; # comment\n c d]"
            for st in (SciMLStyle(), YASStyle())
                test_format(s, target, st; ast=true)
            end

            s = "[#=1=# a #=2=# b #=3=# ;;\n#=4=# c #=5=# d #=6=#]"
            for st in ALL_STYLES
                # Idempotence & semantic invariance
                test_format(s, nothing, st; ast=true)
                # Check that the comment is still there
                out = format_text(s, st)
                for i in 1:6
                    @test occursin("#=$(i)=#", out)
                end
            end
        end
    end
end

@testset "ncat and typed ncat nodes" begin
    @testset "newline and ;;" begin
        s = "[1\n2;;\n3\n4]"
        # Default and Blue style will impose the standard 4-space indentation
        expected_indent4 = "[\n    1\n    2;;\n    3\n    4\n]"
        # SciML and YAS style will indent to the opening brace
        expected_indent1 = "[1\n 2;;\n 3\n 4]"
        for style in (DefaultStyle(), BlueStyle())
            test_format(s, expected_indent4, style)
            test_format(s, expected_indent4, style; margin = 1)
        end
        for style in (SciMLStyle(), YASStyle())
            test_format(s, expected_indent1, style)
            test_format(s, expected_indent1, style; margin = 1)
        end
        # MinimalStyle keeps the brackets where they are, and indents every continuation
        # line by one level
        test_format(s, "[1\n    2;;\n    3\n    4]", MinimalStyle())
    end
end

@testset "MinimalStyle indentation of array literals" begin
    # MinimalStyle leaves the brackets where they are, so every line after the opening
    # bracket gets exactly one extra level of indentation, whichever kind of concatenation
    # the newline stands for.
    @testset "rows are all indented equally" begin
        for (s, expected) in (
            ("[1\n2\n3]", "[1\n    2\n    3]"),
            ("[1\n2;;\n3]", "[1\n    2;;\n    3]"),
            ("[1\n2;;\n3\n4]", "[1\n    2;;\n    3\n    4]"),
            ("[1;\n2;;\n3;\n4]", "[1;\n    2;;\n    3;\n    4]"),
            ("[1;;\n2]", "[1;;\n    2]"),
            ("[1 2;;\n3 4]", "[1 2;;\n    3 4]"),
            # rows nested inside rows
            ("[1\n2;;\n3;;;\n4]", "[1\n    2;;\n    3;;;\n    4]"),
            ("[1\n2;;;\n3\n4]", "[1\n    2;;;\n    3\n    4]"),
            # hcat rows inside a vcat
            ("[1 2\n3 4]", "[1 2\n    3 4]"),
            ("T[1\n2;;\n3]", "T[1\n    2;;\n    3]"),
        )
            @testset let s = s
                test_format(s, expected, MinimalStyle(); ast=true)
            end
        end
    end

    @testset "nested brackets" begin
        s = """
        [[1,
        2]; 3;;
        4]"""
        expected = """
        [[1,
            2]; 3;;
            4]"""
        test_format(s, expected, MinimalStyle(); ast=true)

        s = """
        [[1,
        2] 3
        4 5]"""
        expected = """
        [[1,
            2] 3
            4 5]"""
        test_format(s, expected, MinimalStyle(); ast=true)

        s = """
        [1 2
        [3,
        4] 5]"""
        expected = """
        [1 2
            [3,
                4] 5]"""
        test_format(s, expected, MinimalStyle(); ast=true)
    end

    @testset "indentation is relative to the enclosing scope" begin
        s = """
        function f()
            x = [1
            2;;
            3
            4]
            return x
        end"""
        expected = """
        function f()
            x = [1
                2;;
                3
                4]
            return x
        end"""
        test_format(s, expected, MinimalStyle(); ast=true)
    end
end

@testset "#490" begin
    @testset "DefaultStyle" begin
        str = "[1;; 2]"
        test_format(str, str)

        str_ = """
        [
            1;;
            2
        ]"""
        test_format(str, str_; margin = 1)

        str = "T[1;; 2]"
        test_format(str, str)

        str_ = """
        T[
            1;;
            2
        ]"""
        test_format(str, str_; margin = 1)
    end

    @testset "YASStyle" begin
        str = "[1;; 2]"
        test_format(str, str, YASStyle())

        str_ = """
        [1;;
         2]"""
        test_format(str, str_, YASStyle(); margin = 1)

        str = "T[1;; 2]"
        test_format(str, str, YASStyle())

        str_ = """
        T[1;;
          2]"""
        test_format(str, str_, YASStyle(); margin = 1)
    end
end

@testset "#620" begin
    s = "[1; 0;;]"
    test_format(s, s)
    test_format(s, s, YASStyle())
end

@testset "#582" begin
    test_format("[a;b;;]", "[a; b;;]", YASStyle())
end

@testset "#608" begin
    s1 = """
    hcat([zeros(1); ones(3)], [zeros(2); ones(2)], [zeros(3); ones(1)], [zeros(1); ones(3)], [zeros(2); ones(2)], [zeros(3); ones(1)])
    """
    s2 = """
    hcat([zeros(1); ones(3)], [zeros(2); ones(2)], [zeros(3); ones(1)], [zeros(1); ones(3)],
         [zeros(2); ones(2)], [zeros(3); ones(1)])
    """
    test_format(s1, s2, YASStyle())
end

@testset "#532" begin
    s = "(; a = [1;;], b = cos[2;;])"
    s_nospace = "(; a=[1;;], b=cos[2;;])"
    for style in (DefaultStyle(), SciMLStyle())
        test_format(s, s, style)
    end
    # whitespace_around_kwarg = false
    for style in (BlueStyle(), YASStyle(), MinimalStyle())
        test_format(s, s_nospace, style)
    end
end

@testset "#1113 - comments in vcat array literals" begin
    for s in (
        "X = [\n    # leading\n    1 2\n    3 4\n]\n",
        "X = [\n    1 2\n    3 4\n    # trailing\n]\n",
        "X = [\n    # foo\n    1 2\n    3 4\n    # bar\n]\n",
        "X = [\n    1 2\n    # between\n    3 4\n]\n",
        "X = Int[\n    # foo\n    1 2\n    3 4\n    # bar\n]\n",
    )
        for style in ALL_STYLES
            if style isa DefaultStyle
                # Default should leave it unchanged
                test_format(s, s; ast = true)
            else
                test_format(s, nothing, style; ast = true)
            end
        end
    end
end

@testset "#1113 - comments in ncat array literals" begin
    for s in (
        "[\n    # leading\n    1\n    2;;\n    3\n    4\n]\n",
        "[\n    1\n    2;;\n    3\n    4\n    # trailing\n]\n",
        "[\n    1\n    2;;\n    # between\n    3\n    4\n]\n",
        "[\n    1;\n    2;;\n    # between\n    3;\n    4\n]\n",
    )
        for style in ALL_STYLES
            if style isa DefaultStyle
                # Default should leave it unchanged
                test_format(s, s; ast = true)
            else
                test_format(s, nothing, style; ast = true)
            end
        end
    end

    # this one is still a bit wonky because the \n after 2;; is collapsed
    s = "[\n    # leading\n    1\n    2;;\n    3\n    4;;;\n    5\n    6;;\n    7\n    8\n]\n"
    for style in ALL_STYLES
        test_format(s, nothing, style; ast = true)
        @test occursin("# leading", format_text(s))
    end
    @test_broken format_text(s) == s
end

@testset "#1113 - comments in hcat array literals" begin
    for s in (
        "X = [\n    # only\n    1 2\n]\n",
        "X = [\n    1 2\n    # trailing\n]\n",
        "X = [\n    # foo\n    1 2\n    # bar\n]\n",
        "X = [\n    1 2;;\n    # hi\n    3 4\n]\n",
        "X = [#= inline =# 1 2]\n",
        "X = [1 #= mid =# 2]\n",
    )
        for style in ALL_STYLES
            if style isa DefaultStyle
                test_format(s, s; ast = true)
            else
                test_format(s, nothing, style; ast = true)
            end
        end
    end
end

@testset "1168 vcats in rows of other arrays" begin
    s = """[x [y z
        [q
        s
        # a
        # b
        ]]
        f]"""
    for style in ALL_STYLES
        test_format(s, nothing, style; ast=true)
    end
end

end # module
