import gradio as gr
from PIL import Image
import piexif
import piexif.helper
from datetime import datetime
import os

def add_360_metadata(input_image):
    """Add 360 photo metadata to an image file."""
    try:
        # Open the image
        img = Image.open(input_image)
        
        # Prepare XMP metadata for 360 photo
        xmp_metadata = (
            '<?xpacket begin="" id="W5M0MpCehiHzreSzNTczkc9d"?>\n'
            '<x:xmpmeta xmlns:x="adobe:ns:meta/">\n'
            '<rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">\n'
            '<rdf:Description rdf:about="" xmlns:GPano="http://ns.google.com/photos/1.0/panorama/">\n'
            '<GPano:ProjectionType>equirectangular</GPano:ProjectionType>\n'
            '<GPano:UsePanoramaViewer>True</GPano:UsePanoramaViewer>\n'
            '<GPano:CroppedAreaImageWidthPixels>' + str(img.width) + '</GPano:CroppedAreaImageWidthPixels>\n'
            '<GPano:CroppedAreaImageHeightPixels>' + str(img.height) + '</GPano:CroppedAreaImageHeightPixels>\n'
            '<GPano:FullPanoWidthPixels>' + str(img.width) + '</GPano:FullPanoWidthPixels>\n'
            '<GPano:FullPanoHeightPixels>' + str(img.height) + '</GPano:FullPanoHeightPixels>\n'
            '<GPano:CroppedAreaLeftPixels>0</GPano:CroppedAreaLeftPixels>\n'
            '<GPano:CroppedAreaTopPixels>0</GPano:CroppedAreaTopPixels>\n'
            '</rdf:Description>\n'
            '</rdf:RDF>\n'
            '</x:xmpmeta>\n'
            '<?xpacket end="w"?>'
        )

        # Create temporary filename for output
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_filename = f"360_photo_{timestamp}.jpg"
        
        # Create EXIF dictionary with XMP metadata
        exif_dict = {"0th": {}, "Exif": {}, "GPS": {}, "1st": {}, "thumbnail": None}
        exif_dict["0th"][piexif.ImageIFD.XPComment] = xmp_metadata.encode('utf-16')
        
        # Convert EXIF dictionary to bytes
        exif_bytes = piexif.dump(exif_dict)
        
        # Save image with metadata
        output_path = os.path.join("/tmp", output_filename)
        img.save(output_path, "JPEG", exif=exif_bytes, quality=95)
        
        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 necessary metadata for Google Photos and other 360° viewers. Your image should have a 2:1 aspect ratio for best results.",
    examples=[],
    cache_examples=False
)

# Launch the interface
iface.launch()