package tech.easyflow.agent.runtime;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.Assert;
import org.junit.Test;
import tech.easyflow.agent.entity.Agent;
import tech.easyflow.agent.entity.AgentKnowledgeBinding;
import tech.easyflow.ai.entity.DocumentCollection;
import tech.easyflow.ai.enums.PublishStatus;
import tech.easyflow.ai.service.DocumentCollectionService;
import tech.easyflow.common.entity.LoginAccount;
import tech.easyflow.common.web.exceptions.BusinessException;
import tech.easyflow.system.enums.CategoryResourceType;
import tech.easyflow.system.enums.ResourceAction;
import tech.easyflow.system.permission.resource.VisibilityResource;
import tech.easyflow.system.service.ResourceAccessService;

import java.lang.reflect.Proxy;
import java.math.BigInteger;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Agent 聊天临时能力编排服务测试。
 */
public class AgentChatCapabilityServiceTest {

    @Test
    public void applyShouldAppendPublishedKnowledgeAndSkipBoundDuplicate() {
        DocumentCollectionService documentService = documentService(
                knowledge(1, PublishStatus.PUBLISHED),
                knowledge(2, PublishStatus.PUBLISHED)
        );
        AgentChatCapabilityService service = new AgentChatCapabilityService(
                documentService,
                new AllowResourceAccessService(),
                new ObjectMapper()
        );
        Agent agent = agentWithBoundKnowledge(1);

        AgentChatCapabilityService.AgentChatCapabilityResolution resolution = service.apply(
                agent,
                List.of(capability(1, 2, 2)),
                account()
        );

        List<AgentKnowledgeBinding> bindings = resolution.agent().getKnowledgeBindings();
        Assert.assertEquals(List.of(BigInteger.ONE, BigInteger.valueOf(2)), resolution.extraKnowledgeIds());
        Assert.assertEquals(2, bindings.size());
        Assert.assertEquals(BigInteger.ONE, bindings.get(0).getKnowledgeId());
        Assert.assertEquals(BigInteger.valueOf(2), bindings.get(1).getKnowledgeId());
        Assert.assertEquals("HYBRID", bindings.get(1).getRetrievalMode());
        Assert.assertTrue(bindings.get(1).getEnabled());
    }

    @Test
    public void applyShouldAppendWhenBoundKnowledgeIsDisabled() {
        DocumentCollectionService documentService = documentService(knowledge(2, PublishStatus.PUBLISHED));
        AgentChatCapabilityService service = new AgentChatCapabilityService(
                documentService,
                new AllowResourceAccessService(),
                new ObjectMapper()
        );
        Agent agent = agentWithBoundKnowledge(2);
        agent.getKnowledgeBindings().get(0).setEnabled(false);

        AgentChatCapabilityService.AgentChatCapabilityResolution resolution = service.apply(
                agent,
                List.of(capability(2)),
                account()
        );

        List<AgentKnowledgeBinding> bindings = resolution.agent().getKnowledgeBindings();
        Assert.assertEquals(2, bindings.size());
        Assert.assertFalse(bindings.get(0).getEnabled());
        Assert.assertTrue(bindings.get(1).getEnabled());
        Assert.assertEquals(BigInteger.valueOf(2), bindings.get(1).getKnowledgeId());
    }

    @Test
    public void resolveKnowledgeIdsShouldRejectTooManyItems() {
        AgentChatCapabilityService service = new AgentChatCapabilityService(
                documentService(),
                new AllowResourceAccessService(),
                new ObjectMapper()
        );

        Assert.assertThrows(BusinessException.class,
                () -> service.resolveKnowledgeIds(List.of(capability(1, 2, 3, 4))));
    }

    @Test
    public void applyShouldRejectUnpublishedKnowledge() {
        DocumentCollectionService documentService = documentService(knowledge(2, PublishStatus.OFFLINE));
        AgentChatCapabilityService service = new AgentChatCapabilityService(
                documentService,
                new AllowResourceAccessService(),
                new ObjectMapper()
        );

        Assert.assertThrows(BusinessException.class,
                () -> service.apply(agentWithBoundKnowledge(1), List.of(capability(2)), account()));
    }

    @Test
    public void applyShouldRejectUnauthorizedKnowledge() {
        DocumentCollectionService documentService = documentService(knowledge(2, PublishStatus.PUBLISHED));
        AgentChatCapabilityService service = new AgentChatCapabilityService(
                documentService,
                new DenyResourceAccessService(),
                new ObjectMapper()
        );

        Assert.assertThrows(BusinessException.class,
                () -> service.apply(agentWithBoundKnowledge(1), List.of(capability(2)), account()));
    }

    private static Agent agentWithBoundKnowledge(int knowledgeId) {
        Agent agent = new Agent();
        agent.setId(BigInteger.TEN);
        agent.setTenantId(BigInteger.ONE);
        AgentKnowledgeBinding binding = new AgentKnowledgeBinding();
        binding.setAgentId(agent.getId());
        binding.setKnowledgeId(BigInteger.valueOf(knowledgeId));
        binding.setRetrievalMode("HYBRID");
        binding.setEnabled(true);
        agent.setKnowledgeBindings(List.of(binding));
        return agent;
    }

    private static AgentChatCapability capability(int... ids) {
        AgentChatCapability capability = new AgentChatCapability();
        capability.setType("KNOWLEDGE");
        capability.setResourceIds(java.util.Arrays.stream(ids)
                .mapToObj(BigInteger::valueOf)
                .toList());
        return capability;
    }

    private static DocumentCollection knowledge(int id, PublishStatus status) {
        DocumentCollection collection = new DocumentCollection();
        collection.setId(BigInteger.valueOf(id));
        collection.setTitle("知识库" + id);
        collection.setPublishStatus(status.name());
        return collection;
    }

    private static LoginAccount account() {
        LoginAccount account = new LoginAccount();
        account.setId(BigInteger.valueOf(100));
        account.setTenantId(BigInteger.ONE);
        return account;
    }

    private static DocumentCollectionService documentService(DocumentCollection... collections) {
        Map<BigInteger, DocumentCollection> collectionMap = new HashMap<>();
        for (DocumentCollection collection : collections) {
            collectionMap.put(collection.getId(), collection);
        }
        return (DocumentCollectionService) Proxy.newProxyInstance(
                AgentChatCapabilityServiceTest.class.getClassLoader(),
                new Class<?>[]{DocumentCollectionService.class},
                (proxy, method, args) -> switch (method.getName()) {
                    case "getById" -> collectionMap.get(new BigInteger(String.valueOf(args[0])));
                    case "toPublishedView" -> args[0];
                    case "listByIds" -> ((java.util.Collection<?>) args[0]).stream()
                            .map(id -> collectionMap.get(new BigInteger(String.valueOf(id))))
                            .filter(java.util.Objects::nonNull)
                            .toList();
                    default -> throw new UnsupportedOperationException(method.getName());
                }
        );
    }

    private static class AllowResourceAccessService implements ResourceAccessService {
        @Override public boolean canAccess(CategoryResourceType resourceType, VisibilityResource resource, ResourceAction action) { return true; }
        @Override public boolean canAccess(LoginAccount loginAccount, CategoryResourceType resourceType, VisibilityResource resource, ResourceAction action) { return true; }
        @Override public void assertAccess(CategoryResourceType resourceType, VisibilityResource resource, ResourceAction action, String message) {}
    }

    private static class DenyResourceAccessService extends AllowResourceAccessService {
        @Override public void assertAccess(CategoryResourceType resourceType, VisibilityResource resource, ResourceAction action, String message) {
            throw new BusinessException(message);
        }
    }
}
