#!/usr/bin/env python3
"""Test script for SeaweedFS image upload functionality.

This script tests the complete flow:
1. Load sample chunks with image_base64 data
2. Upload images to SeaweedFS
3. Verify metadata is updated with image_url and image_s3_key
4. Verify image_base64 is removed after upload
"""

import sys
import os
from uuid import uuid4
import base64
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.services.seaweedfs_service import SeaweedFSService


def create_test_image_base64():
    """Create a small test PNG image (1x1 red pixel)."""
    # 1x1 red pixel PNG
    png_bytes = base64.b64decode(
        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
    )
    return base64.b64encode(png_bytes).decode('utf-8')


def create_test_chunks():
    """Create test chunks with image data."""
    test_image = create_test_image_base64()

    chunks = [
        # Text chunk (should be skipped)
        {
            'content': {'text': 'This is a text chunk'},
            'headers': ['text'],
            '_source': {
                'document_id': str(uuid4()),
                'filename': 'test.xlsx',
                'sheet': 'Sheet1',
                'table_id': 0,
                'row': 1,
                'col_range': 'A1:A1',
            },
            '_legacy_metadata': {
                'element_type': 'Text',
                'full_metadata': {
                    'sheet_name': 'Sheet1',
                    'sheet_number': 1,
                }
            }
        },
        # Image chunk 1
        {
            'content': {
                'description': 'Test image 1',
                'image_format': 'png',
            },
            'headers': ['description', 'image_format'],
            '_source': {
                'document_id': str(uuid4()),
                'filename': 'test.xlsx',
                'sheet': 'Sheet1',
                'table_id': 0,
                'row': 5,
                'col_range': 'C5:E10',
            },
            '_legacy_metadata': {
                'element_type': 'Image',
                'full_metadata': {
                    'image_base64': test_image,
                    'mime_type': 'image/png',
                    'image_number': 1,
                    'position_key': 'C5',
                    'sheet_name': 'Sheet1',
                    'width': 100,
                    'height': 100,
                }
            }
        },
        # Image chunk 2
        {
            'content': {
                'description': 'Test image 2',
                'image_format': 'png',
            },
            'headers': ['description', 'image_format'],
            '_source': {
                'document_id': str(uuid4()),
                'filename': 'test.xlsx',
                'sheet': 'Data',
                'table_id': 0,
                'row': 10,
                'col_range': 'F10:H15',
            },
            '_legacy_metadata': {
                'element_type': 'Image',
                'full_metadata': {
                    'image_base64': test_image,
                    'mime_type': 'image/png',
                    'image_number': 2,
                    'position_key': 'F10',
                    'sheet_name': 'Data',
                    'width': 200,
                    'height': 150,
                }
            }
        },
    ]

    return chunks


def test_seaweedfs_upload():
    """Test SeaweedFS upload functionality."""
    print("=" * 70)
    print("SeaweedFS Image Upload Test")
    print("=" * 70)

    # Create test chunks
    print("\n1. Creating test chunks with image data...")
    chunks = create_test_chunks()
    print(f"   ✓ Created {len(chunks)} test chunks")

    # Count images before upload
    image_count_before = sum(
        1 for chunk in chunks
        if chunk.get('_legacy_metadata', {}).get('element_type') == 'Image'
        and 'image_base64' in chunk.get('_legacy_metadata', {}).get('full_metadata', {})
    )
    print(f"   ✓ Found {image_count_before} images with base64 data")

    # Create test document ID
    document_id = uuid4()
    print(f"\n2. Test document ID: {document_id}")

    # Initialize SeaweedFS service
    print("\n3. Initializing SeaweedFS service...")
    try:
        service = SeaweedFSService()
        print("   ✓ SeaweedFS service initialized")
    except Exception as e:
        print(f"   ✗ Failed to initialize SeaweedFS service: {e}")
        print("\n   Make sure SeaweedFS is running and AWS credentials are set:")
        print("   - AWS_ACCESS_KEY_ID")
        print("   - AWS_SECRET_ACCESS_KEY")
        print("   - AWS_ENDPOINT_URL")
        return False

    # Upload images
    print("\n4. Uploading images to SeaweedFS...")
    try:
        updated_chunks = service.upload_images_batch(
            chunks=chunks,
            document_id=document_id,
            bucket='textiq-images'
        )
        print("   ✓ Upload batch completed")
    except Exception as e:
        print(f"   ✗ Upload failed: {e}")
        return False

    # Verify results
    print("\n5. Verifying upload results...")

    uploaded_count = 0
    for i, chunk in enumerate(updated_chunks):
        legacy_meta = chunk.get('_legacy_metadata', {})
        full_metadata = legacy_meta.get('full_metadata', {})
        element_type = legacy_meta.get('element_type', 'Text')

        if element_type == 'Image':
            print(f"\n   Image chunk {i + 1}:")

            # Check if image_base64 was removed
            has_base64 = 'image_base64' in full_metadata
            print(f"      - image_base64 removed: {'✗ NO' if has_base64 else '✓ YES'}")

            # Check if image_s3_key was added
            has_s3_key = 'image_s3_key' in full_metadata
            s3_key = full_metadata.get('image_s3_key', '')
            print(f"      - image_s3_key added: {'✓ YES' if has_s3_key else '✗ NO'}")
            if s3_key:
                print(f"        → {s3_key}")

            # Check if image_url was added
            has_url = 'image_url' in full_metadata
            url = full_metadata.get('image_url', '')
            print(f"      - image_url added: {'✓ YES' if has_url else '✗ NO'}")
            if url:
                print(f"        → {url}")

            if has_s3_key and has_url and not has_base64:
                uploaded_count += 1

    # Summary
    print("\n" + "=" * 70)
    print(f"Upload Summary: {uploaded_count}/{image_count_before} images successfully uploaded")

    if uploaded_count == image_count_before:
        print("✓ All images uploaded successfully!")
        print("=" * 70)
        return True
    else:
        print("✗ Some images failed to upload")
        print("=" * 70)
        return False


if __name__ == "__main__":
    success = test_seaweedfs_upload()
    sys.exit(0 if success else 1)