import io
import gradio as gr
from PIL import Image
import vtracer
import tempfile
import boto3
import os
import requests
from botocore.exceptions import NoCredentialsError

# Access environment variables
AWS_ACCESS_KEY = os.getenv('AWS_ACCESS_KEY_ID')
AWS_SECRET_KEY = os.getenv('AWS_SECRET_ACCESS_KEY')
AWS_REGION = os.getenv('AWS_DEFAULT_REGION')
BUCKET_NAME = os.getenv('BUCKET_NAME')


def upload_to_aws(local_file, s3_file):
    """Uploads a file to AWS S3 and returns the URL."""
    s3 = boto3.client('s3', aws_access_key_id=AWS_ACCESS_KEY, aws_secret_access_key=AWS_SECRET_KEY)
    try:
        s3.upload_file(local_file, BUCKET_NAME, s3_file)
        print("Upload Successful")
        return f"https://{BUCKET_NAME}.s3.amazonaws.com/{s3_file}"
    except FileNotFoundError:
        print("The file was not found")
        return None
    except NoCredentialsError:
        print("Credentials not available")
        return None

def convert_image(image_input, color_mode, hierarchical, mode, filter_speckle,
                  color_precision, layer_difference, corner_threshold,
                  length_threshold, max_iterations, splice_threshold, path_precision):
    """Converts an image from a URL or file path to SVG using vtracer with customizable parameters."""
    
    try:
        # Debugging: Print the received input and its type
        print(f"Received image_input: {image_input}, type: {type(image_input)}")
        
        # Handle both URL strings and Gradio FileData objects
        if isinstance(image_input, dict) and 'path' in image_input:
            # If the input is a Gradio FileData object, use the file path
            image_path = image_input['path']
            with open(image_path, 'rb') as f:
                image_data = f.read()
        elif isinstance(image_input, str):
            # If the input is a URL, fetch the image
            if image_input.startswith(('http://', 'https://')):
                response = requests.get(image_input)
                response.raise_for_status()  # Raise an error for bad status codes
                image_data = response.content
            else:
                return gr.HTML(f'<p style="color: red;">Error: Invalid URL format. Expected a valid HTTP/HTTPS URL.</p>'), None, None
        else:
            return gr.HTML(f'<p style="color: red;">Error: Invalid input type. Expected a URL or file path, got {type(image_input)}</p>'), None, None

        # Load the image using PIL
        image = Image.open(io.BytesIO(image_data))
        print(f"Image format: {image.format}, Size: {image.size}, Mode: {image.mode}")
        image.show()  # Display the image for debugging

        # Convert the image to bytes for vtracer compatibility
        img_byte_array = io.BytesIO()
        img_format = image.format if image.format else 'PNG'  # Default to PNG for lossless quality
        image.save(img_byte_array, format=img_format)
        img_bytes = img_byte_array.getvalue()

        # Debugging: Save the input image to a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_img_file:
            image.save(temp_img_file, format='PNG')
            print(f"Saved input image to: {temp_img_file.name}")

        # Debugging: Print vtracer parameters
        print(f"vtracer parameters: img_format={img_format}, colormode={color_mode}, hierarchical={hierarchical}, mode={mode}, filter_speckle={filter_speckle}, color_precision={color_precision}, layer_difference={layer_difference}, corner_threshold={corner_threshold}, length_threshold={length_threshold}, max_iterations={max_iterations}, splice_threshold={splice_threshold}, path_precision={path_precision}")

        # Perform the conversion
        svg_str = vtracer.convert_raw_image_to_svg(
            img_bytes,
            img_format=img_format.lower(),
            colormode=color_mode.lower(),
            hierarchical=hierarchical.lower(),
            mode=mode.lower(),
            filter_speckle=int(filter_speckle),
            color_precision=int(color_precision),
            layer_difference=int(layer_difference),
            corner_threshold=int(corner_threshold),
            length_threshold=float(length_threshold),
            max_iterations=int(max_iterations),
            splice_threshold=int(splice_threshold),
            path_precision=int(path_precision)
        )
        print(f"Generated SVG: {svg_str}")  # Debugging: Print the generated SVG
    except Exception as e:
        return gr.HTML(f'<p style="color: red;">Error during image processing or SVG conversion: {str(e)}</p>'), None, None

    # Save the SVG string to a temporary file
    with tempfile.NamedTemporaryFile(delete=False, suffix='.svg', mode='w', encoding='utf-8') as temp_file:
        temp_file.write(svg_str)
        temp_file_path = temp_file.name

    try:
        # Upload the SVG file to AWS S3
        s3_file_name = temp_file_path.split('/')[-1]  # Use the temp file name as the S3 file name
        file_url = upload_to_aws(temp_file_path, s3_file_name)
    except Exception as e:
        return gr.HTML(f'<p style="color: red;">Error uploading to AWS: {str(e)}</p>'), None, None

    # Clean up the temporary file
    # os.unlink(temp_file_path)

    # Return the SVG wrapped in HTML and the file URL
    return gr.HTML(f'<svg viewBox="0 0 {image.width} {image.height}">{svg_str}</svg>'), file_url, svg_str
    
    
# Gradio interface
iface = gr.Blocks()

with iface:
    gr.Interface(
        fn=convert_image,
        inputs=[
            gr.Textbox(label="Image URL"),
            gr.Radio(choices=["Color", "Binary"], value="Color", label="Color Mode"),
            gr.Radio(choices=["Stacked", "Cutout"], value="Stacked", label="Hierarchical"),
            gr.Radio(choices=["Spline", "Polygon", "None"], value="Spline", label="Mode"),
            gr.Slider(minimum=1, maximum=10, value=4, step=1, label="Filter Speckle"),
            gr.Slider(minimum=1, maximum=8, value=6, step=1, label="Color Precision"),
            gr.Slider(minimum=1, maximum=32, value=16, step=1, label="Layer Difference"),
            gr.Slider(minimum=10, maximum=90, value=60, step=1, label="Corner Threshold"),
            gr.Slider(minimum=3.5, maximum=10, value=4.0, step=0.5, label="Length Threshold"),
            gr.Slider(minimum=1, maximum=20, value=10, step=1, label="Max Iterations"),
            gr.Slider(minimum=10, maximum=90, value=45, step=1, label="Splice Threshold"),
            gr.Slider(minimum=1, maximum=10, value=8, step=1, label="Path Precision")
        ],
        outputs=[
            gr.HTML(label="SVG Output"),
            gr.Textbox(label="SVG File URL"),
            gr.Textbox(label="StringData")
        ],
        title="Convert Image to SVG vectors",
        description="Upload an image and customize the conversion parameters as needed.<br><h2>Support me USDT (TRC-20): TAe7hsSVWtMEYz3G5V1UiUdYPQVqm28bKx</h2>",
    )

iface.launch()