/*
 * SPDX-License-Identifier: Apache-2.0
 *
 * The OpenSearch Contributors require contributions made to
 * this file be licensed under the Apache-2.0 license or a
 * compatible open source license.
 */

package org.opensearch.indices.pollingingest.mappers;

import org.opensearch.index.IngestionShardPointer;
import org.opensearch.index.Message;
import org.opensearch.index.engine.FakeIngestionSource;
import org.opensearch.indices.pollingingest.ShardUpdateMessage;
import org.opensearch.test.OpenSearchTestCase;

import java.nio.charset.StandardCharsets;
import java.util.Map;

import static org.opensearch.action.index.IndexRequest.UNSET_AUTO_GENERATED_TIMESTAMP;

public class IngestionMessageMapperTests extends OpenSearchTestCase {

    public void testMapperTypeFromStringAndName() {
        assertEquals(IngestionMessageMapper.MapperType.DEFAULT, IngestionMessageMapper.MapperType.fromString("default"));
        assertEquals(IngestionMessageMapper.MapperType.DEFAULT, IngestionMessageMapper.MapperType.fromString("DEFAULT"));
        assertEquals(IngestionMessageMapper.MapperType.RAW_PAYLOAD, IngestionMessageMapper.MapperType.fromString("raw_payload"));
        assertEquals(IngestionMessageMapper.MapperType.RAW_PAYLOAD, IngestionMessageMapper.MapperType.fromString("RAW_PAYLOAD"));

        assertEquals("default", IngestionMessageMapper.MapperType.DEFAULT.getName());
        assertEquals("raw_payload", IngestionMessageMapper.MapperType.RAW_PAYLOAD.getName());
    }

    public void testMapperTypeFromStringInvalid() {
        expectThrows(IllegalArgumentException.class, () -> IngestionMessageMapper.MapperType.fromString("invalid"));
    }

    public void testMapperCreation() {
        IngestionMessageMapper defaultMapper = IngestionMessageMapper.create("default", 0);
        assertNotNull(defaultMapper);
        assertTrue(defaultMapper instanceof DefaultIngestionMessageMapper);

        IngestionMessageMapper rawPayloadMapper = IngestionMessageMapper.create("raw_payload", 0);
        assertNotNull(rawPayloadMapper);
        assertTrue(rawPayloadMapper instanceof RawPayloadIngestionMessageMapper);

        expectThrows(IllegalArgumentException.class, () -> IngestionMessageMapper.create("unknown", 0));
    }

    public void testDefaultMapperWithIdPresent() {
        DefaultIngestionMessageMapper mapper = new DefaultIngestionMessageMapper();
        String payload = "{\"_id\":\"123\",\"_op_type\":\"index\",\"_source\":{\"name\":\"alice\",\"age\":30}}";
        byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);

        IngestionShardPointer pointer = new FakeIngestionSource.FakeIngestionShardPointer(5);
        Message message = new FakeIngestionSource.FakeIngestionMessage(payloadBytes);

        ShardUpdateMessage result = mapper.mapAndProcess(pointer, message);

        assertNotNull(result);
        assertEquals(pointer, result.pointer());
        assertEquals(message, result.originalMessage());
        assertEquals(UNSET_AUTO_GENERATED_TIMESTAMP, result.autoGeneratedIdTimestamp());

        Map<String, Object> parsedMap = result.parsedPayloadMap();
        assertEquals("123", parsedMap.get("_id"));
        assertEquals("index", parsedMap.get("_op_type"));
        assertNotNull(parsedMap.get("_source"));
        assertTrue(parsedMap.get("_source") instanceof Map);

        Map<String, Object> source = (Map<String, Object>) parsedMap.get("_source");
        assertEquals("alice", source.get("name"));
        assertEquals(30, source.get("age"));
    }

    public void testDefaultMapperWithoutId() {
        DefaultIngestionMessageMapper mapper = new DefaultIngestionMessageMapper();
        String payload = "{\"_op_type\":\"index\",\"_source\":{\"name\":\"bob\",\"age\":25}}";
        byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);

        IngestionShardPointer pointer = new FakeIngestionSource.FakeIngestionShardPointer(10);
        Message message = new FakeIngestionSource.FakeIngestionMessage(payloadBytes);

        ShardUpdateMessage result = mapper.mapAndProcess(pointer, message);

        assertNotNull(result);
        assertEquals(pointer, result.pointer());
        assertEquals(message, result.originalMessage());
        assertNotEquals(UNSET_AUTO_GENERATED_TIMESTAMP, result.autoGeneratedIdTimestamp());

        Map<String, Object> parsedMap = result.parsedPayloadMap();
        assertNotNull(parsedMap.get("_id"));
        assertTrue(parsedMap.get("_id") instanceof String);
        String generatedId = (String) parsedMap.get("_id");
        assertFalse(generatedId.isEmpty());
    }

    public void testDefaultMapperWithVersion() {
        DefaultIngestionMessageMapper mapper = new DefaultIngestionMessageMapper();
        String payload = "{\"_id\":\"789\",\"_version\":\"5\",\"_op_type\":\"index\",\"_source\":{\"name\":\"david\"}}";
        byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);

        IngestionShardPointer pointer = new FakeIngestionSource.FakeIngestionShardPointer(20);
        Message message = new FakeIngestionSource.FakeIngestionMessage(payloadBytes);

        ShardUpdateMessage result = mapper.mapAndProcess(pointer, message);

        assertNotNull(result);
        Map<String, Object> parsedMap = result.parsedPayloadMap();
        assertEquals("789", parsedMap.get("_id"));
        assertEquals("5", parsedMap.get("_version"));
    }

    public void testRawPayloadMapper() {
        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper(0);
        String payload = "{\"name\":\"alice\",\"age\":30,\"city\":\"Seattle\"}";
        byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);

        IngestionShardPointer pointer = new FakeIngestionSource.FakeIngestionShardPointer(100);
        Message message = new FakeIngestionSource.FakeIngestionMessage(payloadBytes);

        ShardUpdateMessage result = mapper.mapAndProcess(pointer, message);

        assertNotNull(result);
        assertEquals(pointer, result.pointer());
        assertEquals(message, result.originalMessage());
        assertEquals(UNSET_AUTO_GENERATED_TIMESTAMP, result.autoGeneratedIdTimestamp());

        Map<String, Object> parsedMap = result.parsedPayloadMap();

        // Verify _id is set to shard ID + pointer value
        assertEquals("0-100", parsedMap.get("_id"));

        // Verify _op_type is set to "index"
        assertEquals("index", parsedMap.get("_op_type"));

        // Verify _source contains the original payload
        assertNotNull(parsedMap.get("_source"));
        assertTrue(parsedMap.get("_source") instanceof Map);

        Map<String, Object> source = (Map<String, Object>) parsedMap.get("_source");
        assertEquals("alice", source.get("name"));
        assertEquals(30, source.get("age"));
        assertEquals("Seattle", source.get("city"));

        // Verify _source does not contain metadata fields
        assertFalse(source.containsKey("_id"));
        assertFalse(source.containsKey("_op_type"));
    }

    public void testRawPayloadMapperWithComplexObject() {
        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper(1);
        String payload = "{\"user\":{\"name\":\"bob\",\"email\":\"bob@example.com\"},\"tags\":[\"tag1\",\"tag2\"],\"count\":42}";
        byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);

        IngestionShardPointer pointer = new FakeIngestionSource.FakeIngestionShardPointer(200);
        Message message = new FakeIngestionSource.FakeIngestionMessage(payloadBytes);

        ShardUpdateMessage result = mapper.mapAndProcess(pointer, message);

        assertNotNull(result);
        Map<String, Object> parsedMap = result.parsedPayloadMap();

        assertEquals("1-200", parsedMap.get("_id"));
        assertEquals("index", parsedMap.get("_op_type"));

        Map<String, Object> source = (Map<String, Object>) parsedMap.get("_source");
        assertEquals(3, source.size());
        assertTrue(source.containsKey("user"));
        assertTrue(source.containsKey("tags"));
        assertEquals(42, source.get("count"));
    }

    public void testRawPayloadMapperWithEmptyObject() {
        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper(2);
        String payload = "{}";
        byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);

        IngestionShardPointer pointer = new FakeIngestionSource.FakeIngestionShardPointer(300);
        Message message = new FakeIngestionSource.FakeIngestionMessage(payloadBytes);

        ShardUpdateMessage result = mapper.mapAndProcess(pointer, message);

        assertNotNull(result);
        Map<String, Object> parsedMap = result.parsedPayloadMap();

        assertEquals("2-300", parsedMap.get("_id"));
        assertEquals("index", parsedMap.get("_op_type"));

        Map<String, Object> source = (Map<String, Object>) parsedMap.get("_source");
        assertTrue(source.isEmpty());
    }

    public void testDefaultMapperWithInvalidJson() {
        DefaultIngestionMessageMapper mapper = new DefaultIngestionMessageMapper();
        String payload = "invalid json {{{";
        byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);

        IngestionShardPointer pointer = new FakeIngestionSource.FakeIngestionShardPointer(500);
        Message message = new FakeIngestionSource.FakeIngestionMessage(payloadBytes);

        expectThrows(Exception.class, () -> mapper.mapAndProcess(pointer, message));
    }

    public void testRawPayloadMapperWithInvalidJson() {
        RawPayloadIngestionMessageMapper mapper = new RawPayloadIngestionMessageMapper(3);
        String payload = "not a json";
        byte[] payloadBytes = payload.getBytes(StandardCharsets.UTF_8);

        IngestionShardPointer pointer = new FakeIngestionSource.FakeIngestionShardPointer(600);
        Message message = new FakeIngestionSource.FakeIngestionMessage(payloadBytes);

        expectThrows(Exception.class, () -> mapper.mapAndProcess(pointer, message));
    }
}
