package at.procon.dip.embedding.service;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

import at.procon.dip.domain.access.DocumentVisibility;
import at.procon.dip.domain.document.DocumentFamily;
import at.procon.dip.domain.document.DocumentStatus;
import at.procon.dip.domain.document.DocumentType;
import at.procon.dip.domain.document.SourceType;
import at.procon.dip.domain.document.entity.Document;
import at.procon.dip.embedding.config.EmbeddingPolicyCondition;
import at.procon.dip.embedding.config.EmbeddingPolicyProperties;
import at.procon.dip.embedding.config.EmbeddingPolicyRule;
import at.procon.dip.embedding.config.EmbeddingPolicyUse;
import at.procon.dip.ingestion.spi.OriginalContentStoragePolicy;
import at.procon.dip.ingestion.spi.SourceDescriptor;
import java.time.OffsetDateTime;
import java.util.Map;
import java.util.UUID;
import org.junit.jupiter.api.Test;

class DefaultEmbeddingPolicyResolverTest {

    @Test
    void shouldPreferHintAndOverrideFromAttributes() {
        EmbeddingPolicyProperties properties = baseProperties();

        EmbeddingPolicyRule mailRule = new EmbeddingPolicyRule();
        EmbeddingPolicyUse mailUse = new EmbeddingPolicyUse();
        mailUse.setPolicyKey("mail-default");
        mailUse.setModelKey("e5-default");
        mailUse.setQueryModelKey("e5-default");
        mailUse.setProfileKey("mail-message");
        mailUse.setEnabled(true);
        mailRule.setUse(mailUse);
        properties.getRules().add(mailRule);

        EmbeddingPolicyRule tedRule = new EmbeddingPolicyRule();
        EmbeddingPolicyUse tedUse = new EmbeddingPolicyUse();
        tedUse.setPolicyKey("ted-default");
        tedUse.setModelKey("e5-default");
        tedUse.setQueryModelKey("e5-default");
        tedUse.setProfileKey("ted-semantic");
        tedUse.setEnabled(true);
        tedRule.setUse(tedUse);
        properties.getRules().add(tedRule);

        DefaultEmbeddingPolicyResolver resolver = new DefaultEmbeddingPolicyResolver(properties);

        SourceDescriptor descriptor = sourceDescriptor(SourceType.MAIL, "message/rfc822", Map.of(
                "embeddingPolicyHint", "mail-default",
                "embeddingPolicyKey", "ted-default"
        ));

        var policy = resolver.resolve(document(DocumentFamily.GENERIC, DocumentType.EMAIL, "en"), descriptor);

        assertThat(policy.policyKey()).isEqualTo("ted-default");
    }

    @Test
    void shouldMatchByMimeTypeUsingMediaTypeField() {
        EmbeddingPolicyProperties properties = baseProperties();
        EmbeddingPolicyRule rule = new EmbeddingPolicyRule();
        EmbeddingPolicyCondition when = new EmbeddingPolicyCondition();
        when.setSourceType("MAIL_ATTACHMENT");
        when.setMimeType("application/pdf");
        rule.setWhen(when);
        EmbeddingPolicyUse use = new EmbeddingPolicyUse();
        use.setPolicyKey("mail-attachment-pdf");
        use.setModelKey("e5-default");
        use.setQueryModelKey("e5-default");
        use.setProfileKey("attachment-chunks");
        rule.setUse(use);
        properties.getRules().add(rule);

        DefaultEmbeddingPolicyResolver resolver = new DefaultEmbeddingPolicyResolver(properties);

        var policy = resolver.resolve(document(DocumentFamily.GENERIC, DocumentType.UNKNOWN, "en"),
                sourceDescriptor(SourceType.MAIL, "application/pdf", Map.of()));

        assertThat(policy.policyKey()).isEqualTo("mail-attachment-pdf");
    }

    @Test
    void shouldFailForUnknownOverridePolicy() {
        DefaultEmbeddingPolicyResolver resolver = new DefaultEmbeddingPolicyResolver(baseProperties());
        SourceDescriptor descriptor = sourceDescriptor(SourceType.FILE_SYSTEM, "application/pdf", Map.of(
                "embeddingPolicyKey", "missing-policy"
        ));
        assertThatThrownBy(() -> resolver.resolve(document(DocumentFamily.GENERIC, DocumentType.UNKNOWN, "en"), descriptor))
                .isInstanceOf(IllegalArgumentException.class)
                .hasMessageContaining("Unknown embedding policy key");
    }

    private EmbeddingPolicyProperties baseProperties() {
        EmbeddingPolicyProperties properties = new EmbeddingPolicyProperties();
        EmbeddingPolicyUse defaultPolicy = new EmbeddingPolicyUse();
        defaultPolicy.setPolicyKey("generic-default");
        defaultPolicy.setModelKey("e5-default");
        defaultPolicy.setQueryModelKey("e5-default");
        defaultPolicy.setProfileKey("primary-and-chunks");
        defaultPolicy.setEnabled(true);
        properties.setDefaultPolicy(defaultPolicy);
        return properties;
    }

    private Document document(DocumentFamily family, DocumentType type, String language) {
        return Document.builder()
                .id(UUID.randomUUID())
                .documentFamily(family)
                .documentType(type)
                .languageCode(language)
                .status(DocumentStatus.RECEIVED)
                .visibility(DocumentVisibility.PUBLIC)
                .title("Test document")
                .build();
    }

    private SourceDescriptor sourceDescriptor(SourceType sourceType, String mediaType, Map<String, String> attrs) {
        return new SourceDescriptor(
                null,
                sourceType,
                "source-ref",
                "/tmp/source",
                "source.bin",
                mediaType,
                null,
                null,
                OffsetDateTime.now(),
                OriginalContentStoragePolicy.DEFAULT,
                attrs
        );
    }
}
