from math import cos, radians, sin
from flask import Flask, render_template, request, send_file, jsonify
from PIL import Image, ImageDraw, ImageOps
import gradio as gr
import os
import io
import subprocess
import cairosvg

app = Flask(__name__)

UPLOAD_FOLDER = 'uploads'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

def preprocess_image(image):
    # Ensure the image is square by padding it
    size = max(image.size)
    new_image = Image.new("RGBA", (size, size), (255, 255, 255, 0))
    new_image.paste(image, ((size - image.width) // 2, (size - image.height) // 2))
    return new_image

def convert_webp_to_jpeg(image_data):
    try:
        image = Image.open(io.BytesIO(image_data))
        if image.format == 'WEBP':
            output = io.BytesIO()
            image.convert('RGB').save(output, format='JPEG')
            return output.getvalue()
        return image_data
    except Exception as e:
        raise ValueError(f"Error converting WEBP to JPEG: {e}")

def convert_svg_to_png(image_data):
    try:
        png_data = cairosvg.svg2png(bytestring=image_data)
        return png_data
    except Exception as e:
        raise ValueError(f"Error converting SVG to PNG: {e}")

def create_gif(images):
    frames = []
    duration = 100  # Duration for each frame in milliseconds
    total_frames = 18  # Total number of frames

    try:
        img1_data = images[0].read()
        img2_data = images[1].read()

        # Convert WebP images to JPEG if necessary
        img1_data = convert_webp_to_jpeg(img1_data)
        img2_data = convert_webp_to_jpeg(img2_data)

        # Convert SVG images to PNG if necessary
        if img1_data.strip().startswith(b'<svg'):
            img1_data = convert_svg_to_png(img1_data)
        if img2_data.strip().startswith(b'<svg'):
            img2_data = convert_svg_to_png(img2_data)

        # Open images
        img1 = Image.open(io.BytesIO(img1_data)).convert('RGBA')
        img2 = Image.open(io.BytesIO(img2_data)).convert('RGBA')

        # Preprocess images to make them square and same size
        img1 = preprocess_image(img1)
        img2 = preprocess_image(img2)

        # Set size for the GIF
        size = (256, 256)
        img1 = img1.resize(size, Image.LANCZOS)
        img2 = img2.resize(size, Image.LANCZOS)

        # Calculate step size for consistent speed
        full_width = size[0]
        step = full_width // (total_frames // 2)  # Divide by 2 as we have 2 parts to the animation

        # Generate frames from left to right, reversing direction 32 pixels from the edges
        for i in range(0, full_width, step):
            frame = Image.new('RGBA', size)
            frame.paste(img1, (0, 0))
            frame.paste(img2.crop((i, 0, full_width, size[1])), (i, 0), mask=img2.crop((i, 0, full_width, size[1])))

            draw = ImageDraw.Draw(frame)
            draw.line((i, 0, i, size[1]), fill=(0, 255, 0), width=2)

            # Convert frame to P mode which is a palette-based image
            frame = frame.convert('P', palette=Image.ADAPTIVE)
            frames.append(frame)

        # Generate frames from right to left, reversing direction 32 pixels from the edges
        for i in range(full_width, step, -step):
            frame = Image.new('RGBA', size)
            frame.paste(img1, (0, 0))
            frame.paste(img2.crop((i, 0, full_width, size[1])), (i, 0), mask=img2.crop((i, 0, full_width, size[1])))

            draw = ImageDraw.Draw(frame)
            draw.line((i, 0, i, size[1]), fill=(0, 255, 0), width=2)

            # Convert frame to P mode which is a palette-based image
            frame = frame.convert('P', palette=Image.ADAPTIVE)
            frames.append(frame)

        # Save as GIF
        output = io.BytesIO()
        frames[0].save(output, format='GIF', save_all=True, append_images=frames[1:], duration=duration, loop=0, optimize=True)
        output.seek(0)

        return output
    except Exception as e:
        raise ValueError(f"Error creating GIF: {e}")

def create_rotating_gif(images):
    frames = []
    duration = 100  # Duration for each frame in milliseconds
    total_frames = 18  # Total number of frames for a smooth rotation

    try:
        img1_data = images[0].read()
        img2_data = images[1].read()

        # Convert WebP images to JPEG if necessary
        img1_data = convert_webp_to_jpeg(img1_data)
        img2_data = convert_webp_to_jpeg(img2_data)

        # Convert SVG images to PNG if necessary
        if img1_data.strip().startswith(b'<svg'):
            img1_data = convert_svg_to_png(img1_data)
        if img2_data.strip().startswith(b'<svg'):
            img2_data = convert_svg_to_png(img2_data)

        # Open images
        img1 = Image.open(io.BytesIO(img1_data)).convert('RGBA')
        img2 = Image.open(io.BytesIO(img2_data)).convert('RGBA')

        # Preprocess images to make them square and same size
        img1 = preprocess_image(img1)
        img2 = preprocess_image(img2)

        # Set size for the GIF
        size = (256, 256)
        img1 = img1.resize(size, Image.LANCZOS)
        img2 = img2.resize(size, Image.LANCZOS)

        # Create a mask that is 2x the size of the image
        mask_size = (size[0] * 2, size[1] * 2)
        mask = Image.new('L', mask_size, 0)
        draw = ImageDraw.Draw(mask)
        draw.rectangle([size[0], 0, mask_size[0], mask_size[1]], fill=255)

        center_x, center_y = size[0] // 2, size[1] // 2

        for angle in range(0, 360, 360 // total_frames):
            rotated_mask = mask.rotate(angle, center=(mask_size[0] // 2, mask_size[1] // 2), expand=False)
            cropped_mask = rotated_mask.crop((size[0] // 2, size[1] // 2, size[0] // 2 + size[0], size[1] // 2 + size[1]))
            frame = Image.composite(img1, img2, cropped_mask)

            # Draw the green dividing lines in the reverse direction
            draw = ImageDraw.Draw(frame)
            reverse_angle = -angle + 90  # Reverse the direction of the line and add 90 degrees
            end_x1 = center_x + int(size[0] * 1.5 * cos(radians(reverse_angle)))
            end_y1 = center_y + int(size[1] * 1.5 * sin(radians(reverse_angle)))
            end_x2 = center_x - int(size[0] * 1.5 * cos(radians(reverse_angle)))
            end_y2 = center_y - int(size[1] * 1.5 * sin(radians(reverse_angle)))
            draw.line([center_x, center_y, end_x1, end_y1], fill=(0, 255, 0), width=3)  # Line in one direction
            draw.line([center_x, center_y, end_x2, end_y2], fill=(0, 255, 0), width=3)  # Mirrored tail

            # Convert frame to P mode which is a palette-based image
            frame = frame.convert('P', palette=Image.ADAPTIVE)
            frames.append(frame)

        # Save as GIF
        output = io.BytesIO()
        frames[0].save(output, format='GIF', save_all=True, append_images=frames[1:], duration=duration, loop=0, optimize=True)
        output.seek(0)

        return output
    except Exception as e:
        raise ValueError(f"Error creating rotating GIF: {e}")

def compress_gif(input_path, output_path, lossy=80, colors=256):
    # Optimize the GIF using gifsicle with lossy compression
    subprocess.run(['gifsicle', '--lossy={}'.format(lossy), '--colors', str(colors), input_path, '-o', output_path])
    print(f"Compressed GIF saved to {output_path}")

def create_gif_gradio(image1, image2, transition_type):
    # Convert Gradio image inputs to file-like objects
    img1_io = io.BytesIO()
    img2_io = io.BytesIO()
    image1.save(img1_io, format='PNG')
    image2.save(img2_io, format='PNG')
    img1_io.seek(0)
    img2_io.seek(0)

    images = [img1_io, img2_io]

    if transition_type == 'rotate':
        gif_data = create_rotating_gif(images)
    else:
        gif_data = create_gif(images)

    # Save the GIF to a temporary file
    temp_input_path = "temp_input.gif"
    with open(temp_input_path, "wb") as f:
        f.write(gif_data.getbuffer())

    # Compress the GIF
    compressed_output_path = "compressed_output.gif"
    compress_gif(temp_input_path, compressed_output_path)

    return compressed_output_path

# Gradio interface
iface = gr.Interface(
    fn=create_gif_gradio,
    inputs=[
        gr.Image(type="pil", label="Image 1"),
        gr.Image(type="pil", label="Image 2"),
        gr.Radio(["default", "rotate"], label="Transition Type")
    ],
    outputs=gr.Image(type="filepath", label="Generated GIF"),
    title="GIF Generator",
    description="Upload two images and select a transition type to generate an animated GIF."
)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/gradio', methods=['GET', 'POST'])
def gradio_interface():
    return gr.routes.App.create_app(iface)()

@app.route('/upload_image', methods=['POST'])
def upload_image():
    try:
        image = request.files['image']
        if image:
            filepath = os.path.join(UPLOAD_FOLDER, image.filename)
            image.save(filepath)
            return jsonify({'success': True, 'filepath': filepath})
        else:
            return jsonify({'success': False, 'error': 'No image uploaded.'})
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})

@app.route('/generate_gif', methods=['POST'])
def generate_gif():
    try:
        images = request.files.getlist('images')
        if len(images) != 2:
            raise ValueError("Two images are required.")

        # Determine the type of transition
        transition_type = request.form.get('transition_type', 'default')

        # Create the GIF based on the transition type
        if transition_type == 'rotate':
            gif_data = create_rotating_gif(images)
        else:
            gif_data = create_gif(images)

        # Save the GIF to a temporary file
        temp_input_path = "temp_input.gif"
        with open(temp_input_path, "wb") as f:
            f.write(gif_data.getbuffer())

        # Compress the GIF
        compressed_output_path = "compressed_output.gif"
        compress_gif(temp_input_path, compressed_output_path)

        # Send the compressed GIF as the response
        return send_file(compressed_output_path, mimetype='image/gif', as_attachment=True, download_name='animation.gif')
    except Exception as e:
        return jsonify({'error': str(e)}), 400

if __name__ == '__main__':
    app.run(debug=True)