# (C) Copyright 2018- ECMWF.
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.

"""
Specific test battery for the OMNI parser frontend.
"""

import pytest

from loki import Module, Subroutine
from loki.backend import fgen
from loki.expression import symbols as sym
from loki.frontend import OMNI, HAVE_OMNI
from loki.ir import nodes as ir, FindNodes, FindVariables


@pytest.mark.skipif(not HAVE_OMNI, reason='Test tequires OMNI frontend.')
def test_derived_type_definitions(tmp_path):
    """ Test correct parsing of derived type declarations. """
    fcode = """
module omni_derived_type_mod
  type explicit
      real(kind=8) :: scalar, vector(3), matrix(3, 3)
  end type explicit

  type deferred
      real(kind=8), allocatable :: scalar, vector(:), matrix(:, :)
  end type deferred

  type ranged
      real(kind=8) :: scalar, vector(1:3), matrix(0:3, 0:3)
  end type ranged
end module omni_derived_type_mod
"""
    # Parse the source and validate the IR
    module = Module.from_source(fcode, frontend=OMNI, xmods=[tmp_path])

    assert len(module.typedefs) == 3
    explicit_symbols = FindVariables(unique=False).visit(module['explicit'].body)
    assert explicit_symbols == ('scalar', 'vector(3)', 'matrix(3, 3)')

    deferred_symbols = FindVariables(unique=False).visit(module['deferred'].body)
    assert deferred_symbols == ('scalar', 'vector(:)', 'matrix(:, :)')

    ranged_symbols = FindVariables(unique=False).visit(module['ranged'].body)
    assert ranged_symbols == ('scalar', 'vector(3)', 'matrix(0:3, 0:3)')


@pytest.mark.skipif(not HAVE_OMNI, reason='Test tequires OMNI frontend.')
def test_array_dimensions(tmp_path):
    """ Test correct parsing of derived type declarations. """
    fcode = """
subroutine omni_array_indexing(n, a, b)
  integer, intent(in) :: n
  real(kind=8), intent(inout) :: a(3), b(n)
  real(kind=8) :: c(n, n)
  real(kind=8) :: d(1:n, 0:n)

  a(:) = 11.
  b(1:n) = 42.
  c(2:n, 0:n) = 66.
  d(:, 0:n) = 68.
end subroutine omni_array_indexing
"""
    # Parse the source and validate the IR
    routine = Subroutine.from_source(fcode, frontend=OMNI, xmods=[tmp_path])

    # OMNI separate declarations per variable
    decls = FindNodes(ir.VariableDeclaration).visit(routine.spec)
    assert len(decls) == 5
    assert decls[0].symbols == ('n',)
    assert decls[1].symbols == ('a(3)',)
    assert decls[2].symbols == ('b(n)',)
    assert decls[3].symbols == ('c(n, n)',)
    assert decls[4].symbols == ('d(n, 0:n)',)

    assigns = FindNodes(ir.Assignment).visit(routine.body)
    assert len(assigns) == 4
    assert assigns[0].lhs == 'a(:)'
    assert assigns[1].lhs == 'b(1:n)'
    assert assigns[2].lhs == 'c(2:n, 0:n)'
    assert assigns[3].lhs == 'd(:, 0:n)'


@pytest.mark.skipif(not HAVE_OMNI, reason='Test tequires OMNI frontend.')
def test_labelled_format_with_quoted_items():
    """Test parsing of labelled FORMAT statements with quoted comma literals."""
    fcode = """
subroutine omni_labelled_format()
  implicit none
1003 format(5x,'NUMPROC=',i0,', NUMOMP=',i0,', NGPTOTG=',i0,', NPROMA=',i0,', NGPBLKS=',i0)
end subroutine omni_labelled_format
"""
    routine = Subroutine.from_source(fcode, frontend=OMNI)
    stmt = FindNodes(ir.FormatStmt).visit(routine.body)[0]

    assert stmt.label == '1003'
    assert stmt.values == (
        '5x', sym.StringLiteral('NUMPROC='), 'i0', sym.StringLiteral(', NUMOMP='), 'i0',
        sym.StringLiteral(', NGPTOTG='), 'i0', sym.StringLiteral(', NPROMA='), 'i0',
        sym.StringLiteral(', NGPBLKS='), 'i0'
    )
    assert fgen(routine.body) == (
        "1003 FORMAT(5x, 'NUMPROC=', i0, ', NUMOMP=', i0, ', NGPTOTG=', i0, "
        "', NPROMA=', i0, ', NGPBLKS=', i0)"
    )
