import json
from unittest.mock import Mock, patch
from uuid import uuid4

import pytest
from sqlalchemy.orm import Session

from core.rag.index_processor.constant.index_type import IndexTechniqueType
from graphon.model_runtime.entities.model_entities import ModelType
from models.account import (
    Account,
    AccountStatus,
    Tenant,
    TenantAccountJoin,
    TenantAccountRole,
    TenantStatus,
)
from models.dataset import Dataset, ExternalKnowledgeApis, ExternalKnowledgeBindings
from models.enums import DataSourceType
from services.dataset_service import DatasetService
from services.errors.account import NoPermissionError


class DatasetUpdateTestDataFactory:
    """Factory class for creating real test data for dataset update integration tests."""

    @staticmethod
    def create_account_with_tenant(
        db_session_with_containers: Session, role: TenantAccountRole = TenantAccountRole.OWNER
    ) -> tuple[Account, Tenant]:
        """Create a real account and tenant with the given role."""
        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.commit()

        tenant = Tenant(name=f"tenant-{account.id}", status=TenantStatus.NORMAL)
        db_session_with_containers.add(tenant)
        db_session_with_containers.commit()

        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_dataset(
        db_session_with_containers: Session,
        tenant_id: str,
        created_by: str,
        provider: str = "vendor",
        name: str = "old_name",
        description: str = "old_description",
        indexing_technique: str = IndexTechniqueType.HIGH_QUALITY,
        retrieval_model: str = "old_model",
        permission: str = "only_me",
        embedding_model_provider: str | None = None,
        embedding_model: str | None = None,
        collection_binding_id: str | None = None,
    ) -> Dataset:
        """Create a real dataset."""
        dataset = Dataset(
            tenant_id=tenant_id,
            name=name,
            description=description,
            data_source_type=DataSourceType.UPLOAD_FILE,
            indexing_technique=indexing_technique,
            created_by=created_by,
            provider=provider,
            retrieval_model=retrieval_model,
            permission=permission,
            embedding_model_provider=embedding_model_provider,
            embedding_model=embedding_model,
            collection_binding_id=collection_binding_id,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.commit()
        return dataset

    @staticmethod
    def create_external_binding(
        db_session_with_containers: Session,
        tenant_id: str,
        dataset_id: str,
        created_by: str,
        external_knowledge_id: str = "old_knowledge_id",
        external_knowledge_api_id: str | None = None,
    ) -> ExternalKnowledgeBindings:
        """Create a real external knowledge binding."""
        if external_knowledge_api_id is None:
            external_knowledge_api_id = str(uuid4())
        binding = ExternalKnowledgeBindings(
            tenant_id=tenant_id,
            dataset_id=dataset_id,
            created_by=created_by,
            external_knowledge_id=external_knowledge_id,
            external_knowledge_api_id=external_knowledge_api_id,
        )
        db_session_with_containers.add(binding)
        db_session_with_containers.commit()
        return binding

    @staticmethod
    def create_external_knowledge_api(
        db_session_with_containers: Session,
        tenant_id: str,
        created_by: str,
        api_id: str | None = None,
        name: str = "test-api",
    ) -> ExternalKnowledgeApis:
        """Create a real external knowledge API template for tenant-scoped update validation."""
        external_api = ExternalKnowledgeApis(
            tenant_id=tenant_id,
            created_by=created_by,
            updated_by=created_by,
            name=name,
            description="test description",
            settings=json.dumps(
                {
                    "endpoint": "https://example.com",
                    "api_key": "test-api-key",
                }
            ),
        )
        if api_id is not None:
            external_api.id = api_id
        db_session_with_containers.add(external_api)
        db_session_with_containers.commit()
        return external_api


class TestDatasetServiceUpdateDataset:
    """
    Comprehensive integration tests for DatasetService.update_dataset method.

    This test suite covers all supported scenarios including:
    - External dataset updates
    - Internal dataset updates with different indexing techniques
    - Embedding model updates
    - Permission checks
    - Error conditions and edge cases
    """

    # ==================== External Dataset Tests ====================

    def test_update_external_dataset_success(self, db_session_with_containers: Session):
        """Test successful update of external dataset."""
        user, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            provider="external",
            name="old_name",
            description="old_description",
            retrieval_model="old_model",
        )
        binding = DatasetUpdateTestDataFactory.create_external_binding(
            db_session_with_containers,
            tenant_id=tenant.id,
            dataset_id=dataset.id,
            created_by=user.id,
        )
        binding_id = binding.id
        db_session_with_containers.expunge(binding)
        external_api = DatasetUpdateTestDataFactory.create_external_knowledge_api(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
        )

        update_data = {
            "name": "new_name",
            "description": "new_description",
            "external_retrieval_model": "new_model",
            "permission": "only_me",
            "external_knowledge_id": "new_knowledge_id",
            "external_knowledge_api_id": external_api.id,
        }

        result = DatasetService.update_dataset(dataset.id, update_data, user)

        db_session_with_containers.refresh(dataset)
        updated_binding = db_session_with_containers.query(ExternalKnowledgeBindings).filter_by(id=binding_id).first()

        assert dataset.name == "new_name"
        assert dataset.description == "new_description"
        assert dataset.retrieval_model == "new_model"
        assert updated_binding is not None
        assert updated_binding.external_knowledge_id == "new_knowledge_id"
        assert updated_binding.external_knowledge_api_id == update_data["external_knowledge_api_id"]
        assert result.id == dataset.id

    def test_update_external_dataset_missing_knowledge_id_error(self, db_session_with_containers: Session):
        """Test error when external knowledge id is missing."""
        user, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            provider="external",
        )
        DatasetUpdateTestDataFactory.create_external_binding(
            db_session_with_containers,
            tenant_id=tenant.id,
            dataset_id=dataset.id,
            created_by=user.id,
        )

        update_data = {"name": "new_name", "external_knowledge_api_id": str(uuid4())}

        with pytest.raises(ValueError) as context:
            DatasetService.update_dataset(dataset.id, update_data, user)

        assert "External knowledge id is required" in str(context.value)
        db_session_with_containers.rollback()

    def test_update_external_dataset_missing_api_id_error(self, db_session_with_containers: Session):
        """Test error when external knowledge api id is missing."""
        user, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            provider="external",
        )
        DatasetUpdateTestDataFactory.create_external_binding(
            db_session_with_containers,
            tenant_id=tenant.id,
            dataset_id=dataset.id,
            created_by=user.id,
        )

        update_data = {"name": "new_name", "external_knowledge_id": "knowledge_id"}

        with pytest.raises(ValueError) as context:
            DatasetService.update_dataset(dataset.id, update_data, user)

        assert "External knowledge api id is required" in str(context.value)
        db_session_with_containers.rollback()

    def test_update_external_dataset_binding_not_found_error(self, db_session_with_containers: Session):
        """Test error when external knowledge binding is not found."""
        user, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            provider="external",
        )
        external_api = DatasetUpdateTestDataFactory.create_external_knowledge_api(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
        )

        update_data = {
            "name": "new_name",
            "external_knowledge_id": "knowledge_id",
            "external_knowledge_api_id": external_api.id,
        }

        with pytest.raises(ValueError) as context:
            DatasetService.update_dataset(dataset.id, update_data, user)

        assert "External knowledge binding not found" in str(context.value)
        db_session_with_containers.rollback()

    # ==================== Internal Dataset Basic Tests ====================

    def test_update_internal_dataset_basic_success(self, db_session_with_containers: Session):
        """Test successful update of internal dataset with basic fields."""
        user, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        existing_binding_id = str(uuid4())
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            provider="vendor",
            indexing_technique=IndexTechniqueType.HIGH_QUALITY,
            embedding_model_provider="openai",
            embedding_model="text-embedding-ada-002",
            collection_binding_id=existing_binding_id,
        )

        update_data = {
            "name": "new_name",
            "description": "new_description",
            "indexing_technique": IndexTechniqueType.HIGH_QUALITY,
            "retrieval_model": "new_model",
            "embedding_model_provider": "openai",
            "embedding_model": "text-embedding-ada-002",
        }

        result = DatasetService.update_dataset(dataset.id, update_data, user)
        db_session_with_containers.refresh(dataset)

        assert dataset.name == "new_name"
        assert dataset.description == "new_description"
        assert dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY
        assert dataset.retrieval_model == "new_model"
        assert dataset.embedding_model_provider == "openai"
        assert dataset.embedding_model == "text-embedding-ada-002"
        assert result.id == dataset.id

    def test_update_internal_dataset_filter_none_values(self, db_session_with_containers: Session):
        """Test that None values are filtered out except for description field."""
        user, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        existing_binding_id = str(uuid4())
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            provider="vendor",
            indexing_technique=IndexTechniqueType.HIGH_QUALITY,
            embedding_model_provider="openai",
            embedding_model="text-embedding-ada-002",
            collection_binding_id=existing_binding_id,
        )

        update_data = {
            "name": "new_name",
            "description": None,
            "indexing_technique": IndexTechniqueType.HIGH_QUALITY,
            "retrieval_model": "new_model",
            "embedding_model_provider": None,
            "embedding_model": None,
        }

        result = DatasetService.update_dataset(dataset.id, update_data, user)
        db_session_with_containers.refresh(dataset)

        assert dataset.name == "new_name"
        assert dataset.description is None
        assert dataset.embedding_model_provider == "openai"
        assert dataset.embedding_model == "text-embedding-ada-002"
        assert dataset.retrieval_model == "new_model"
        assert result.id == dataset.id

    # ==================== Indexing Technique Switch Tests ====================

    def test_update_internal_dataset_indexing_technique_to_economy(self, db_session_with_containers: Session):
        """Test updating internal dataset indexing technique to economy."""
        user, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        existing_binding_id = str(uuid4())
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            provider="vendor",
            indexing_technique=IndexTechniqueType.HIGH_QUALITY,
            embedding_model_provider="openai",
            embedding_model="text-embedding-ada-002",
            collection_binding_id=existing_binding_id,
        )

        update_data = {
            "indexing_technique": IndexTechniqueType.ECONOMY,
            "retrieval_model": "new_model",
        }

        with patch("services.dataset_service.deal_dataset_vector_index_task") as mock_task:
            result = DatasetService.update_dataset(dataset.id, update_data, user)
            mock_task.delay.assert_called_once_with(dataset.id, "remove")

        db_session_with_containers.refresh(dataset)
        assert dataset.indexing_technique == IndexTechniqueType.ECONOMY
        assert dataset.embedding_model is None
        assert dataset.embedding_model_provider is None
        assert dataset.collection_binding_id is None
        assert dataset.retrieval_model == "new_model"
        assert result.id == dataset.id

    def test_update_internal_dataset_indexing_technique_to_high_quality(self, db_session_with_containers: Session):
        """Test updating internal dataset indexing technique to high_quality."""
        user, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            provider="vendor",
            indexing_technique=IndexTechniqueType.ECONOMY,
        )

        embedding_model = Mock()
        embedding_model.model_name = "text-embedding-ada-002"
        embedding_model.provider = "openai"

        binding = Mock()
        binding.id = str(uuid4())

        update_data = {
            "indexing_technique": IndexTechniqueType.HIGH_QUALITY,
            "embedding_model_provider": "openai",
            "embedding_model": "text-embedding-ada-002",
            "retrieval_model": "new_model",
        }

        with (
            patch("services.dataset_service.current_user", user),
            patch("services.dataset_service.ModelManager.for_tenant") as mock_model_manager,
            patch(
                "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding"
            ) as mock_get_binding,
            patch("services.dataset_service.deal_dataset_vector_index_task") as mock_task,
        ):
            mock_model_manager.return_value.get_model_instance.return_value = embedding_model
            mock_get_binding.return_value = binding

            result = DatasetService.update_dataset(dataset.id, update_data, user)

            mock_model_manager.return_value.get_model_instance.assert_called_once_with(
                tenant_id=tenant.id,
                provider="openai",
                model_type=ModelType.TEXT_EMBEDDING,
                model="text-embedding-ada-002",
            )
            mock_get_binding.assert_called_once_with("openai", "text-embedding-ada-002")
            mock_task.delay.assert_called_once_with(dataset.id, "add")

        db_session_with_containers.refresh(dataset)
        assert dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY
        assert dataset.embedding_model == "text-embedding-ada-002"
        assert dataset.embedding_model_provider == "openai"
        assert dataset.collection_binding_id == binding.id
        assert dataset.retrieval_model == "new_model"
        assert result.id == dataset.id

    # ==================== Embedding Model Update Tests ====================

    def test_update_internal_dataset_keep_existing_embedding_model_when_indexing_technique_unchanged(
        self, db_session_with_containers
    ):
        """Test preserving embedding settings when indexing technique remains unchanged."""
        user, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        existing_binding_id = str(uuid4())
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            provider="vendor",
            indexing_technique=IndexTechniqueType.HIGH_QUALITY,
            embedding_model_provider="openai",
            embedding_model="text-embedding-ada-002",
            collection_binding_id=existing_binding_id,
        )

        update_data = {
            "name": "new_name",
            "indexing_technique": IndexTechniqueType.HIGH_QUALITY,
            "retrieval_model": "new_model",
        }

        result = DatasetService.update_dataset(dataset.id, update_data, user)
        db_session_with_containers.refresh(dataset)

        assert dataset.name == "new_name"
        assert dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY
        assert dataset.embedding_model_provider == "openai"
        assert dataset.embedding_model == "text-embedding-ada-002"
        assert dataset.collection_binding_id == existing_binding_id
        assert dataset.retrieval_model == "new_model"
        assert result.id == dataset.id

    def test_update_internal_dataset_embedding_model_update(self, db_session_with_containers: Session):
        """Test updating internal dataset with new embedding model."""
        user, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        existing_binding_id = str(uuid4())
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            provider="vendor",
            indexing_technique=IndexTechniqueType.HIGH_QUALITY,
            embedding_model_provider="openai",
            embedding_model="text-embedding-ada-002",
            collection_binding_id=existing_binding_id,
        )

        embedding_model = Mock()
        embedding_model.model_name = "text-embedding-3-small"
        embedding_model.provider = "openai"

        binding = Mock()
        binding.id = str(uuid4())

        update_data = {
            "indexing_technique": IndexTechniqueType.HIGH_QUALITY,
            "embedding_model_provider": "openai",
            "embedding_model": "text-embedding-3-small",
            "retrieval_model": "new_model",
        }

        with (
            patch("services.dataset_service.current_user", user),
            patch("services.dataset_service.ModelManager.for_tenant") as mock_model_manager,
            patch(
                "services.dataset_service.DatasetCollectionBindingService.get_dataset_collection_binding"
            ) as mock_get_binding,
            patch("services.dataset_service.deal_dataset_vector_index_task") as mock_task,
            patch("services.dataset_service.regenerate_summary_index_task") as mock_regenerate_task,
        ):
            mock_model_manager.return_value.get_model_instance.return_value = embedding_model
            mock_get_binding.return_value = binding

            result = DatasetService.update_dataset(dataset.id, update_data, user)

            mock_model_manager.return_value.get_model_instance.assert_called_once_with(
                tenant_id=tenant.id,
                provider="openai",
                model_type=ModelType.TEXT_EMBEDDING,
                model="text-embedding-3-small",
            )
            mock_get_binding.assert_called_once_with("openai", "text-embedding-3-small")
            mock_task.delay.assert_called_once_with(dataset.id, "update")
            mock_regenerate_task.delay.assert_called_once_with(
                dataset.id,
                regenerate_reason="embedding_model_changed",
                regenerate_vectors_only=True,
            )

        db_session_with_containers.refresh(dataset)
        assert dataset.embedding_model == "text-embedding-3-small"
        assert dataset.embedding_model_provider == "openai"
        assert dataset.collection_binding_id == binding.id
        assert dataset.retrieval_model == "new_model"
        assert result.id == dataset.id

    # ==================== Error Handling Tests ====================

    def test_update_dataset_not_found_error(self, db_session_with_containers: Session):
        """Test error when dataset is not found."""
        user, _ = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        update_data = {"name": "new_name"}

        with pytest.raises(ValueError) as context:
            DatasetService.update_dataset(str(uuid4()), update_data, user)

        assert "Dataset not found" in str(context.value)

    def test_update_dataset_permission_error(self, db_session_with_containers: Session):
        """Test error when user doesn't have permission."""
        owner, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(
            db_session_with_containers, role=TenantAccountRole.OWNER
        )
        outsider, _ = DatasetUpdateTestDataFactory.create_account_with_tenant(
            db_session_with_containers, role=TenantAccountRole.NORMAL
        )
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=owner.id,
            provider="vendor",
            permission="only_me",
        )

        update_data = {"name": "new_name"}

        with pytest.raises(NoPermissionError):
            DatasetService.update_dataset(dataset.id, update_data, outsider)

    def test_update_internal_dataset_embedding_model_error(self, db_session_with_containers: Session):
        """Test error when embedding model is not available."""
        user, tenant = DatasetUpdateTestDataFactory.create_account_with_tenant(db_session_with_containers)
        dataset = DatasetUpdateTestDataFactory.create_dataset(
            db_session_with_containers,
            tenant_id=tenant.id,
            created_by=user.id,
            provider="vendor",
            indexing_technique=IndexTechniqueType.ECONOMY,
        )

        update_data = {
            "indexing_technique": IndexTechniqueType.HIGH_QUALITY,
            "embedding_model_provider": "invalid_provider",
            "embedding_model": "invalid_model",
            "retrieval_model": "new_model",
        }

        with (
            patch("services.dataset_service.current_user", user),
            patch("services.dataset_service.ModelManager.for_tenant") as mock_model_manager,
        ):
            mock_model_manager.return_value.get_model_instance.side_effect = Exception("No Embedding Model available")

            with pytest.raises(Exception) as context:
                DatasetService.update_dataset(dataset.id, update_data, user)

        assert "No Embedding Model available".lower() in str(context.value).lower()