# -*- coding: utf-8 -*-
# Copyright 2021-2022 CERN
#
# Licensed 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.
#
# Authors:
# - Radu Carpa <radu.carpa@cern.ch>, 2021-2022

import pytest

import itertools
from datetime import datetime, timedelta
from random import randint
from unittest.mock import patch
from sqlalchemy import delete

from rucio.core import distance as distance_core
from rucio.core import request as request_core
from rucio.core import rse as rse_core
from rucio.core import replica as replica_core
from rucio.core import rule as rule_core
from rucio.core import config as core_config
from rucio.daemons.conveyor.submitter import submitter
from rucio.daemons.reaper.reaper import reaper
from rucio.db.sqla.models import Request, Source
from rucio.db.sqla.constants import RequestState
from rucio.db.sqla.session import read_session, transactional_session


@pytest.mark.noparallel(reason="multiple submitters cannot be run in parallel due to partial job assignment by hash")
@pytest.mark.parametrize("file_config_mock", [
    # Run test twice: with, and without, temp tables
    {
        "overrides": [
            ('core', 'use_temp_tables', 'True'),
        ]
    },
    {
        "overrides": [
            ('core', 'use_temp_tables', 'False'),
        ]
    }
], indirect=True)
def test_request_submitted_in_order(rse_factory, did_factory, root_account, file_config_mock):

    src_rses = [rse_factory.make_posix_rse() for _ in range(2)]
    dst_rses = [rse_factory.make_posix_rse() for _ in range(3)]
    for _, src_rse_id in src_rses:
        for _, dst_rse_id in dst_rses:
            distance_core.add_distance(src_rse_id=src_rse_id, dest_rse_id=dst_rse_id, ranking=10)
            distance_core.add_distance(src_rse_id=dst_rse_id, dest_rse_id=src_rse_id, ranking=10)

    # Create a certain number of files on source RSEs with replication rules towards random destination RSEs
    nb_files = 15
    dids = []
    requests = []
    src_rses_iterator = itertools.cycle(src_rses)
    dst_rses_iterator = itertools.cycle(dst_rses)
    for _ in range(nb_files):
        src_rse_name, src_rse_id = next(src_rses_iterator)
        dst_rse_name, dst_rse_id = next(dst_rses_iterator)
        did = did_factory.upload_test_file(rse_name=src_rse_name)
        rule_core.add_rule(dids=[did], account=root_account, copies=1, rse_expression=dst_rse_name, grouping='ALL', weight=None, lifetime=None, locked=False, subscription_id=None)
        requests.append(request_core.get_request_by_did(rse_id=dst_rse_id, **did))
        dids.append(did)

    # Forge request creation time to a random moment in the past hour
    @transactional_session
    def _forge_requests_creation_time(session=None):
        base_time = datetime.utcnow().replace(microsecond=0, minute=0) - timedelta(hours=1)
        assigned_times = set()
        for request in requests:
            request_creation_time = None
            while not request_creation_time or request_creation_time in assigned_times:
                # Ensure uniqueness to avoid multiple valid submission orders and make tests deterministic with simple sorting techniques
                request_creation_time = base_time + timedelta(minutes=randint(0, 3600))
            assigned_times.add(request_creation_time)
            session.query(Request).filter(Request.id == request['id']).update({'created_at': request_creation_time})
            request['created_at'] = request_creation_time

    _forge_requests_creation_time()
    requests = sorted(requests, key=lambda r: r['created_at'])

    for request in requests:
        assert request_core.get_request(request_id=request['id'])['state'] == RequestState.QUEUED

    requests_id_in_submission_order = []
    with patch('rucio.transfertool.mock.MockTransfertool.submit') as mock_transfertool_submit:
        # Record the order of requests passed to MockTranfertool.submit()
        mock_transfertool_submit.side_effect = lambda transfers, job_params, timeout: requests_id_in_submission_order.extend([t.rws.request_id for t in transfers])

        submitter(once=True, rses=[{'id': rse_id} for _, rse_id in dst_rses], partition_wait_time=None, transfertool='mock', transfertype='single', filter_transfertool=None)

    for request in requests:
        assert request_core.get_request(request_id=request['id'])['state'] == RequestState.SUBMITTED

    # Requests must be submitted in the order of their creation
    assert requests_id_in_submission_order == [r['id'] for r in requests]


@pytest.mark.noparallel(reason="multiple submitters cannot be run in parallel due to partial job assignment by hash")
@pytest.mark.parametrize("core_config_mock", [
    # Run test twice: with, and without, temp tables
    {
        "table_content": [
            ('transfers', 'use_multihop', True),
            ('core', 'use_temp_tables', False),
        ]
    },
    {
        "table_content": [
            ('transfers', 'use_multihop', True),
            ('core', 'use_temp_tables', True),
        ]
    }
], indirect=True)
@pytest.mark.parametrize("caches_mock", [{"caches_to_mock": [
    'rucio.core.rse_expression_parser.REGION',  # The list of multihop RSEs is retrieved by rse expression
    'rucio.core.config.REGION',
]}], indirect=True)
def test_multihop_sources_created(rse_factory, did_factory, root_account, core_config_mock, caches_mock, metrics_mock):
    """
    Ensure that multihop transfers are handled and intermediate request correctly created
    """
    src_rse_name, src_rse_id = rse_factory.make_posix_rse()
    _, jump_rse1_id = rse_factory.make_posix_rse()
    _, jump_rse2_id = rse_factory.make_posix_rse()
    _, jump_rse3_id = rse_factory.make_posix_rse()
    dst_rse_name, dst_rse_id = rse_factory.make_posix_rse()

    jump_rses = [jump_rse1_id, jump_rse2_id, jump_rse3_id]
    all_rses = jump_rses + [src_rse_id, dst_rse_id]

    for rse_id in jump_rses:
        rse_core.add_rse_attribute(rse_id, 'available_for_multihop', True)

    rse_tombstone_delay = 3600
    rse_multihop_tombstone_delay = 12 * 3600
    default_multihop_tombstone_delay = 24 * 3600

    # if both attributes are set, the multihop one will take precedence
    rse_core.add_rse_attribute(jump_rse1_id, 'tombstone_delay', rse_tombstone_delay)
    rse_core.add_rse_attribute(jump_rse1_id, 'multihop_tombstone_delay', rse_multihop_tombstone_delay)

    # if multihop delay not set, it's the default multihop takes precedence. Not normal tombstone delay.
    rse_core.add_rse_attribute(jump_rse2_id, 'tombstone_delay', rse_tombstone_delay)
    core_config.set(section='transfers', option='multihop_tombstone_delay', value=default_multihop_tombstone_delay)

    # if multihop delay is set to 0, the replica will have no tombstone
    rse_core.add_rse_attribute(jump_rse3_id, 'multihop_tombstone_delay', 0)

    distance_core.add_distance(src_rse_id, jump_rse1_id, ranking=10)
    distance_core.add_distance(jump_rse1_id, jump_rse2_id, ranking=10)
    distance_core.add_distance(jump_rse2_id, jump_rse3_id, ranking=10)
    distance_core.add_distance(jump_rse3_id, dst_rse_id, ranking=10)

    did = did_factory.upload_test_file(src_rse_name)
    rule_core.add_rule(dids=[did], account=root_account, copies=1, rse_expression=dst_rse_name, grouping='ALL', weight=None, lifetime=None, locked=False, subscription_id=None)

    submitter(once=True, rses=[{'id': rse_id} for rse_id in all_rses], partition_wait_time=None, transfertool='mock', transfertype='single', filter_transfertool=None)

    # Ensure that each intermediate request was correctly created
    for rse_id in jump_rses:
        assert request_core.get_request_by_did(rse_id=rse_id, **did)

    @read_session
    def __number_sources(rse_id, scope, name, session=None):
        return session.query(Source). \
            filter(Source.rse_id == rse_id). \
            filter(Source.scope == scope). \
            filter(Source.name == name). \
            count()

    # Ensure that sources where created for transfers
    for rse_id in [src_rse_id, jump_rse1_id, jump_rse2_id]:
        assert __number_sources(rse_id, **did) == 2
    assert __number_sources(jump_rse3_id, **did) == 1

    # Ensure the tombstone is correctly set on intermediate replicas
    expected_tombstone = datetime.utcnow() + timedelta(seconds=rse_multihop_tombstone_delay)
    replica = replica_core.get_replica(jump_rse1_id, **did)
    assert expected_tombstone - timedelta(minutes=5) < replica['tombstone'] < expected_tombstone + timedelta(minutes=5)

    expected_tombstone = datetime.utcnow() + timedelta(seconds=default_multihop_tombstone_delay)
    replica = replica_core.get_replica(jump_rse2_id, **did)
    assert expected_tombstone - timedelta(minutes=5) < replica['tombstone'] < expected_tombstone + timedelta(minutes=5)

    replica = replica_core.get_replica(jump_rse3_id, **did)
    assert replica['tombstone'] is None

    # Ensure that prometheus metrics were correctly registered. Only one submission, mock transfertool groups everything into one job.
    assert metrics_mock.get_sample_value('rucio_core_request_submit_transfer_total') == 1


@pytest.mark.noparallel(reason="multiple submitters cannot be run in parallel due to partial job assignment by hash")
@pytest.mark.parametrize("core_config_mock", [{"table_content": [
    ('transfers', 'use_multihop', True),
]}], indirect=True)
@pytest.mark.parametrize("caches_mock", [{"caches_to_mock": [
    'rucio.core.config.REGION',
    'rucio.daemons.reaper.reaper.REGION',
]}], indirect=True)
def test_source_avoid_deletion(vo, caches_mock, core_config_mock, rse_factory, did_factory, root_account, file_factory):
    """ Test that sources on a file block it from deletion """

    _, reaper_region = caches_mock
    src_rse1, src_rse1_id = rse_factory.make_mock_rse()
    src_rse2, src_rse2_id = rse_factory.make_mock_rse()
    dst_rse, dst_rse_id = rse_factory.make_mock_rse()
    all_rses = [src_rse1_id, src_rse2_id, dst_rse_id]
    any_source = f'{src_rse1}|{src_rse2}'

    for rse_id in [src_rse1_id, src_rse2_id]:
        rse_core.set_rse_limits(rse_id=rse_id, name='MinFreeSpace', value=1)
        rse_core.set_rse_usage(rse_id=rse_id, source='storage', used=1, free=0)
    distance_core.add_distance(src_rse1_id, dst_rse_id, ranking=20)
    distance_core.add_distance(src_rse2_id, dst_rse_id, ranking=10)

    # Upload a test file to both rses without registering
    did = did_factory.random_did()

    # Register replica on one source RSE
    replica_core.add_replica(rse_id=src_rse1_id, account=root_account, bytes_=1, tombstone=datetime(year=1970, month=1, day=1), **did)
    rule_core.add_rule(dids=[did], account=root_account, copies=1, rse_expression=dst_rse, grouping='ALL', weight=None, lifetime=None, locked=False, subscription_id=None)

    # Reaper will not delete a file which only has one replica if there is any pending transfer for it
    reaper_region.invalidate()
    reaper(once=True, rses=[], include_rses=any_source, exclude_rses=None)
    replica = next(iter(replica_core.list_replicas(dids=[did], rse_expression=any_source)))
    assert len(replica['pfns']) == 1

    # Register replica on second source rse
    replica_core.add_replica(rse_id=src_rse2_id, account=root_account, bytes_=1, tombstone=datetime(year=1970, month=1, day=1), **did)
    replica = next(iter(replica_core.list_replicas(dids=[did], rse_expression=any_source)))
    assert len(replica['pfns']) == 2

    # Submit the transfer. This will create the sources.
    submitter(once=True, rses=[{'id': rse_id} for rse_id in all_rses], partition_wait_time=None, transfertool='mock', transfertype='single', filter_transfertool=None)

    # None of the replicas will be removed. They are protected by an entry in the sources table
    reaper_region.invalidate()
    reaper(once=True, rses=[], include_rses=any_source, exclude_rses=None)
    replica = next(iter(replica_core.list_replicas(dids=[did], rse_expression=any_source)))
    assert len(replica['pfns']) == 2

    @transactional_session
    def __delete_sources(rse_id, scope, name, session=None):
        session.execute(
            delete(Source).where(Source.rse_id == rse_id,
                                 Source.scope == scope,
                                 Source.name == name))

    # Deletion succeeds for one replica (second still protected by existing request)
    __delete_sources(src_rse1_id, **did)
    __delete_sources(src_rse2_id, **did)
    reaper_region.invalidate()
    reaper(once=True, rses=[], include_rses=any_source, exclude_rses=None)
    replica = next(iter(replica_core.list_replicas(dids=[did], rse_expression=any_source)))
    assert len(replica['pfns']) == 1


@pytest.mark.noparallel(reason="multiple submitters cannot be run in parallel due to partial job assignment by hash")
@pytest.mark.parametrize("core_config_mock", [{"table_content": [
    ('transfers', 'use_multihop', True)
]}], indirect=True)
@pytest.mark.parametrize("caches_mock", [{"caches_to_mock": [
    'rucio.core.rse_expression_parser.REGION',  # The list of multihop RSEs is retrieved by rse expression
    'rucio.core.transfer.REGION_SHORT',
]}], indirect=True)
def test_ignore_availability(rse_factory, did_factory, root_account, core_config_mock, caches_mock):

    def __setup_test():
        src_rse, src_rse_id = rse_factory.make_posix_rse()
        dst_rse, dst_rse_id = rse_factory.make_posix_rse()

        distance_core.add_distance(src_rse_id, dst_rse_id, ranking=10)
        did = did_factory.upload_test_file(src_rse)
        rule_core.add_rule(dids=[did], account=root_account, copies=1, rse_expression=dst_rse, grouping='ALL', weight=None, lifetime=None, locked=False, subscription_id=None)

        rse_core.update_rse(src_rse_id, {'availability_read': False})

        return src_rse_id, dst_rse_id, did

    src_rse_id, dst_rse_id, did = __setup_test()
    submitter(once=True, rses=[{'id': rse_id} for rse_id in (src_rse_id, dst_rse_id)], partition_wait_time=None, transfertool='mock', transfertype='single')
    request = request_core.get_request_by_did(rse_id=dst_rse_id, **did)
    assert request['state'] == RequestState.NO_SOURCES

    src_rse_id, dst_rse_id, did = __setup_test()
    submitter(once=True, rses=[{'id': rse_id} for rse_id in (src_rse_id, dst_rse_id)], partition_wait_time=None, transfertool='mock', transfertype='single', ignore_availability=True)
    request = request_core.get_request_by_did(rse_id=dst_rse_id, **did)
    assert request['state'] == RequestState.SUBMITTED


@pytest.mark.noparallel(reason="multiple submitters cannot be run in parallel due to partial job assignment by hash")
def test_globus(rse_factory, did_factory, root_account):
    """
    Test bulk submissions with globus transfertool.
    Rely on mocks, because we don't contact a real globus server in tests
    """
    # +------+    +------+
    # |      |    |      |
    # | RSE1 +--->| RSE2 |
    # |      |    |      |
    # +------+    +------+
    #
    # +------+    +------+
    # |      |    |      |
    # | RSE3 +--->| RSE4 |
    # |      |    |      |
    # +------+    +------+
    rse1, rse1_id = rse_factory.make_posix_rse()
    rse2, rse2_id = rse_factory.make_posix_rse()
    rse3, rse3_id = rse_factory.make_posix_rse()
    rse4, rse4_id = rse_factory.make_posix_rse()
    all_rses = [rse1_id, rse2_id, rse3_id, rse4_id]

    distance_core.add_distance(rse1_id, rse2_id, ranking=10)
    distance_core.add_distance(rse3_id, rse4_id, ranking=10)
    for rse_id in all_rses:
        rse_core.add_rse_attribute(rse_id, 'globus_endpoint_id', rse_id)

    # Single submission
    did1 = did_factory.upload_test_file(rse1)
    rule_core.add_rule(dids=[did1], account=root_account, copies=1, rse_expression=rse2, grouping='ALL', weight=None, lifetime=None, locked=False, subscription_id=None)
    did2 = did_factory.upload_test_file(rse3)
    rule_core.add_rule(dids=[did2], account=root_account, copies=1, rse_expression=rse4, grouping='ALL', weight=None, lifetime=None, locked=False, subscription_id=None)
    with patch('rucio.transfertool.globus.bulk_submit_xfer') as mock_bulk_submit:
        mock_bulk_submit.return_value = 0
        submitter(once=True, rses=[{'id': rse_id} for rse_id in all_rses], group_bulk=10, partition_wait_time=None, transfertool='globus', transfertype='single', filter_transfertool=None)
        # Called separately for each job
        assert len(mock_bulk_submit.call_args_list) == 2
        (submitjob,), _kwargs = mock_bulk_submit.call_args_list[0]
        assert len(submitjob) == 1

    # Bulk submission
    did1 = did_factory.upload_test_file(rse1)
    rule_core.add_rule(dids=[did1], account=root_account, copies=1, rse_expression=rse2, grouping='ALL', weight=None, lifetime=None, locked=False, subscription_id=None)
    did2 = did_factory.upload_test_file(rse3)
    rule_core.add_rule(dids=[did2], account=root_account, copies=1, rse_expression=rse4, grouping='ALL', weight=None, lifetime=None, locked=False, subscription_id=None)
    with patch('rucio.transfertool.globus.bulk_submit_xfer') as mock_bulk_submit:
        mock_bulk_submit.return_value = 0
        submitter(once=True, rses=[{'id': rse_id} for rse_id in all_rses], group_bulk=10, partition_wait_time=None, transfertool='globus', transfertype='bulk', filter_transfertool=None)

        mock_bulk_submit.assert_called_once()
        (submitjob,), _kwargs = mock_bulk_submit.call_args_list[0]

        # both jobs were grouped together and submitted in one call
        assert len(submitjob) == 2

        job_did1 = next(iter(filter(lambda job: did1['name'] in job['sources'][0], submitjob)))
        assert len(job_did1['sources']) == 1
        assert len(job_did1['destinations']) == 1
        assert job_did1['metadata']['src_rse'] == rse1
        assert job_did1['metadata']['dst_rse'] == rse2
        assert job_did1['metadata']['name'] == did1['name']
        assert job_did1['metadata']['source_globus_endpoint_id'] == rse1_id
        assert job_did1['metadata']['dest_globus_endpoint_id'] == rse2_id

        job_did2 = next(iter(filter(lambda job: did2['name'] in job['sources'][0], submitjob)))
        assert len(job_did2['sources']) == 1
        assert len(job_did2['destinations']) == 1
        assert job_did2['metadata']['src_rse'] == rse3
        assert job_did2['metadata']['dst_rse'] == rse4
        assert job_did2['metadata']['name'] == did2['name']
    request = request_core.get_request_by_did(rse_id=rse2_id, **did1)
    assert request['state'] == RequestState.SUBMITTED
    request = request_core.get_request_by_did(rse_id=rse4_id, **did2)
    assert request['state'] == RequestState.SUBMITTED
