import gradio as gr
from PIL import Image
import os
from datetime import datetime
import struct

def create_xmp_block(width, height):
    """Create XMP metadata block following ExifTool's exact format."""
    xmp = (
        f'<?xpacket begin="" id="W5M0MpCehiHzreSzNTczkc9d"?>\n'
        f'<x:xmpmeta xmlns:x="adobe:ns:meta/" x:xmptk="ExifTool">\n'
        f'<rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">\n'
        f'<rdf:Description rdf:about=""\n'
        f'xmlns:GPano="http://ns.google.com/photos/1.0/panorama/"\n'
        f'GPano:ProjectionType="equirectangular"\n'
        f'GPano:UsePanoramaViewer="True"\n'
        f'GPano:FullPanoWidthPixels="{width}"\n'
        f'GPano:FullPanoHeightPixels="{height}"\n'
        f'GPano:CroppedAreaImageWidthPixels="{width}"\n'
        f'GPano:CroppedAreaImageHeightPixels="{height}"\n'
        f'GPano:CroppedAreaLeftPixels="0"\n'
        f'GPano:CroppedAreaTopPixels="0"/>\n'
        f'</rdf:RDF>\n'
        f'</x:xmpmeta>\n'
        f'<?xpacket end="w"?>'
    )
    return xmp

def write_xmp_to_jpg(input_path, output_path, width, height):
    """Write XMP metadata to JPEG file following ExifTool's method."""
    # Read the original JPEG
    with open(input_path, 'rb') as f:
        data = f.read()
    
    # Find the start of image marker
    if data[0:2] != b'\xFF\xD8':
        raise ValueError("Not a valid JPEG file")
    
    # Create XMP data
    xmp_data = create_xmp_block(width, height)
    
    # Create APP1 segment for XMP
    app1_marker = b'\xFF\xE1'
    xmp_header = b'http://ns.adobe.com/xap/1.0/\x00'
    xmp_bytes = xmp_data.encode('utf-8')
    length = len(xmp_header) + len(xmp_bytes) + 2  # +2 for length bytes
    length_bytes = struct.pack('>H', length)
    
    # Construct new file content
    output = bytearray()
    output.extend(data[0:2])  # SOI marker
    output.extend(app1_marker)
    output.extend(length_bytes)
    output.extend(xmp_header)
    output.extend(xmp_bytes)
    output.extend(data[2:])  # Rest of the original file
    
    # Write the new file
    with open(output_path, 'wb') as f:
        f.write(output)

def add_360_metadata(input_image):
    """Add 360 photo metadata to an image file."""
    try:
        # Open and verify the image
        img = Image.open(input_image)
        if img.width != 2 * img.height:
            raise gr.Error("Image must have 2:1 aspect ratio for equirectangular projection")
        
        # Create output filename
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_filename = f"360_photo_{timestamp}.jpg"
        output_path = os.path.join("/tmp", output_filename)
        
        # First save as high-quality JPEG
        img.save(output_path, "JPEG", quality=95)
        
        # Then inject XMP metadata directly into JPEG file
        write_xmp_to_jpg(output_path, output_path, img.width, img.height)
        
        return output_path
    
    except Exception as e:
        raise gr.Error(f"Error processing image: {str(e)}")

# Create Gradio interface
iface = gr.Interface(
    fn=add_360_metadata,
    inputs=gr.Image(type="filepath", label="Upload 360° Photo"),
    outputs=gr.Image(type="filepath", label="360° Photo with Metadata"),
    title="360° Photo Metadata Adder",
    description=(
        "Upload an equirectangular 360° photo to add metadata for Google Photos and other 360° viewers.\n"
        "Important: Image must have 2:1 aspect ratio (width = 2 × height)."
    ),
    examples=[],
    cache_examples=False
)

# Launch the interface
iface.launch()