"""
Comprehensive integration tests for DatasetService retrieval/list methods.

This test suite covers:
- get_datasets - pagination, search, filtering, permissions
- get_dataset - single dataset retrieval
- get_datasets_by_ids - bulk retrieval
- get_process_rules - dataset processing rules
- get_dataset_queries - dataset query history
- get_related_apps - apps using the dataset
"""

import json
from uuid import uuid4

from sqlalchemy.orm import Session

from core.rag.index_processor.constant.index_type import IndexTechniqueType
from models import AccountStatus, CreatorUserRole, TenantStatus
from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import (
    AppDatasetJoin,
    Dataset,
    DatasetPermission,
    DatasetPermissionEnum,
    DatasetProcessRule,
    DatasetQuery,
)
from models.enums import DatasetQuerySource, DataSourceType, ProcessRuleMode, TagType
from models.model import Tag, TagBinding
from services.dataset_service import DatasetService, DocumentService


class DatasetRetrievalTestDataFactory:
    """Factory class for creating database-backed test data for dataset retrieval integration tests."""

    @staticmethod
    def create_account_with_tenant(
        db_session_with_containers: Session, role: TenantAccountRole = TenantAccountRole.NORMAL
    ) -> tuple[Account, Tenant]:
        """Create an account and tenant with the specified role."""
        account = Account(
            email=f"{uuid4()}@example.com",
            name=f"user-{uuid4()}",
            interface_language="en-US",
            status=AccountStatus.ACTIVE,
        )
        tenant = Tenant(
            name=f"tenant-{uuid4()}",
            status=TenantStatus.NORMAL,
        )
        db_session_with_containers.add_all([account, tenant])
        db_session_with_containers.flush()

        join = TenantAccountJoin(
            tenant_id=tenant.id,
            account_id=account.id,
            role=role,
            current=True,
        )
        db_session_with_containers.add(join)
        db_session_with_containers.commit()

        account.current_tenant = tenant
        return account, tenant

    @staticmethod
    def create_account_in_tenant(
        db_session_with_containers: Session, tenant: Tenant, role: TenantAccountRole = TenantAccountRole.OWNER
    ) -> Account:
        """Create an account and add it to an existing tenant."""
        account = Account(
            email=f"{uuid4()}@example.com",
            name=f"user-{uuid4()}",
            interface_language="en-US",
            status=AccountStatus.ACTIVE,
        )
        db_session_with_containers.add(account)
        db_session_with_containers.flush()

        join = TenantAccountJoin(
            tenant_id=tenant.id,
            account_id=account.id,
            role=role,
            current=True,
        )
        db_session_with_containers.add(join)
        db_session_with_containers.commit()

        account.current_tenant = tenant
        return account

    @staticmethod
    def create_dataset(
        db_session_with_containers: Session,
        tenant_id: str,
        created_by: str,
        name: str = "Test Dataset",
        permission: DatasetPermissionEnum = DatasetPermissionEnum.ONLY_ME,
    ) -> Dataset:
        """Create a dataset."""
        dataset = Dataset(
            tenant_id=tenant_id,
            name=name,
            description="desc",
            data_source_type=DataSourceType.UPLOAD_FILE,
            indexing_technique=IndexTechniqueType.HIGH_QUALITY,
            created_by=created_by,
            permission=permission,
            provider="vendor",
            retrieval_model={"top_k": 2},
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.commit()
        return dataset

    @staticmethod
    def create_dataset_permission(
        db_session_with_containers: Session, dataset_id: str, tenant_id: str, account_id: str
    ) -> DatasetPermission:
        """Create a dataset permission."""
        permission = DatasetPermission(
            dataset_id=dataset_id,
            tenant_id=tenant_id,
            account_id=account_id,
            has_permission=True,
        )
        db_session_with_containers.add(permission)
        db_session_with_containers.commit()
        return permission

    @staticmethod
    def create_process_rule(
        db_session_with_containers: Session, dataset_id: str, created_by: str, mode: ProcessRuleMode, rules: dict
    ) -> DatasetProcessRule:
        """Create a dataset process rule."""
        process_rule = DatasetProcessRule(
            dataset_id=dataset_id,
            created_by=created_by,
            mode=mode,
            rules=json.dumps(rules),
        )
        db_session_with_containers.add(process_rule)
        db_session_with_containers.commit()
        return process_rule

    @staticmethod
    def create_dataset_query(
        db_session_with_containers: Session, dataset_id: str, created_by: str, content: str
    ) -> DatasetQuery:
        """Create a dataset query."""
        dataset_query = DatasetQuery(
            dataset_id=dataset_id,
            content=content,
            source=DatasetQuerySource.APP,
            source_app_id=None,
            created_by_role=CreatorUserRole.ACCOUNT,
            created_by=created_by,
        )
        db_session_with_containers.add(dataset_query)
        db_session_with_containers.commit()
        return dataset_query

    @staticmethod
    def create_app_dataset_join(db_session_with_containers: Session, dataset_id: str) -> AppDatasetJoin:
        """Create an app-dataset join."""
        join = AppDatasetJoin(
            app_id=str(uuid4()),
            dataset_id=dataset_id,
        )
        db_session_with_containers.add(join)
        db_session_with_containers.commit()
        return join

    @staticmethod
    def create_tag_binding(db_session_with_containers: Session, tenant_id: str, created_by: str, target_id: str) -> Tag:
        """Create a knowledge tag and bind it to the target dataset."""
        tag = Tag(
            tenant_id=tenant_id,
            type=TagType.KNOWLEDGE,
            name=f"tag-{uuid4()}",
            created_by=created_by,
        )
        db_session_with_containers.add(tag)
        db_session_with_containers.flush()

        binding = TagBinding(
            tenant_id=tenant_id,
            tag_id=tag.id,
            target_id=target_id,
            created_by=created_by,
        )
        db_session_with_containers.add(binding)
        db_session_with_containers.commit()
        return tag


class TestDatasetServiceGetDatasets:
    """
    Comprehensive integration tests for DatasetService.get_datasets method.

    This test suite covers:
    - Pagination
    - Search functionality
    - Tag filtering
    - Permission-based filtering (ONLY_ME, ALL_TEAM, PARTIAL_TEAM)
    - Role-based filtering (OWNER, DATASET_OPERATOR, NORMAL)
    - include_all flag
    """

    # ==================== Basic Retrieval Tests ====================

    def test_get_datasets_basic_pagination(self, db_session_with_containers: Session):
        """Test basic pagination without user or filters."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        page = 1
        per_page = 20

        for i in range(5):
            DatasetRetrievalTestDataFactory.create_dataset(
                db_session_with_containers,
                tenant_id=tenant.id,
                created_by=account.id,
                name=f"Dataset {i}",
                permission=DatasetPermissionEnum.ALL_TEAM,
            )

        # Act
        datasets, total = DatasetService.get_datasets(page, per_page, tenant_id=tenant.id)

        # Assert
        assert len(datasets) == 5
        assert total == 5

    def test_get_datasets_with_search(self, db_session_with_containers: Session):
        """Test get_datasets with search keyword."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        page = 1
        per_page = 20
        search = "test"

        DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            name="Test Dataset",
            permission=DatasetPermissionEnum.ALL_TEAM,
        )
        DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            name="Another Dataset",
            permission=DatasetPermissionEnum.ALL_TEAM,
        )

        # Act
        datasets, total = DatasetService.get_datasets(page, per_page, tenant_id=tenant.id, search=search)

        # Assert
        assert len(datasets) == 1
        assert total == 1

    def test_get_datasets_with_tag_filtering(self, db_session_with_containers: Session):
        """Test get_datasets with tag_ids filtering."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        page = 1
        per_page = 20

        dataset_1 = DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            permission=DatasetPermissionEnum.ALL_TEAM,
        )
        dataset_2 = DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            permission=DatasetPermissionEnum.ALL_TEAM,
        )

        tag_1 = DatasetRetrievalTestDataFactory.create_tag_binding(
            db_session_with_containers, tenant.id, account.id, dataset_1.id
        )
        tag_2 = DatasetRetrievalTestDataFactory.create_tag_binding(
            db_session_with_containers, tenant.id, account.id, dataset_2.id
        )
        db_session_with_containers.add(
            TagBinding(
                tenant_id=tenant.id,
                tag_id=tag_2.id,
                target_id=dataset_1.id,
                created_by=account.id,
            )
        )
        db_session_with_containers.commit()
        tag_ids = [tag_1.id, tag_2.id]

        # Act
        datasets, total = DatasetService.get_datasets(page, per_page, tenant_id=tenant.id, tag_ids=tag_ids)

        # Assert
        assert len(datasets) == 1
        assert total == 1
        assert datasets[0].id == dataset_1.id

    def test_get_datasets_with_empty_tag_ids(self, db_session_with_containers: Session):
        """Test get_datasets with empty tag_ids skips tag filtering and returns all matching datasets."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        page = 1
        per_page = 20
        tag_ids = []

        for i in range(3):
            DatasetRetrievalTestDataFactory.create_dataset(
                db_session_with_containers,
                tenant_id=tenant.id,
                created_by=account.id,
                name=f"dataset-{i}",
                permission=DatasetPermissionEnum.ALL_TEAM,
            )

        # Act
        datasets, total = DatasetService.get_datasets(page, per_page, tenant_id=tenant.id, tag_ids=tag_ids)

        # Assert
        # When tag_ids is empty, tag filtering is skipped, so normal query results are returned
        assert len(datasets) == 3
        assert total == 3

    # ==================== Permission-Based Filtering Tests ====================

    def test_get_datasets_without_user_shows_only_all_team(self, db_session_with_containers: Session):
        """Test that without user, only ALL_TEAM datasets are shown."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        page = 1
        per_page = 20

        DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            permission=DatasetPermissionEnum.ALL_TEAM,
        )
        DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=account.id,
            permission=DatasetPermissionEnum.ONLY_ME,
        )

        # Act
        datasets, total = DatasetService.get_datasets(page, per_page, tenant_id=tenant.id, user=None)

        # Assert
        assert len(datasets) == 1
        assert total == 1

    def test_get_datasets_owner_with_include_all(self, db_session_with_containers: Session):
        """Test that OWNER with include_all=True sees all datasets."""
        # Arrange
        owner, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(
            db_session_with_containers, role=TenantAccountRole.OWNER
        )

        for i, permission in enumerate(
            [DatasetPermissionEnum.ONLY_ME, DatasetPermissionEnum.ALL_TEAM, DatasetPermissionEnum.PARTIAL_TEAM]
        ):
            DatasetRetrievalTestDataFactory.create_dataset(
                db_session_with_containers,
                tenant_id=tenant.id,
                created_by=owner.id,
                name=f"dataset-{i}",
                permission=permission,
            )

        # Act
        datasets, total = DatasetService.get_datasets(
            page=1,
            per_page=20,
            tenant_id=tenant.id,
            user=owner,
            include_all=True,
        )

        # Assert
        assert len(datasets) == 3
        assert total == 3

    def test_get_datasets_normal_user_only_me_permission(self, db_session_with_containers: Session):
        """Test that normal user sees ONLY_ME datasets they created."""
        # Arrange
        user, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(
            db_session_with_containers, role=TenantAccountRole.NORMAL
        )

        DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            permission=DatasetPermissionEnum.ONLY_ME,
        )

        # Act
        datasets, total = DatasetService.get_datasets(page=1, per_page=20, tenant_id=tenant.id, user=user)

        # Assert
        assert len(datasets) == 1
        assert total == 1

    def test_get_datasets_normal_user_all_team_permission(self, db_session_with_containers: Session):
        """Test that normal user sees ALL_TEAM datasets."""
        # Arrange
        user, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(
            db_session_with_containers, role=TenantAccountRole.NORMAL
        )
        owner = DatasetRetrievalTestDataFactory.create_account_in_tenant(
            db_session_with_containers, tenant, role=TenantAccountRole.OWNER
        )

        DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=owner.id,
            permission=DatasetPermissionEnum.ALL_TEAM,
        )

        # Act
        datasets, total = DatasetService.get_datasets(page=1, per_page=20, tenant_id=tenant.id, user=user)

        # Assert
        assert len(datasets) == 1
        assert total == 1

    def test_get_datasets_normal_user_partial_team_with_permission(self, db_session_with_containers: Session):
        """Test that normal user sees PARTIAL_TEAM datasets they have permission for."""
        # Arrange
        user, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(
            db_session_with_containers, role=TenantAccountRole.NORMAL
        )
        owner = DatasetRetrievalTestDataFactory.create_account_in_tenant(
            db_session_with_containers, tenant, role=TenantAccountRole.OWNER
        )

        dataset = DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=owner.id,
            permission=DatasetPermissionEnum.PARTIAL_TEAM,
        )
        DatasetRetrievalTestDataFactory.create_dataset_permission(
            db_session_with_containers, dataset.id, tenant.id, user.id
        )

        # Act
        datasets, total = DatasetService.get_datasets(page=1, per_page=20, tenant_id=tenant.id, user=user)

        # Assert
        assert len(datasets) == 1
        assert total == 1

    def test_get_datasets_dataset_operator_with_permissions(self, db_session_with_containers: Session):
        """Test that DATASET_OPERATOR only sees datasets they have explicit permission for."""
        # Arrange
        operator, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(
            db_session_with_containers, role=TenantAccountRole.DATASET_OPERATOR
        )
        owner = DatasetRetrievalTestDataFactory.create_account_in_tenant(
            db_session_with_containers, tenant, role=TenantAccountRole.OWNER
        )

        dataset = DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=owner.id,
            permission=DatasetPermissionEnum.ONLY_ME,
        )
        DatasetRetrievalTestDataFactory.create_dataset_permission(
            db_session_with_containers, dataset.id, tenant.id, operator.id
        )

        # Act
        datasets, total = DatasetService.get_datasets(page=1, per_page=20, tenant_id=tenant.id, user=operator)

        # Assert
        assert len(datasets) == 1
        assert total == 1

    def test_get_datasets_dataset_operator_without_permissions(self, db_session_with_containers: Session):
        """Test that DATASET_OPERATOR without permissions returns empty result."""
        # Arrange
        operator, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(
            db_session_with_containers, role=TenantAccountRole.DATASET_OPERATOR
        )
        owner = DatasetRetrievalTestDataFactory.create_account_in_tenant(
            db_session_with_containers, tenant, role=TenantAccountRole.OWNER
        )
        DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=owner.id,
            permission=DatasetPermissionEnum.ALL_TEAM,
        )

        # Act
        datasets, total = DatasetService.get_datasets(page=1, per_page=20, tenant_id=tenant.id, user=operator)

        # Assert
        assert datasets == []
        assert total == 0


class TestDatasetServiceGetDataset:
    """Comprehensive integration tests for DatasetService.get_dataset method."""

    def test_get_dataset_success(self, db_session_with_containers: Session):
        """Test successful retrieval of a single dataset."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers, tenant_id=tenant.id, created_by=account.id
        )

        # Act
        result = DatasetService.get_dataset(dataset.id)

        # Assert
        assert result is not None
        assert result.id == dataset.id

    def test_get_dataset_not_found(self, db_session_with_containers: Session):
        """Test retrieval when dataset doesn't exist."""
        # Arrange
        dataset_id = str(uuid4())

        # Act
        result = DatasetService.get_dataset(dataset_id)

        # Assert
        assert result is None


class TestDatasetServiceGetDatasetsByIds:
    """Comprehensive integration tests for DatasetService.get_datasets_by_ids method."""

    def test_get_datasets_by_ids_success(self, db_session_with_containers: Session):
        """Test successful bulk retrieval of datasets by IDs."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        datasets = [
            DatasetRetrievalTestDataFactory.create_dataset(
                db_session_with_containers, tenant_id=tenant.id, created_by=account.id
            )
            for _ in range(3)
        ]
        dataset_ids = [dataset.id for dataset in datasets]

        # Act
        result_datasets, total = DatasetService.get_datasets_by_ids(dataset_ids, tenant.id)

        # Assert
        assert len(result_datasets) == 3
        assert total == 3
        assert all(dataset.id in dataset_ids for dataset in result_datasets)

    def test_get_datasets_by_ids_empty_list(self, db_session_with_containers: Session):
        """Test get_datasets_by_ids with empty list returns empty result."""
        # Arrange
        tenant_id = str(uuid4())
        dataset_ids = []

        # Act
        datasets, total = DatasetService.get_datasets_by_ids(dataset_ids, tenant_id)

        # Assert
        assert datasets == []
        assert total == 0

    def test_get_datasets_by_ids_none_list(self, db_session_with_containers: Session):
        """Test get_datasets_by_ids with None returns empty result."""
        # Arrange
        tenant_id = str(uuid4())

        # Act
        datasets, total = DatasetService.get_datasets_by_ids(None, tenant_id)

        # Assert
        assert datasets == []
        assert total == 0


class TestDatasetServiceGetProcessRules:
    """Comprehensive integration tests for DatasetService.get_process_rules method."""

    def test_get_process_rules_with_existing_rule(self, db_session_with_containers: Session):
        """Test retrieval of process rules when rule exists."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers, tenant_id=tenant.id, created_by=account.id
        )

        rules_data = {
            "pre_processing_rules": [{"id": "remove_extra_spaces", "enabled": True}],
            "segmentation": {"delimiter": "\n", "max_tokens": 500},
        }
        DatasetRetrievalTestDataFactory.create_process_rule(
            db_session_with_containers,
            dataset_id=dataset.id,
            created_by=account.id,
            mode=ProcessRuleMode.CUSTOM,
            rules=rules_data,
        )

        # Act
        result = DatasetService.get_process_rules(dataset.id)

        # Assert
        assert result["mode"] == "custom"
        assert result["rules"] == rules_data

    def test_get_process_rules_without_existing_rule(self, db_session_with_containers: Session):
        """Test retrieval of process rules when no rule exists (returns defaults)."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers, tenant_id=tenant.id, created_by=account.id
        )

        # Act
        result = DatasetService.get_process_rules(dataset.id)

        # Assert
        assert result["mode"] == DocumentService.DEFAULT_RULES["mode"]
        assert "rules" in result
        assert result["rules"] == DocumentService.DEFAULT_RULES["rules"]


class TestDatasetServiceGetDatasetQueries:
    """Comprehensive integration tests for DatasetService.get_dataset_queries method."""

    def test_get_dataset_queries_success(self, db_session_with_containers: Session):
        """Test successful retrieval of dataset queries."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers, tenant_id=tenant.id, created_by=account.id
        )
        page = 1
        per_page = 20

        for i in range(3):
            DatasetRetrievalTestDataFactory.create_dataset_query(
                db_session_with_containers,
                dataset_id=dataset.id,
                created_by=account.id,
                content=f"query-{i}",
            )

        # Act
        queries, total = DatasetService.get_dataset_queries(dataset.id, page, per_page)

        # Assert
        assert len(queries) == 3
        assert total == 3
        assert all(query.dataset_id == dataset.id for query in queries)

    def test_get_dataset_queries_empty_result(self, db_session_with_containers: Session):
        """Test retrieval when no queries exist."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers, tenant_id=tenant.id, created_by=account.id
        )
        page = 1
        per_page = 20

        # Act
        queries, total = DatasetService.get_dataset_queries(dataset.id, page, per_page)

        # Assert
        assert queries == []
        assert total == 0


class TestDatasetServiceGetRelatedApps:
    """Comprehensive integration tests for DatasetService.get_related_apps method."""

    def test_get_related_apps_success(self, db_session_with_containers: Session):
        """Test successful retrieval of related apps."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers, tenant_id=tenant.id, created_by=account.id
        )

        for _ in range(2):
            DatasetRetrievalTestDataFactory.create_app_dataset_join(db_session_with_containers, dataset.id)

        # Act
        result = DatasetService.get_related_apps(dataset.id)

        # Assert
        assert len(result) == 2
        assert all(join.dataset_id == dataset.id for join in result)

    def test_get_related_apps_empty_result(self, db_session_with_containers: Session):
        """Test retrieval when no related apps exist."""
        # Arrange
        account, tenant = DatasetRetrievalTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetRetrievalTestDataFactory.create_dataset(
            db_session_with_containers, tenant_id=tenant.id, created_by=account.id
        )

        # Act
        result = DatasetService.get_related_apps(dataset.id)

        # Assert
        assert result == []
