#!/bin/bash

. ./test_shared_functions.sh

JSON_FILE=./test_data/Instituto_Superior_Tecnico.md.json
HTML_FILE=./test_data/AMERICA-VENTURE_HIGHWAY.html
MP3_FILE=./test_data/flight_of_the_bumblebee_2.mp3
JPG_FILE=./test_data/China_Beijing_Summer_Palace2.jpg
EXCEL_FILE=./test_data/Worldbank-Datasheet-02082022.xlsx
TEXT_FILE=./test_data/America-Horse_with_no_name.txt
OPEN_OFFICE_FILE=./test_data/GAR_ARchitecture.odg

echo "# 1. Create Document Tests"
echo "=== Testing Document Creation ==="

echo "# Test 1.1: Create document with file, use different formats"
make_request "POST" "/documents" \
    "-F file=@${TEXT_FILE}" \
    "-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_FILE=${DOCUMENT_ID}

make_request "POST" "/documents" \
    "-F file=@${JSON_FILE}" \
    "-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_JSON_FILE=${DOCUMENT_ID}

make_request "POST" "/documents" \
    "-F file=@${HTML_FILE}" \
    "-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_HTML_FILE=${DOCUMENT_ID}

make_request "POST" "/documents" \
    "-F file=@${MP3_FILE}" \
    "-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_MP3_FILE=${DOCUMENT_ID}

make_request "POST" "/documents" \
    "-F file=@${JPG_FILE}" \
    "-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_JPG_FILE=${DOCUMENT_ID}

make_request "POST" "/documents" \
    "-F file=@${EXCEL_FILE}" \
    "-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_EXCEL_FILE=${DOCUMENT_ID}

make_request "POST" "/documents" \
    "-F file=@${OPEN_OFFICE_FILE}" \
    "-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_OPEN_OFFICE_FILE=${DOCUMENT_ID}

echo "# Test 1.2: Create document with raw text"
make_request "POST" "/documents" \
    '-F raw_text="On the first part of the journey
I was looking at all the life
There were plants and birds and rocks and things
There was sand and hills and rings
The first thing I met was a fly with a buzz
And the sky with no clouds
The heat was hot and the ground was dry
But the air was full of sound"' \
    "-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_RAW_TEXT=${DOCUMENT_ID}

echo "# Test 1.3: Create document with chunks"
make_request "POST" "/documents" \
    "-F chunks=\"[\\\"I've been through the desert on a horse with no name\\\",\\\"It felt good to be out of the rain\\\"]\"" \
    "-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_CHUNKS=${DOCUMENT_ID}

echo "# Test 1.4: Create document with metadata"
make_request "POST" "/documents" \
    "-F file=@${TEXT_FILE} -F metadata='{\"title\":\"Horse with No Name\",\"author\":\"America\"}'" \
    "-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_FILE_WITH_METADATA=${DOCUMENT_ID}

echo "# Test 1.5: Create document with custom ingestion config"
make_request "POST" "/documents" \
    "-F file=@${TEXT_FILE} -F ingestion_mode=hi-res -F ingestion_config='{\"chunk_size\":1000}'" \
    "-H 'Content-Type: multipart/form-data'"
DOCUMENT_ID_FROM_FILE_WITH_CUSTOM_INGEST=${DOCUMENT_ID}

echo "# 2. List Documents Tests"
echo "=== Testing Document Listing ==="

echo "# Test 2.1: List all documents"
make_request "GET" "/documents" ""

echo "# Test 2.2: List documents with pagination"
make_request "GET" "/documents?offset=0&limit=10" ""

echo "# Test 2.3: List specific documents by IDs"
make_request "GET" "/documents?ids=${DOCUMENT_ID_FROM_CHUNKS},${DOCUMENT_ID_FROM_FILE_WITH_METADATA}" ""

echo "# Test 2.4: List documents with summary embeddings"
make_request "GET" "/documents?include_summary_embeddings=1" ""

echo "# 3. Get Document Tests"
echo "=== Testing Document Retrieval ==="

echo "# Test 3.1: Get document by ID"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}" ""

echo "# Test 3.2: Get non-existent document"
make_request "GET" "/documents/00000000-0000-0000-0000-000000000000" ""

echo "# 4. List Document Chunks Tests"
echo "=== Testing Document Chunks ==="

echo "# Test 4.1: List chunks with pagination"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_CHUNKS}/chunks?offset=0&limit=10" ""

echo "# Test 4.2: List chunks with vectors"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_CHUNKS}/chunks?include_vectors=true" ""

echo "# 5. Download Document Tests"
echo "=== Testing Document Download ==="

echo "# Test 5.1: Download document"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE}/download" ""

echo "# 6. Delete Document Tests"
echo "=== Testing Document Deletion ==="

echo "# Test 6.1: Delete document by ID"
make_request "DELETE" "/documents/${DOCUMENT_ID_FROM_RAW_TEXT}" ""

echo "# Test 6.2: Delete documents by filter"
make_request "DELETE" "/documents/by-filter" \
    "-d '{\"document_type\":{\"\$eq\":\"txt\"}}'"

echo "# 7. Document Collections Tests"
echo "=== Testing Document Collections ==="

echo "# Test 7.1: List document collections"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_CHUNKS}/collections?offset=0&limit=10" ""

echo "# 8. Entity Extraction Tests"
echo "=== Testing Entity Extraction ==="

echo "# Test 8.1: Extract entities with default settings"
make_request "POST" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/extract" \
    "-d '{\"run_type\":\"RUN\"}'"

echo "# Test 8.2: Extract entities with custom settings"
make_request "POST" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/extract" \
    "-d '{\"run_type\":\"RUN\",\"settings\":{\"model\":\"gpt-4\"}}'"

echo "# 9. Entity Retrieval Tests"
echo "=== Testing Entity Retrieval ==="

echo "# Test 9.1: Get entities with pagination"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/entities?offset=0&limit=10" ""

echo "# Test 9.2: Get entities with embeddings"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/entities?include_embeddings=true" ""

echo "# 10. Relationship Tests"
echo "=== Testing Relationships ==="

echo "# Test 10.1: Get relationships with pagination"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/relationships?offset=0&limit=10" ""

echo "# Test 10.2: Get relationships with filters"
make_request "GET" "/documents/${DOCUMENT_ID_FROM_FILE_WITH_METADATA}/relationships?entity_names=person,company&relationship_types=works_for" ""

echo "# 11. Document Search Tests"
echo "=== Testing Document Search ==="

echo "# Test 11.1: Basic search"
make_request "POST" "/documents/search" \
    "-d '{\"query\":\"test query\",\"search_mode\":\"basic\"}'"

echo "# Test 11.2: Advanced search with custom settings"
make_request "POST" "/documents/search" \
    "-d '{\"query\":\"test query\",\"search_mode\":\"advanced\",\"search_settings\":{\"filters\":{\"document_type\":\"pdf\"}}}'"

echo "# Test 11.3: Custom search with complex filters"
make_request "POST" "/documents/search" \
    "-d '{\"query\":\"test query\",\"search_mode\":\"custom\",\"search_settings\":{\"filters\":{\"created_at\":{\"\$gt\":\"2023-01-01\"}}}}'"
