#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""Unit tests for the processes module."""

# pytype: skip-file

import glob
import os
import random
import re
import shutil
import socketserver
import subprocess
import tempfile
import threading
import unittest

from apache_beam.utils import subprocess_server


class JavaJarServerTest(unittest.TestCase):
  def test_gradle_jar_release(self):
    self.assertEqual(
        'https://repo.maven.apache.org/maven2/org/apache/beam/'
        'beam-sdks-java-fake/VERSION/beam-sdks-java-fake-VERSION.jar',
        subprocess_server.JavaJarServer.path_to_beam_jar(
            ':sdks:java:fake:fatJar', version='VERSION'))
    self.assertEqual(
        'https://repo.maven.apache.org/maven2/org/apache/beam/'
        'beam-sdks-java-fake/VERSION/beam-sdks-java-fake-A-VERSION.jar',
        subprocess_server.JavaJarServer.path_to_beam_jar(
            ':sdks:java:fake:fatJar', appendix='A', version='VERSION'))
    self.assertEqual(
        'https://repo.maven.apache.org/maven2/org/apache/beam/'
        'beam-sdks-java-fake/VERSION/beam-sdks-java-fake-A-VERSION.jar',
        subprocess_server.JavaJarServer.path_to_beam_jar(
            ':gradle:target:doesnt:matter',
            appendix='A',
            version='VERSION',
            artifact_id='beam-sdks-java-fake'))
    self.assertEqual(
        'https://repository.apache.org/content/groups/staging/org/apache/beam/'
        'beam-sdks-java-fake/2.99.9/beam-sdks-java-fake-2.99.9.jar',
        subprocess_server.JavaJarServer.path_to_beam_jar(
            ':sdks:java:fake:fatJar', version='2.99.9rc2'))

  def test_gradle_jar_dev(self):
    with self.assertRaisesRegex(
        Exception,
        re.escape(os.path.join('sdks',
                               'java',
                               'fake',
                               'build',
                               'libs',
                               'beam-sdks-java-fake-VERSION-SNAPSHOT.jar')) +
        ' not found.'):
      subprocess_server.JavaJarServer.path_to_beam_jar(
          ':sdks:java:fake:fatJar', version='VERSION.dev')
    with self.assertRaisesRegex(
        Exception,
        re.escape(os.path.join('sdks',
                               'java',
                               'fake',
                               'build',
                               'libs',
                               'beam-sdks-java-fake-A-VERSION-SNAPSHOT.jar')) +
        ' not found.'):
      subprocess_server.JavaJarServer.path_to_beam_jar(
          ':sdks:java:fake:fatJar', appendix='A', version='VERSION.dev')
    with self.assertRaisesRegex(
        Exception,
        re.escape(os.path.join('sdks',
                               'java',
                               'fake',
                               'build',
                               'libs',
                               'fake-artifact-id-A-VERSION-SNAPSHOT.jar')) +
        ' not found.'):
      subprocess_server.JavaJarServer.path_to_beam_jar(
          ':sdks:java:fake:fatJar',
          appendix='A',
          version='VERSION.dev',
          artifact_id='fake-artifact-id')

  def test_beam_services(self):
    with subprocess_server.JavaJarServer.beam_services({':some:target': 'foo'}):
      self.assertEqual(
          'foo',
          subprocess_server.JavaJarServer.path_to_beam_jar(':some:target'))

  def test_local_jar(self):
    class Handler(socketserver.BaseRequestHandler):
      timeout = 1

      def handle(self):
        data = self.request.recv(1024)
        if 'User-Agent: Apache Beam SDK for Python' in str(data):
          self.request.sendall(b'HTTP/1.1 200 OK\n\ndata')
        else:
          self.request.sendall(b'HTTP/1.1 400 BAD REQUEST\n\n')

    port, = subprocess_server.pick_port(None)
    server = socketserver.TCPServer(('localhost', port), Handler)
    t = threading.Thread(target=server.handle_request)
    t.daemon = True
    t.start()

    with tempfile.TemporaryDirectory() as temp_dir:
      subprocess_server.JavaJarServer.local_jar(
          'http://localhost:%s/path/to/file.jar' % port, temp_dir)
      with open(os.path.join(temp_dir, 'file.jar')) as fin:
        self.assertEqual(fin.read(), 'data')

  def test_local_jar_fallback_to_google_maven_mirror(self):
    """Test that Google Maven mirror is used as fallback 
    when Maven Central fails."""
    class MavenCentralHandler(socketserver.BaseRequestHandler):
      timeout = 1

      def handle(self):
        # Simulate Maven Central returning 403 Forbidden
        self.request.sendall(b'HTTP/1.1 403 Forbidden\n\n')

    # Set up Maven Central server (will return 403)
    maven_port, = subprocess_server.pick_port(None)
    maven_server = socketserver.TCPServer(('localhost', maven_port),
                                          MavenCentralHandler)
    maven_thread = threading.Thread(target=maven_server.handle_request)
    maven_thread.daemon = True
    maven_thread.start()

    # Temporarily replace the Maven Central constant to use our test server
    original_maven_central = (
        subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY)

    try:
      subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY = (
          f'http://localhost:{maven_port}/maven2')

      with tempfile.TemporaryDirectory() as temp_dir:
        # Use a Maven Central URL that will trigger the fallback to real
        # Google mirror
        maven_url = (
            f'http://localhost:{maven_port}/maven2/org/apache/beam/'
            f'beam-sdks-java-extensions-schemaio-expansion-service/2.63.0/'
            f'beam-sdks-java-extensions-schemaio-expansion-service-2.63.0.jar')

        # This should fail on our mock Maven Central and fallback to the
        # real Google mirror
        jar_path = subprocess_server.JavaJarServer.local_jar(
            maven_url, temp_dir)

        # Verify the file was downloaded successfully (from the real Google
        # mirror)
        self.assertTrue(os.path.exists(jar_path))
        jar_size = os.path.getsize(jar_path)
        self.assertTrue(jar_size > 0)  # Should have actual content

    finally:
      # Restore original constants
      subprocess_server.JavaJarServer.MAVEN_CENTRAL_REPOSITORY = (
          original_maven_central)

  @unittest.skipUnless(shutil.which('javac'), 'missing java jdk')
  def test_classpath_jar(self):
    with tempfile.TemporaryDirectory() as temp_dir:
      try:
        # Avoid having to prefix everything in our test strings.
        oldwd = os.getcwd()
        os.chdir(temp_dir)

        with open('Main.java', 'w') as fout:
          fout.write(
              """
public class Main {
  public static void main(String[] args) { Other.greet(); }
}
          """)

        with open('Other.java', 'w') as fout:
          fout.write(
              """
public class Other {
  public static void greet() { System.out.println("You got me!"); }
}
          """)

        os.mkdir('jars')
        # Using split just for readability/copyability.
        subprocess.check_call('javac Main.java Other.java'.split())
        subprocess.check_call('jar cfe jars/Main.jar Main Main.class'.split())
        subprocess.check_call('jar cf jars/Other.jar Other.class'.split())
        # Make sure the java and class files don't get picked up.
        for path in glob.glob('*.*'):
          os.unlink(path)

        # These should fail.
        self.assertNotEqual(
            subprocess.call('java -jar jars/Main.jar'.split()), 0)
        self.assertNotEqual(
            subprocess.call('java -jar jars/Other.jar'.split()), 0)

        os.mkdir('beam_temp')
        composite_jar = subprocess_server.JavaJarServer.make_classpath_jar(
            'jars/Main.jar', ['jars/Other.jar'], cache_dir='beam_temp')
        # This, however, should work.
        subprocess.check_call(f'java -jar {composite_jar}'.split())

      finally:
        os.chdir(oldwd)


class CacheTest(unittest.TestCase):
  @staticmethod
  def with_prefix(prefix):
    return '%s-%s' % (prefix, random.random())

  def test_memoization(self):
    cache = subprocess_server._SharedCache(self.with_prefix, lambda x: None)
    try:
      token = cache.register()
      a = cache.get('a')
      self.assertEqual(a[0], 'a')
      self.assertEqual(cache.get('a'), a)
      b = cache.get('b')
      self.assertEqual(b[0], 'b')
      self.assertEqual(cache.get('b'), b)
    finally:
      cache.purge(token)

  def test_purged(self):
    cache = subprocess_server._SharedCache(self.with_prefix, lambda x: None)
    try:
      token = cache.register()
      a = cache.get('a')
      self.assertEqual(cache.get('a'), a)
    finally:
      cache.purge(token)

    try:
      token = cache.register()
      new_a = cache.get('a')
      self.assertNotEqual(new_a, a)
    finally:
      cache.purge(token)

  def test_multiple_owners(self):
    cache = subprocess_server._SharedCache(self.with_prefix, lambda x: None)
    try:
      owner1 = cache.register()
      a = cache.get('a')
      try:
        self.assertEqual(cache.get('a'), a)
        owner2 = cache.register()
        b = cache.get('b')
        self.assertEqual(cache.get('b'), b)
      finally:
        cache.purge(owner2)
      self.assertEqual(cache.get('a'), a)
      self.assertEqual(cache.get('b'), b)
    finally:
      cache.purge(owner1)

    try:
      owner3 = cache.register()
      self.assertNotEqual(cache.get('a'), a)
      self.assertNotEqual(cache.get('b'), b)
    finally:
      cache.purge(owner3)

  def test_interleaved_owners(self):
    cache = subprocess_server._SharedCache(self.with_prefix, lambda x: None)
    owner1 = cache.register()
    a = cache.get('a')
    self.assertEqual(cache.get('a'), a)

    owner2 = cache.register()
    b = cache.get('b')
    self.assertEqual(cache.get('b'), b)

    cache.purge(owner1)
    self.assertNotEqual(cache.get('a'), a)
    self.assertEqual(cache.get('b'), b)

    cache.purge(owner2)
    owner3 = cache.register()
    self.assertNotEqual(cache.get('b'), b)
    cache.purge(owner3)


if __name__ == '__main__':
  unittest.main()