import streamlit as st
from streamlit.elements import image as st_image
import base64
import io
from PIL import Image, ImageDraw, ImageFont, ImageFilter, ImageEnhance, ImageOps, ImageChops, ImageColor
import numpy as np
import os
import json
from rembg import remove
import requests
from pathlib import Path
from streamlit_drawable_canvas import st_canvas

# ------------------------------
# Monkey-Patch to Add image_to_url
# ------------------------------

def image_to_url(image, *args, **kwargs):
    """
    Converts a PIL Image to a data URL.
    """
    buffered = io.BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/png;base64,{img_str}"

# Monkey-patch the image_to_url method
st_image.image_to_url = image_to_url

# ------------------------------
# Utility Functions
# ------------------------------

def download_file(url, dest_path):
    """
    Downloads a file from a URL to the specified destination path.
    """
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()
        with open(dest_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        st.info(f"Downloaded: {dest_path.name}")
    except Exception as e:
        st.error(f"Failed to download {url}: {e}")

def setup_fonts():
    """
    Ensures that all required fonts are downloaded and available locally.
    Uses open-source fonts from Google Fonts.
    """
    fonts_dir = Path("fonts")
    fonts_dir.mkdir(exist_ok=True)

    # Define open-source fonts to download
    font_info = {
        "Roboto": {
            "name": "Roboto-Regular.ttf",
            "url": "https://github.com/google/fonts/raw/main/apache/roboto/Roboto-Regular.ttf"
        },
        "OpenSans": {
            "name": "OpenSans-Regular.ttf",
            "url": "https://github.com/google/fonts/raw/main/apache/opensans/OpenSans-Regular.ttf"
        },
        "Lato": {
            "name": "Lato-Regular.ttf",
            "url": "https://github.com/google/fonts/raw/main/ofl/lato/Lato-Regular.ttf"
        },
        "Montserrat": {
            "name": "Montserrat-Regular.ttf",
            "url": "https://github.com/google/fonts/raw/main/ofl/montserrat/Montserrat-Regular.ttf"
        },
        "SourceSansPro": {
            "name": "SourceSansPro-Regular.ttf",
            "url": "https://github.com/google/fonts/raw/main/ofl/sourcesanspro/SourceSansPro-Regular.ttf"
        }
    }

    # Download fonts if not already present
    for font_key, font_details in font_info.items():
        font_path = fonts_dir / font_details["name"]
        if not font_path.exists():
            st.info(f"Downloading {font_key} font...")
            download_file(font_details["url"], font_path)
        else:
            st.success(f"Font {font_key} already exists.")

    # Return a mapping of font names to paths
    font_paths = {font_key: str(fonts_dir / details["name"]) for font_key, details in font_info.items()}
    return font_paths

def setup_stickers():
    """
    Ensures that all required stickers are downloaded and available locally.
    Downloads a set of free stickers from a reliable source.
    """
    stickers_dir = Path("stickers")
    stickers_dir.mkdir(exist_ok=True)

    # Define stickers to download
    sticker_info = {
        "Star": {
            "name": "Star.png",
            "url": "https://github.com/huggingface/datasets/raw/main/datasets/emoji/images/Star.png"
        },
        "Heart": {
            "name": "Heart.png",
            "url": "https://github.com/huggingface/datasets/raw/main/datasets/emoji/images/Heart.png"
        },
        "Smile": {
            "name": "Smile.png",
            "url": "https://github.com/huggingface/datasets/raw/main/datasets/emoji/images/Smile.png"
        }
    }

    # Download stickers if not already present
    for sticker_key, sticker_details in sticker_info.items():
        sticker_path = stickers_dir / sticker_details["name"]
        if not sticker_path.exists():
            st.info(f"Downloading {sticker_key} sticker...")
            download_file(sticker_details["url"], sticker_path)
        else:
            st.success(f"Sticker {sticker_key} already exists.")

# ------------------------------
# Initialize Application
# ------------------------------

# Set Streamlit page configuration
st.set_page_config(page_title="🎨 Pro Image Editor", layout="wide")

# Display a loading spinner while setting up fonts and stickers
with st.spinner('Setting up fonts and stickers...'):
    font_paths = setup_fonts()
    setup_stickers()

# Initialize session state for layers and history
if 'layers' not in st.session_state:
    st.session_state['layers'] = []
if 'history' not in st.session_state:
    st.session_state['history'] = []
if 'redo_stack' not in st.session_state:
    st.session_state['redo_stack'] = []

# Initialize session state for text addition
if 'text_input' not in st.session_state:
    st.session_state['text_input'] = ""
if 'click_position' not in st.session_state:
    st.session_state['click_position'] = None

# ------------------------------
# Layer Management Functions
# ------------------------------

def add_layer(layer):
    st.session_state['layers'].append(layer)
    st.session_state['history'].append(('add', layer))
    st.session_state['redo_stack'] = []

def remove_layer(index):
    layer = st.session_state['layers'].pop(index)
    st.session_state['history'].append(('remove', layer, index))
    st.session_state['redo_stack'] = []

def reorder_layers(from_idx, to_idx):
    layer = st.session_state['layers'].pop(from_idx)
    st.session_state['layers'].insert(to_idx, layer)
    st.session_state['history'].append(('reorder', from_idx, to_idx))
    st.session_state['redo_stack'] = []

def toggle_visibility(index):
    st.session_state['layers'][index]['visible'] = not st.session_state['layers'][index].get('visible', True)
    st.session_state['history'].append(('toggle', index))
    st.session_state['redo_stack'] = []

def undo():
    if not st.session_state['history']:
        st.warning("No actions to undo.")
        return
    action = st.session_state['history'].pop()
    st.session_state['redo_stack'].append(action)
    action_type = action[0]
    if action_type == 'add':
        st.session_state['layers'].pop()
    elif action_type == 'remove':
        _, layer, index = action
        st.session_state['layers'].insert(index, layer)
    elif action_type == 'reorder':
        _, from_idx, to_idx = action
        layer = st.session_state['layers'].pop(to_idx)
        st.session_state['layers'].insert(from_idx, layer)
    elif action_type == 'toggle':
        _, index = action
        st.session_state['layers'][index]['visible'] = not st.session_state['layers'][index].get('visible', True)

def redo():
    if not st.session_state['redo_stack']:
        st.warning("No actions to redo.")
        return
    action = st.session_state['redo_stack'].pop()
    st.session_state['history'].append(action)
    action_type = action[0]
    if action_type == 'add':
        _, layer = action
        st.session_state['layers'].append(layer)
    elif action_type == 'remove':
        _, _, index = action
        st.session_state['layers'].pop(index)
    elif action_type == 'reorder':
        _, from_idx, to_idx = action
        layer = st.session_state['layers'].pop(from_idx)
        st.session_state['layers'].insert(to_idx, layer)
    elif action_type == 'toggle':
        _, index = action
        st.session_state['layers'][index]['visible'] = not st.session_state['layers'][index].get('visible', True)

# ------------------------------
# Image Processing Functions
# ------------------------------

def apply_filters(image, filters):
    if filters.get('grayscale'):
        image = ImageOps.grayscale(image).convert("RGBA")
    if filters.get('sepia'):
        sepia_image = ImageOps.colorize(ImageOps.grayscale(image), '#704214', '#C0A080')
        image = sepia_image.convert("RGBA")
    if filters.get('blur'):
        image = image.filter(ImageFilter.GaussianBlur(radius=filters.get('blur_radius', 2)))
    if filters.get('brightness') != 1.0:
        enhancer = ImageEnhance.Brightness(image)
        image = enhancer.enhance(filters['brightness'])
    if filters.get('contrast') != 1.0:
        enhancer = ImageEnhance.Contrast(image)
        image = enhancer.enhance(filters['contrast'])
    return image

def auto_enhance(image):
    enhancer = ImageEnhance.Color(image)
    enhanced_image = enhancer.enhance(1.5)  # Example enhancement factor
    enhancer = ImageEnhance.Brightness(enhanced_image)
    enhanced_image = enhancer.enhance(1.2)
    enhancer = ImageEnhance.Contrast(enhanced_image)
    enhanced_image = enhancer.enhance(1.3)
    return enhanced_image

def remove_background(image):
    img_bytes = io.BytesIO()
    image.save(img_bytes, format='PNG')
    img_bytes = img_bytes.getvalue()
    output = remove(img_bytes)
    new_image = Image.open(io.BytesIO(output)).convert("RGBA")
    return new_image

def blend_images(base, overlay, mode):
    if mode == "Multiply":
        return ImageChops.multiply(base, overlay)
    elif mode == "Screen":
        return ImageChops.screen(base, overlay)
    elif mode == "Overlay":
        return ImageChops.overlay(base, overlay)
    else:
        return Image.alpha_composite(base, overlay)

# ------------------------------
# Sidebar Components
# ------------------------------

# 1. Upload Base Image
st.sidebar.header("1. Upload Base Image")
uploaded_file = st.sidebar.file_uploader("Choose an image file", type=["jpg", "jpeg", "png"])

# 2. Add a Layer
st.sidebar.header("2. Add a Layer")
layer_type = st.sidebar.selectbox("Select Layer Type", ["Text", "Image", "Shape", "Sticker"])

if layer_type == "Text":
    text_content = st.sidebar.text_input("Enter Text", "Your text here")
    font_size = st.sidebar.slider("Font Size", 10, 200, 40)
    font_color = st.sidebar.color_picker("Font Color", "#000000")
    font_choice = st.sidebar.selectbox(
        "Select Font",
        list(font_paths.keys())
    )
    selected_font_path = font_paths.get(font_choice, list(font_paths.values())[0])
    x_position = st.sidebar.slider("X Position", 0, 2000, 50)
    y_position = st.sidebar.slider("Y Position", 0, 2000, 50)
    opacity = st.sidebar.slider("Opacity", 0, 100, 100)
    alignment = st.sidebar.selectbox("Text Alignment", ["left", "center", "right"])
    blending_mode = st.sidebar.selectbox("Blending Mode", ["Normal", "Multiply", "Screen", "Overlay"])

    if st.sidebar.button("Add Text Layer"):
        layer = {
            'type': 'Text',
            'content': text_content,
            'x': x_position,
            'y': y_position,
            'font_size': font_size,
            'font_color': font_color,
            'font_path': selected_font_path,
            'opacity': opacity,
            'visible': True,
            'alignment': alignment,
            'blending_mode': blending_mode
        }
        add_layer(layer)
        st.sidebar.success("Text layer added!")

elif layer_type == "Image":
    overlay_file = st.sidebar.file_uploader("Upload an image for the layer", type=["jpg", "jpeg", "png"], key="overlay")
    if overlay_file:
        overlay_image = Image.open(overlay_file).convert("RGBA")
        img_width, img_height = overlay_image.size
        max_width = st.sidebar.slider("Width", 10, 2000, img_width)
        max_height = st.sidebar.slider("Height", 10, 2000, img_height)
        overlay_image = overlay_image.resize((max_width, max_height))
        x_position = st.sidebar.slider("X Position", 0, 2000, 50)
        y_position = st.sidebar.slider("Y Position", 0, 2000, 50)
        opacity = st.sidebar.slider("Opacity", 0, 100, 100)
        rotation = st.sidebar.slider("Rotation (degrees)", 0, 360, 0)
        blending_mode = st.sidebar.selectbox("Blending Mode", ["Normal", "Multiply", "Screen", "Overlay"])

        if st.sidebar.button("Add Image Layer"):
            overlay_image = overlay_image.rotate(rotation, expand=True)
            layer = {
                'type': 'Image',
                'content': overlay_image,
                'x': x_position,
                'y': y_position,
                'opacity': opacity,
                'visible': True,
                'blending_mode': blending_mode
            }
            add_layer(layer)
            st.sidebar.success("Image layer added!")

elif layer_type == "Shape":
    shape_type = st.sidebar.selectbox("Shape Type", ["Rectangle", "Circle", "Line"])
    shape_color = st.sidebar.color_picker("Shape Color", "#FF0000")
    x1 = st.sidebar.slider("Start X", 0, 2000, 100)
    y1 = st.sidebar.slider("Start Y", 0, 2000, 100)
    x2 = st.sidebar.slider("End X", 0, 2000, 300)
    y2 = st.sidebar.slider("End Y", 0, 2000, 300)
    thickness = st.sidebar.slider("Thickness", 1, 50, 5)
    opacity = st.sidebar.slider("Opacity", 0, 100, 100)
    blending_mode = st.sidebar.selectbox("Blending Mode", ["Normal", "Multiply", "Screen", "Overlay"])

    if st.sidebar.button("Add Shape Layer"):
        layer = {
            'type': 'Shape',
            'shape': shape_type,
            'color': shape_color,
            'coordinates': (x1, y1, x2, y2),
            'thickness': thickness,
            'opacity': opacity,
            'visible': True,
            'blending_mode': blending_mode
        }
        add_layer(layer)
        st.sidebar.success("Shape layer added!")

elif layer_type == "Sticker":
    sticker_options = ["Star", "Heart", "Smile"]
    sticker = st.sidebar.selectbox("Select Sticker", sticker_options)
    sticker_path = Path(f"stickers/{sticker}.png")
    if sticker_path.exists():
        sticker_image = Image.open(sticker_path).convert("RGBA")
        img_width, img_height = sticker_image.size
        max_width = st.sidebar.slider("Width", 10, 500, img_width)
        max_height = st.sidebar.slider("Height", 10, 500, img_height)
        sticker_image = sticker_image.resize((max_width, max_height))
        x_position = st.sidebar.slider("X Position", 0, 2000, 50)
        y_position = st.sidebar.slider("Y Position", 0, 2000, 50)
        opacity = st.sidebar.slider("Opacity", 0, 100, 100)
        rotation = st.sidebar.slider("Rotation (degrees)", 0, 360, 0)
        blending_mode = st.sidebar.selectbox("Blending Mode", ["Normal", "Multiply", "Screen", "Overlay"])

        if st.sidebar.button("Add Sticker Layer"):
            sticker_image = sticker_image.rotate(rotation, expand=True)
            layer = {
                'type': 'Sticker',
                'content': sticker_image,
                'x': x_position,
                'y': y_position,
                'opacity': opacity,
                'visible': True,
                'blending_mode': blending_mode
            }
            add_layer(layer)
            st.sidebar.success("Sticker layer added!")
    else:
        st.sidebar.error("Sticker image not found.")

# 3. Manage Layers
st.sidebar.header("3. Manage Layers")
if st.session_state['layers']:
    st.sidebar.subheader("Layers")
    for idx, layer in enumerate(reversed(st.session_state['layers'])):
        layer_num = len(st.session_state['layers']) - idx
        with st.sidebar.expander(f"Layer {layer_num} ({layer['type']})"):
            visibility = st.checkbox("Visible", value=layer.get('visible', True), key=f"visibility_{idx}")
            if visibility != layer.get('visible', True):
                toggle_visibility(len(st.session_state['layers']) - idx - 1)
            if st.button("Remove", key=f"remove_{idx}"):
                remove_layer(len(st.session_state['layers']) - idx - 1)
            # Reordering
            new_order = st.selectbox(
                "Move to position",
                options=list(range(1, len(st.session_state['layers'])+1)), 
                index=len(st.session_state['layers'])-idx-1, 
                key=f"move_{idx}"
            )
            if new_order != (len(st.session_state['layers']) - idx):
                reorder_layers(len(st.session_state['layers']) - idx - 1, new_order - 1)

# 4. History (Undo/Redo)
st.sidebar.header("4. History")
col1, col2 = st.sidebar.columns(2)
with col1:
    if st.sidebar.button("Undo"):
        undo()
with col2:
    if st.sidebar.button("Redo"):
        redo()

# 5. Filters & Adjustments
st.sidebar.header("5. Filters & Adjustments")
apply_grayscale = st.sidebar.checkbox("Grayscale")
apply_sepia = st.sidebar.checkbox("Sepia")
apply_blur = st.sidebar.checkbox("Blur")
blur_radius = st.sidebar.slider("Blur Radius", 0.0, 10.0, 2.0)
brightness = st.sidebar.slider("Brightness", 0.5, 2.0, 1.0)
contrast = st.sidebar.slider("Contrast", 0.5, 2.0, 1.0)
apply_auto_enhance = st.sidebar.checkbox("Auto Enhance")
remove_bg = st.sidebar.checkbox("Remove Background")

filters = {
    'grayscale': apply_grayscale,
    'sepia': apply_sepia,
    'blur': apply_blur,
    'blur_radius': blur_radius,
    'brightness': brightness,
    'contrast': contrast
}

# 6. Advanced Tools
st.sidebar.header("6. Advanced Tools")
if st.sidebar.button("Remove Background"):
    if uploaded_file:
        try:
            base_image = Image.open(uploaded_file).convert("RGBA")
            canvas = remove_background(base_image)
            st.session_state['layers'] = []  # Clear existing layers
            st.success("Background removed successfully!")
        except Exception as e:
            st.error(f"Background removal failed: {e}")
    else:
        st.warning("Please upload an image first.")

# 7. Save/Load Project
st.sidebar.header("7. Save/Load Project")
project_name = st.sidebar.text_input("Project Name")
if st.sidebar.button("Save Project") and project_name:
    projects_dir = Path("projects")
    projects_dir.mkdir(exist_ok=True)
    project_path = projects_dir / f"{project_name}.json"
    try:
        with open(project_path, "w") as f:
            json.dump(st.session_state['layers'], f)
        st.sidebar.success(f"Project '{project_name}' saved!")
    except Exception as e:
        st.sidebar.error(f"Failed to save project: {e}")

load_project = st.sidebar.file_uploader("Load Project", type=["json"], key="load_project")
if load_project:
    try:
        loaded_layers = json.load(load_project)
        st.session_state['layers'] = loaded_layers
        st.sidebar.success("Project loaded successfully!")
    except Exception as e:
        st.sidebar.error(f"Failed to load project: {e}")

# ------------------------------
# Main Area: Display and Edit Image
# ------------------------------

if uploaded_file:
    base_image = Image.open(uploaded_file).convert("RGBA")
    canvas_image = base_image.copy()

    # Apply auto enhance if selected
    if apply_auto_enhance:
        canvas_image = auto_enhance(canvas_image)

    # Initialize drawing context
    draw = ImageDraw.Draw(canvas_image, 'RGBA')

    # ------------------------------
    # Interactive Text Addition
    # ------------------------------

    st.header("🖌️ Interactive Text Addition")

    # Display the canvas with drawable features
    canvas_result = st_canvas(
        fill_color="rgba(255, 255, 255, 0)",  # Transparent fill
        stroke_width=0,
        stroke_color="#000000",
        background_image=canvas_image,
        update_streamlit=True,
        height=canvas_image.height,
        width=canvas_image.width,
        drawing_mode="point",
        key="canvas",
        # Hide the drawing toolbar
        display_toolbar=False,
    )

    if canvas_result.json_data is not None:
        objects = canvas_result.json_data["objects"]
        if objects:
            last_object = objects[-1]
            if last_object["type"] == "point":
                x, y = last_object["left"], last_object["top"]
                st.session_state['click_position'] = (int(x), int(y))
                st.session_state['text_input'] = st.text_input("Enter Text for the Clicked Position:", "")
                if st.button("Add Text Here"):
                    text_content = st.session_state['text_input']
                    if text_content:
                        # Define default properties for interactive text
                        layer = {
                            'type': 'Text',
                            'content': text_content,
                            'x': st.session_state['click_position'][0],
                            'y': st.session_state['click_position'][1],
                            'font_size': 40,
                            'font_color': "#000000",
                            'font_path': list(font_paths.values())[0],  # Default to the first font
                            'opacity': 100,
                            'visible': True,
                            'alignment': "left",
                            'blending_mode': "Normal"
                        }
                        add_layer(layer)
                        st.sidebar.success("Interactive text added!")
                        # Reset click position and text input
                        st.session_state['click_position'] = None
                        st.session_state['text_input'] = ""
                    else:
                        st.warning("Please enter some text.")

    # ------------------------------
    # Process Layers and Render Image
    # ------------------------------

    for layer in st.session_state['layers']:
        if not layer.get('visible', True):
            continue
        if layer['type'] == "Text":
            try:
                font = ImageFont.truetype(layer['font_path'], layer['font_size'])
            except IOError:
                font = ImageFont.load_default()
                st.warning(f"Font not found: {layer['font_path']}. Using default font.")
            text_color = layer['font_color']
            text_opacity = int(255 * (layer.get('opacity', 100) / 100))
            try:
                text_color_rgba = ImageColor.getcolor(text_color, "RGBA")
            except ValueError:
                text_color_rgba = (0, 0, 0, text_opacity)
            text_color_rgba = text_color_rgba[:3] + (text_opacity,)
            text_layer = Image.new('RGBA', canvas_image.size, (255,255,255,0))
            text_draw = ImageDraw.Draw(text_layer)
            alignment = layer.get('alignment', 'left')
            text_position = (layer['x'], layer['y'])
            text_draw.text(text_position, layer['content'], fill=text_color_rgba, font=font, align=alignment)
            if layer.get('blending_mode') and layer['blending_mode'] != "Normal":
                canvas_image = blend_images(canvas_image, text_layer, layer['blending_mode'])
            else:
                canvas_image = Image.alpha_composite(canvas_image, text_layer)
        elif layer['type'] in ["Image", "Sticker"]:
            overlay = layer['content'].copy()
            opacity = layer.get('opacity', 100)
            if opacity < 100:
                alpha = overlay.split()[3]
                alpha = alpha.point(lambda p: p * opacity / 100)
                overlay.putalpha(alpha)
            if layer.get('blending_mode') and layer['blending_mode'] != "Normal":
                canvas_image = blend_images(canvas_image, overlay, layer['blending_mode'])
            else:
                canvas_image.paste(overlay, (layer['x'], layer['y']), overlay)
        elif layer['type'] == "Shape":
            shape_layer = Image.new('RGBA', canvas_image.size, (255,255,255,0))
            shape_draw = ImageDraw.Draw(shape_layer)
            color = layer['color']
            opacity = layer.get('opacity', 100)
            try:
                shape_color_rgba = ImageColor.getcolor(color, "RGBA")
            except ValueError:
                shape_color_rgba = (255, 0, 0, int(255 * (opacity / 100)))
            shape_color_rgba = shape_color_rgba[:3] + (int(255 * (opacity / 100)),)
            x1, y1, x2, y2 = layer['coordinates']
            thickness = layer['thickness']
            if layer['shape'] == "Rectangle":
                shape_draw.rectangle([x1, y1, x2, y2], outline=shape_color_rgba, width=thickness)
            elif layer['shape'] == "Circle":
                shape_draw.ellipse([x1, y1, x2, y2], outline=shape_color_rgba, width=thickness)
            elif layer['shape'] == "Line":
                shape_draw.line([x1, y1, x2, y2], fill=shape_color_rgba, width=thickness)
            if layer.get('blending_mode') and layer['blending_mode'] != "Normal":
                canvas_image = blend_images(canvas_image, shape_layer, layer['blending_mode'])
            else:
                canvas_image = Image.alpha_composite(canvas_image, shape_layer)

    # Apply filters
    canvas_image = apply_filters(canvas_image, filters)

    # Remove background if selected
    if remove_bg:
        try:
            canvas_image = remove_background(canvas_image)
            st.success("Background removed successfully!")
        except Exception as e:
            st.error(f"Background removal failed: {e}")

    # Display the final image
    st.image(canvas_image, caption="Edited Image", use_column_width=True)

    # Download options
    buffer = io.BytesIO()
    canvas_image.save(buffer, format="PNG")
    st.download_button(
        label="📥 Download Edited Image",
        data=buffer.getvalue(),
        file_name="edited_image.png",
        mime="image/png"
    )

else:
    st.warning("Please upload an image to start editing.")