# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# Copyright 2025 The TransferQueue Team
#
# 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.

import logging
import sys
from pathlib import Path

import pytest
import ray
import torch

parent_dir = Path(__file__).resolve().parent.parent
sys.path.append(str(parent_dir))

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

from transfer_queue import TransferQueueController  # noqa: E402
from transfer_queue.utils.utils import ProductionStatus  # noqa: E402


@pytest.fixture(scope="function")
def ray_setup():
    if ray.is_initialized():
        ray.shutdown()
    ray.init(
        ignore_reinit_error=True,
        runtime_env={"env_vars": {"RAY_DEBUG": "1", "RAY_DEDUP_LOGS": "0"}},
        log_to_driver=True,
    )
    yield
    if ray.is_initialized():
        ray.shutdown()
        logger.info("Ray has been shut down completely after test")


class TestTransferQueueController:
    def test_controller_with_single_partition(self, ray_setup):
        gbs = 8
        num_n_samples = 4

        tq_controller = TransferQueueController.remote()

        # Test get metadata in insert mode
        partition_id = "train_0"
        data_fields = ["prompt_ids", "attention_mask"]
        metadata = ray.get(
            tq_controller.get_metadata.remote(
                data_fields=data_fields,
                batch_size=gbs * num_n_samples,
                partition_id=partition_id,
                mode="insert",
            )
        )

        assert metadata.global_indexes == list(range(gbs * num_n_samples))
        assert metadata.samples[0].partition_id == "train_0"
        assert sum([int(sample.fields.get("prompt_ids").production_status) for sample in metadata.samples]) == int(
            ProductionStatus.NOT_PRODUCED
        )
        assert sum([int(sample.fields.get("attention_mask").production_status) for sample in metadata.samples]) == int(
            ProductionStatus.NOT_PRODUCED
        )
        partition_index_range = ray.get(tq_controller.get_partition_index_range.remote(partition_id))
        assert partition_index_range == set(range(gbs * num_n_samples))

        print(" Initial get metadata correct")

        # Test update production status
        dtypes = {k: {"prompt_ids": "torch.int64", "attention_mask": "torch.bool"} for k in metadata.global_indexes}
        shapes = {k: {"prompt_ids": (32,), "attention_mask": (32,)} for k in metadata.global_indexes}
        success = ray.get(
            tq_controller.update_production_status.remote(
                partition_id=partition_id,
                global_indexes=metadata.global_indexes,
                field_names=metadata.field_names,
                dtypes=dtypes,
                shapes=shapes,
            )
        )
        assert success
        partition = ray.get(tq_controller.get_partition_snapshot.remote(partition_id))
        assert partition.production_status is not None
        assert partition.production_status.size(0) == gbs * num_n_samples

        # Test for get production status
        production_status = ray.get(
            tq_controller.get_production_status.remote(
                partition_id=partition_id,
                data_fields=data_fields,
            )
        )
        assert production_status

        # Total fields should match the number of fields we added
        assert partition.total_fields_num == len(data_fields)

        # Allocated fields should be at least the number of actual fields
        assert partition.allocated_fields_num >= partition.total_fields_num

        # Check production status for the fields we added
        assert torch.equal(
            sum(partition.production_status[:, : len(data_fields)]),
            torch.Tensor([gbs * num_n_samples, gbs * num_n_samples]),
        )

        # Any additional allocated fields should be zero (unused)
        if partition.allocated_fields_num > len(data_fields):
            assert torch.equal(
                sum(partition.production_status[:, len(data_fields) :]),
                torch.zeros(1 * (partition.allocated_fields_num - len(data_fields))),
            )

        print(f" Updated production status for partition {partition_id}")

        # Test for get consumption status
        consumption_status = ray.get(
            tq_controller.get_consumption_status.remote(
                partition_id=partition_id,
                task_name="generate_sequences",
            )
        )
        assert torch.equal(consumption_status, torch.zeros(gbs * num_n_samples))

        # Test get metadate in fetch mode
        gen_meta = ray.get(
            tq_controller.get_metadata.remote(
                data_fields=["prompt_ids"],
                batch_size=gbs * num_n_samples,
                partition_id=partition_id,
                mode="fetch",
                task_name="generate_sequences",
            )
        )

        assert gen_meta.global_indexes == list(range(gbs * num_n_samples))
        assert gen_meta.samples[0].partition_id == "train_0"
        assert gen_meta.field_names == ["prompt_ids"]
        partition = ray.get(tq_controller.get_partition_snapshot.remote(partition_id))
        assert torch.equal(partition.consumption_status["generate_sequences"], torch.ones(gbs * num_n_samples))
        print(" Get metadata in fetch mode correct")

        # Test for get consumption status
        consumption_status = ray.get(
            tq_controller.get_consumption_status.remote(
                partition_id=partition_id,
                task_name="generate_sequences",
            )
        )
        assert torch.equal(consumption_status, torch.ones(gbs * num_n_samples))

        # Test get clear meta
        clear_meta = ray.get(
            tq_controller.get_metadata.remote(
                data_fields=[],
                partition_id=partition_id,
                mode="insert",
            )
        )
        assert clear_meta.global_indexes == list(range(gbs * num_n_samples))
        assert [sample.fields for sample in clear_meta.samples] == [{}] * (gbs * num_n_samples)
        print(" Clear metadata correct")

        # Test clear_partition
        ray.get(tq_controller.clear_partition.remote(partition_id))
        partition = ray.get(tq_controller.get_partition_snapshot.remote(partition_id))
        partition_index_range = ray.get(tq_controller.get_partition_index_range.remote(partition_id))
        assert partition_index_range == set()
        assert partition is None
        print(" Clear partition correct")

    def test_controller_with_multi_partitions(self, ray_setup):
        gbs_1 = 8
        num_n_samples_1 = 4
        partition_id_1 = "train_0"

        gbs_2 = 16
        num_n_samples_2 = 1
        partition_id_2 = "val_0"

        gbs_3 = 32
        num_n_samples_3 = 2
        partition_id_3 = "train_1"

        tq_controller = TransferQueueController.remote()

        # Test get metadata in insert mode
        data_fields = ["prompt_ids", "attention_mask"]
        metadata = ray.get(
            tq_controller.get_metadata.remote(
                data_fields=data_fields,
                batch_size=gbs_1 * num_n_samples_1,
                partition_id=partition_id_1,
                mode="insert",
            )
        )

        # Test update production status
        dtypes = {k: {"prompt_ids": "torch.int64", "attention_mask": "torch.bool"} for k in metadata.global_indexes}
        shapes = {k: {"prompt_ids": (32,), "attention_mask": (32,)} for k in metadata.global_indexes}
        success = ray.get(
            tq_controller.update_production_status.remote(
                partition_id=partition_id_1,
                global_indexes=metadata.global_indexes,
                field_names=metadata.field_names,
                dtypes=dtypes,
                shapes=shapes,
            )
        )
        assert success

        # Test get metadate in fetch mode
        gen_meta = ray.get(
            tq_controller.get_metadata.remote(
                data_fields=["prompt_ids"],
                batch_size=gbs_1 * num_n_samples_1,
                partition_id=partition_id_1,
                mode="fetch",
                task_name="generate_sequences",
            )
        )
        assert gen_meta

        # Test get clear meta
        clear_meta = ray.get(
            tq_controller.get_metadata.remote(
                data_fields=[],
                partition_id=partition_id_1,
                mode="insert",
            )
        )
        assert clear_meta

        # =========================partition 2=============================#
        data_fields = ["prompt_ids", "attention_mask"]
        val_metadata = ray.get(
            tq_controller.get_metadata.remote(
                data_fields=data_fields,
                batch_size=gbs_2 * num_n_samples_2,
                partition_id=partition_id_2,
                mode="insert",
            )
        )

        part1_index_range = gbs_1 * num_n_samples_1
        part2_index_range = gbs_2 * num_n_samples_2
        assert val_metadata.global_indexes == list(range(part1_index_range, part2_index_range + part1_index_range))
        assert val_metadata.samples[0].partition_id == "val_0"
        assert sum([int(sample.fields.get("prompt_ids").production_status) for sample in val_metadata.samples]) == int(
            ProductionStatus.NOT_PRODUCED
        )
        assert sum(
            [int(sample.fields.get("attention_mask").production_status) for sample in val_metadata.samples]
        ) == int(ProductionStatus.NOT_PRODUCED)
        partition_index_range = ray.get(tq_controller.get_partition_index_range.remote(partition_id_2))
        assert partition_index_range == set(range(part1_index_range, part2_index_range + part1_index_range))

        # Update production status
        dtypes = {k: {"prompt_ids": "torch.int64", "attention_mask": "torch.bool"} for k in val_metadata.global_indexes}
        shapes = {k: {"prompt_ids": (32,), "attention_mask": (32,)} for k in val_metadata.global_indexes}
        success = ray.get(
            tq_controller.update_production_status.remote(
                partition_id=partition_id_2,
                global_indexes=val_metadata.global_indexes,
                field_names=val_metadata.field_names,
                dtypes=dtypes,
                shapes=shapes,
            )
        )
        assert success

        # Clear partition 1
        partition_index_range_1 = ray.get(tq_controller.get_partition_index_range.remote(partition_id_1))
        assert partition_index_range_1
        ray.get(tq_controller.clear_partition.remote(partition_id_1))
        partition_1_after_clear = ray.get(tq_controller.get_partition_snapshot.remote(partition_id_1))
        partition_index_range_1_after_clear = ray.get(tq_controller.get_partition_index_range.remote(partition_id_1))

        assert not partition_index_range_1_after_clear
        assert partition_1_after_clear is None
        assert partition_index_range_1_after_clear == set()

        partition_2 = ray.get(tq_controller.get_partition_snapshot.remote(partition_id_2))
        partition_index_range_2 = ray.get(tq_controller.get_partition_index_range.remote(partition_id_2))
        assert partition_index_range_2 == set([32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47])
        assert torch.all(
            partition_2.production_status[list(partition_index_range_2), : len(val_metadata.field_names)] == 1
        )
        print(" Only clear partition 1 correct")

        # =========================partition 3=============================#
        metadata_2 = ray.get(
            tq_controller.get_metadata.remote(
                data_fields=data_fields,
                batch_size=gbs_3 * num_n_samples_3,
                partition_id=partition_id_3,
                mode="insert",
            )
        )
        assert metadata_2.global_indexes == list(range(32)) + list(range(48, 80))
        assert metadata_2.samples[0].partition_id == "train_1"
        assert sum([int(sample.fields.get("prompt_ids").production_status) for sample in metadata_2.samples]) == int(
            ProductionStatus.NOT_PRODUCED
        )
        assert sum(
            [int(sample.fields.get("attention_mask").production_status) for sample in metadata_2.samples]
        ) == int(ProductionStatus.NOT_PRODUCED)
        partition_index_range = ray.get(tq_controller.get_partition_index_range.remote(partition_id_3))
        assert partition_index_range == set(list(range(32)) + list(range(48, 80)))
        print(" Correctly assign partition_3")

    def test_controller_clear_meta(self, ray_setup):
        """Test clear_meta functionality for individual samples"""
        gbs = 4
        num_n_samples = 2
        partition_id = "test_clear_meta"

        tq_controller = TransferQueueController.remote()

        # Create metadata in insert mode
        data_fields = ["prompt_ids", "attention_mask"]
        metadata = ray.get(
            tq_controller.get_metadata.remote(
                data_fields=data_fields,
                batch_size=gbs * num_n_samples,
                partition_id=partition_id,
                mode="insert",
            )
        )

        assert metadata.global_indexes == list(range(gbs * num_n_samples))

        # Update production status
        dtypes = {k: {"prompt_ids": "torch.int64", "attention_mask": "torch.bool"} for k in metadata.global_indexes}
        shapes = {k: {"prompt_ids": (32,), "attention_mask": (32,)} for k in metadata.global_indexes}
        success = ray.get(
            tq_controller.update_production_status.remote(
                partition_id=partition_id,
                global_indexes=metadata.global_indexes,
                field_names=metadata.field_names,
                dtypes=dtypes,
                shapes=shapes,
            )
        )
        assert success

        # Get partition snapshot before clear
        partition_before = ray.get(tq_controller.get_partition_snapshot.remote(partition_id))
        assert partition_before is not None
        assert len(partition_before.global_indexes) == gbs * num_n_samples
        assert set(partition_before.global_indexes) == set(range(gbs * num_n_samples))

        # Test clear_meta - clear first 4 samples (indexes 0-3)
        global_indexes_to_clear = [0, 1, 2, 3, 6]
        partition_ids_to_clear = [partition_id] * len(global_indexes_to_clear)

        ray.get(
            tq_controller.clear_meta.remote(
                global_indexes=global_indexes_to_clear,
                partition_ids=partition_ids_to_clear,
            )
        )

        # Check that only the cleared samples are affected
        partition_after = ray.get(tq_controller.get_partition_snapshot.remote(partition_id))
        assert partition_after is not None

        # Verify production status is cleared for the specified indexes
        assert set(partition_after.global_indexes) == set([4, 5, 7])

        print(" Clear meta correct")