import gradio as gr
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles
from fastapi.responses import RedirectResponse
import shutil
import os
import uuid
from datetime import datetime
from PIL import Image
from collections import defaultdict
import base64
from io import BytesIO
import requests
import json
from dotenv import load_dotenv
import os


# Directories
TEMP_DIR = "cloud_storage/temp"
PERM_DIR = "cloud_storage/permanent"
THUMB_DIR = "cloud_storage/thumbs"
os.makedirs(TEMP_DIR, exist_ok=True)
os.makedirs(PERM_DIR, exist_ok=True)
os.makedirs(THUMB_DIR, exist_ok=True)


# Load variables from .env file
load_dotenv()

# Sightengine API credentials (set as environment variables)
API_USER = os.getenv('SIGHTENGINE_API_USER')
API_SECRET = os.getenv('SIGHTENGINE_API_SECRET')

# FastAPI app
app = FastAPI()

# Redirect to '/gradio' endpoint of FastAPI.
@app.get("/")
async def root():
    return RedirectResponse(url="/gradio/?__theme=light", status_code=307)
    
# Serve images and thumbnails as static files
app.mount("/images/temp", StaticFiles(directory=TEMP_DIR), name="temp_images")
app.mount("/images/permanent", StaticFiles(directory=PERM_DIR), name="perm_images")
app.mount("/images/thumbs", StaticFiles(directory=THUMB_DIR), name="thumbnails")

# In-memory analytics store
analytics = defaultdict(lambda: {"views": 0, "likes": 0, "upload_time": None, "size": 0, "thumb": None, "tags": []})

# Generate thumbnail for an image
def generate_thumbnail(src_path, filename):
    thumb_path = os.path.join(THUMB_DIR, filename)
    try:
        with Image.open(src_path) as img:
            img.thumbnail((300, 300))
            img.save(thumb_path, quality=85)
        return f"/images/thumbs/{filename}"
    except Exception as e:
        print(f"Error generating thumbnail: {e}")
        return None

# Generate data URI for preview
def generate_preview(image_path):
    try:
        with Image.open(image_path) as img:
            img.thumbnail((100, 100))
            buffered = BytesIO()
            img.save(buffered, format="PNG")
            return f"data:image/png;base64,{base64.b64encode(buffered.getvalue()).decode()}"
    except Exception as e:
        print(f"Error generating preview: {e}")
        return None

# Generate HTML for previews
def generate_preview_html(image_paths):
    if not image_paths:
        return "<p style='color:#666;font-family:Arial,sans-serif;'>No images selected.</p>"
    html = "<div style='display: flex; flex-wrap: wrap; gap: 10px; padding: 10px;'>"
    for path in image_paths:
        data_uri = generate_preview(path)
        filename = os.path.basename(path)
        short_name = filename[:8] + "..." if len(filename) > 10 else filename
        if data_uri:
            html += f"""
            <div style='text-align: center;'>
                <img src='{data_uri}' style='width: 100px; height: 100px; object-fit: cover; border-radius: 5px;'>
                <p style='font-size: 12px; color: #333; margin: 5px 0;'>{short_name}</p>
            </div>
            """
    html += "</div>"
    return html

# Save images to chosen directory with NSFW moderation
def save_image(image_paths, store_type, tags):
    if not image_paths:
        return "No images uploaded.", "", get_gallery_html(), get_analytics_html(), gr.update(value=None), gr.update(value=None)
    
    save_dir = TEMP_DIR if store_type == "Temporary" else PERM_DIR
    tag_list = [tag.strip() for tag in tags.split(",")] if tags else []
    saved_count = 0
    rejected = []
    
    for image_path in image_paths:
        # Moderate the image using Sightengine API
        params = {
            'models': 'nudity-2.1',
            'api_user': API_USER,
            'api_secret': API_SECRET
        }
        files = {'media': open(image_path, 'rb')}
        try:
            r = requests.post('https://api.sightengine.com/1.0/check.json', files=files, data=params)
            output = json.loads(r.text)
            
            if output['status'] == 'success':
                nudity = output['nudity']
                # Check if the image is NSFW
                if nudity['none'] < 0.9 or any(score > 0.1 for score in [
                    nudity.get('sexual_activity', 0),
                    nudity.get('sexual_display', 0),
                    nudity.get('erotica', 0),
                    nudity.get('very_suggestive', 0),
                    nudity.get('suggestive', 0),
                    nudity.get('mildly_suggestive', 0)
                ]):
                    rejected.append(os.path.basename(image_path))
                    continue
                else:
                    # Save the image if safe
                    ext = os.path.splitext(image_path)[1]
                    new_filename = f"{uuid.uuid4().hex}{ext}"
                    dest_path = os.path.join(save_dir, new_filename)
                    shutil.copy(image_path, dest_path)
                    thumb_src = generate_thumbnail(dest_path, new_filename)
                    analytics[(new_filename, store_type)] = {
                        "views": 0,
                        "likes": 0,
                        "upload_time": datetime.utcnow().isoformat(),
                        "size": os.path.getsize(dest_path),
                        "thumb": thumb_src,
                        "tags": tag_list
                    }
                    saved_count += 1
            else:
                print(f"Error moderating image {image_path}: {output.get('error', 'Unknown error')}")
                rejected.append(os.path.basename(image_path))
        except Exception as e:
            print(f"Error processing image {image_path}: {e}")
            rejected.append(os.path.basename(image_path))
    
    status_message = f"Saved {saved_count} out of {len(image_paths)} images to {store_type} storage."
    if rejected:
        status_message += f" Rejected: {', '.join(rejected)}"
    
    return (status_message, 
            "", 
            get_gallery_html(), 
            get_analytics_html(), 
            gr.update(value=None), 
            gr.update(value=None))

# Delete selected image
def delete_image(filename, storage):
    dir_path = TEMP_DIR if storage == "Temporary" else PERM_DIR
    file_path = os.path.join(dir_path, filename)
    thumb_path = os.path.join(THUMB_DIR, filename)
    if os.path.exists(file_path):
        os.remove(file_path)
        if os.path.exists(thumb_path):
            os.remove(thumb_path)
        del analytics[(filename, storage)]
        return f"Deleted {filename} from {storage} storage.", get_gallery_html(), get_analytics_html()
    return "Image not found.", get_gallery_html(), get_analytics_html()

# Like an image
def like_image(filename, storage):
    analytics[(filename, storage)]["likes"] += 1
    return f"Liked {filename}.", get_gallery_html(), get_analytics_html()

# Build gallery HTML with masonry-like layout and updated hover behavior
def get_gallery_html():
    def make_img_tag(src, filename, storage):
        meta = analytics[(filename, storage)]
        thumb_src = meta.get("thumb", src)
        views = meta["views"]
        likes = meta["likes"]
        upload_time = meta["upload_time"]
        size_kb = meta["size"] // 1024
        tags = ", ".join(meta["tags"]) if meta["tags"] else "No tags"
        delete_event = f"document.querySelector('#delete_{filename}_{storage}').click()"
        like_event = f"document.querySelector('#like_{filename}_{storage}').click()"
        short_name = filename[:8] + "..." if len(filename) > 10 else filename
        share_link = f"{src}"
        
        modal_data = f"{short_name}|{storage}|{views}|{size_kb}|{upload_time}|{tags}"
        
        return f'''
        <div class="gallery-item">
            <div class="image-container">
                <img src="{thumb_src}" alt="{short_name}" loading="lazy" onclick="openModal('{src}', '{modal_data}')">
                <div class="overlay">
                    <a href="{src}" download class="download-btn">⬇</a>
                    <div class="action-buttons">
                        <button onclick="{delete_event}; event.stopPropagation();" class="delete-btn">🗑️</button>
                        <button onclick="navigator.clipboard.writeText('{share_link}'); event.stopPropagation();" class="share-btn">🔗</button>
                        <button onclick="{like_event}; event.stopPropagation();" class="like-btn">❤️ {likes}</button>
                    </div>
                </div>
            </div>
        </div>
        '''

    temp_html = ""
    perm_html = ""
    for (filename, storage), meta in analytics.items():
        src = f"/images/{'temp' if storage == 'Temporary' else 'permanent'}/{filename}"
        img_tag = make_img_tag(src, filename, storage)
        if storage == "Temporary":
            temp_html += img_tag
        else:
            perm_html += img_tag

    css = '''
    <style>
        body {
            background: linear-gradient(135deg, #ece9e6, #ffffff);
            margin: 0;
            font-family: 'Arial', sans-serif;
        }
        .gallery-section {
            margin: 40px 20px;
            padding: 20px;
            background: #ffffff;
            border-radius: 15px;
            box-shadow: 0 8px 16px rgba(0,0,0,0.1);
        }
        .gallery-section h2 {
            color: #1a1a1a;
            text-align: center;
            margin-bottom: 30px;
            font-size: 1.8em;
            font-weight: 600;
            letter-spacing: 0.5px;
        }
        .gallery-container {
            column-count: 4;
            column-gap: 15px;
            padding: 10px;
        }
        .gallery-item {
            position: relative;
            overflow: hidden;
            border-radius: 8px;
            cursor: pointer;
            margin-bottom: 15px;
            break-inside: avoid;
            opacity: 0;
            animation: fadeIn 0.5s ease-out forwards;
        }
        @keyframes fadeIn {
            from { opacity: 0; transform: translateY(20px); }
            to { opacity: 1; transform: translateY(0); }
        }
        .gallery-item:nth-child(odd) {
            animation-delay: 0.1s;
        }
        .gallery-item:nth-child(even) {
            animation-delay: 0.2s;
        }
        .image-container {
            position: relative;
            overflow: hidden;
            border-radius: 8px;
        }
        .image-container img {
            width: 100%;
            height: auto;
            display: block;
            border-radius: 8px;
            transition: transform 0.3s ease;
        }
        .gallery-item:hover .image-container img {
            transform: scale(1.1);
        }
        .overlay {
            position: absolute;
            top: 0;
            left: 0;
            right: 0;
            bottom: 0;
            background: transparent;
            opacity: 0;
            transition: opacity 0.3s ease;
            border-radius: 8px;
        }
        .gallery-item:hover .overlay {
            opacity: 1;
        }
        .download-btn {
            position: absolute;
            top: 10px;
            right: 10px;
            color: white;
            font-size: 1.2em;
            text-decoration: none;
            background: #007bff;
            padding: 5px 10px;
            border-radius: 5px;
            z-index: 10;
        }
        .download-btn:hover {
            background: #0056b3;
        }
        .action-buttons {
            position: absolute;
            bottom: 10px;
            left: 0;
            right: 0;
            display: flex;
            justify-content: center;
            align-items: center;
            gap: 15px;
        }
        .action-buttons button {
            background: rgba(0,0,0,0.7);
            border: none;
            color: white;
            font-size: 1.1em;
            cursor: pointer;
            padding: 8px 14px;
            border-radius: 5px;
            transition: background 0.2s ease;
        }
        .delete-btn { background: #dc3545; }
        .share-btn { background: #28a745; }
        .like-btn { background: #ff2e63; }
        .action-buttons button:hover {
            filter: brightness(1.15);
        }
        #imageModal {
            display: none;
            position: fixed;
            top: 0;
            left: 0;
            width: 100%;
            height: 100%;
            background: rgba(0,0,0,0.8);
            backdrop-filter: blur(10px);
            justify-content: center;
            align-items: center;
            z-index: 1000;
            animation: modalFadeIn 0.3s ease-out;
        }
        @keyframes modalFadeIn {
            from { opacity: 0; }
            to { opacity: 1; }
        }
        #modalContent {
            display: flex;
            flex-direction: column;
            align-items: center;
            gap: 20px;
            max-width: 90%;
            max-height: 90%;
        }
        #modalImage {
            max-width: 80%;
            max-height: 70%;
            border-radius: 10px;
            box-shadow: 0 4px 20px rgba(0,0,0,0.3);
            transform: scale(0.8);
            animation: modalImageScale 0.3s ease-out forwards;
        }
        @keyframes modalImageScale {
            to { transform: scale(1); }
        }
        #modalInfo {
            background: rgba(255,255,255,0.9);
            padding: 15px 20px;
            border-radius: 8px;
            max-width: 80%;
            text-align: center;
        }
        #modalInfo p {
            color: #333;
            font-size: 0.9em;
            margin: 5px 0;
            font-weight: 400;
        }
        #modalDownload {
            background: #007bff;
            color: white;
            font-size: 1em;
            text-decoration: none;
            padding: 8px 16px;
            border-radius: 5px;
            transition: background 0.2s ease;
        }
        #modalDownload:hover {
            background: #0056b3;
        }
        #closeModal {
            position: absolute;
            top: 20px;
            right: 20px;
            color: white;
            font-size: 2em;
            cursor: pointer;
            transition: color 0.2s ease;
        }
        #closeModal:hover {
            color: #ccc;
        }
        @media (max-width: 1024px) {
            .gallery-container {
                column-count: 3;
            }
        }
        @media (max-width: 768px) {
            .gallery-container {
                column-count: 2;
            }
            .gallery-section {
                margin: 20px 10px;
                padding: 15px;
            }
            #modalImage {
                max-width: 85%;
                max-height: 65%;
            }
            #modalInfo {
                max-width: 85%;
            }
        }
        @media (max-width: 480px) {
            .gallery-container {
                column-count: 1;
            }
            .action-buttons {
                gap: 10px;
            }
            .action-buttons button {
                padding: 6px 12px;
                font-size: 1em;
            }
            #modalImage {
                max-width: 90%;
                max-height: 60%;
            }
        }
    </style>
    '''

    javascript = '''
    <script>
        function openModal(src, data) {
            const [name, storage, views, size, uploadTime, tags] = data.split('|');
            document.getElementById('modalImage').src = src;
            document.getElementById('modalInfo').innerHTML = `
                <p>Name: ${name} (${storage})</p>
                <p>Views: ${views} | Size: ${size} KB</p>
                <p>Uploaded: ${uploadTime}</p>
                <p>Tags: ${tags}</p>
            `;
            document.getElementById('modalDownload').href = src;
            document.getElementById('imageModal').style.display = 'flex';
        }
        document.addEventListener('DOMContentLoaded', function() {
            const modal = document.getElementById('imageModal');
            const closeModal = document.getElementById('closeModal');
            if (closeModal) {
                closeModal.addEventListener('click', function() {
                    modal.style.display = 'none';
                });
            }
            if (modal) {
                modal.addEventListener('click', function(e) {
                    if (e.target.id === 'imageModal') {
                        modal.style.display = 'none';
                    }
                });
            }
        });
    </script>
    '''

    html = f"""
    {css}
    <div class="gallery-section">
        <h2>🕒 Temporary Uploads</h2>
        <div class="gallery-container">{temp_html}</div>
    </div>
    <div class="gallery-section">
        <h2>📁 Permanent Uploads</h2>
        <div class="gallery-container">{perm_html}</div>
    </div>
    <div id="imageModal">
        <span id="closeModal">×</span>
        <div id="modalContent">
            <img id="modalImage" src="" alt="Enlarged Image">
            <div id="modalInfo"></div>
            <a id="modalDownload" download>⬇ Download</a>
        </div>
    </div>
    {javascript}
    """
    return html

# Build analytics HTML for top 5 images
def get_analytics_html():
    sorted_images = sorted(analytics.items(), key=lambda x: x[1]["likes"], reverse=True)[:5]
    analytics_html = "<h3 style='color:#1a1a1a;margin-top:30px;font-family:Arial,sans-serif;'>📊 Top 5 Most Liked Images</h3><ul style='list-style:none;padding:0;'>"
    for (filename, storage), data in sorted_images:
        short_name = filename[:8] + "..." if len(filename) > 10 else filename
        analytics_html += f"<li style='font-size:14px;color:#555;margin:10px 0;font-family:Arial,sans-serif;'>{short_name} ({storage}): {data['likes']} likes</li>"
    analytics_html += "</ul>" if sorted_images else "<p style='color:#666;font-family:Arial,sans-serif;'>No analytics data yet.</p>"
    return analytics_html

# Gradio interface with enhanced features and UI
with gr.Blocks(theme=gr.themes.Soft()) as gradio_app:
    gr.Markdown("# 🌟 Image Upload Gallery")
    
    # Upload Section
    gr.Markdown("## 📤 Upload Images")
    with gr.Row():
        with gr.Column():
            image_input = gr.File(label="Upload images (multiple allowed)", type="filepath", file_count="multiple")
            store_option = gr.Radio(["Temporary", "Permanent"], value="Temporary", label="Storage Type")
            tags_input = gr.Textbox(label="Tags (comma-separated)", placeholder="e.g., nature, portrait, art")
        with gr.Column():
            preview_output = gr.HTML(label="Upload Preview")
    
    upload_btn = gr.Button("Upload", variant="primary")
    upload_status = gr.Textbox(label="Status", interactive=False, placeholder="Action status will appear here...")

    # Gallery Section
    gr.Markdown("## 🖼️ Gallery\nExplore and manage your images in a stunning layout.")
    gallery_html = gr.HTML(get_gallery_html)

    # Analytics Section
    gr.Markdown("## 📊 Analytics")
    analytics_html = gr.HTML(get_analytics_html)

    # Hidden buttons for event triggering
    hidden_components = []
    for (filename, storage), _ in list(analytics.items()):
        delete_btn = gr.Button(f"Delete {filename}", elem_id=f"delete_{filename}_{storage}", visible=False)
        like_btn = gr.Button(f"Like {filename}", elem_id=f"like_{filename}_{storage}", visible=False)
        delete_btn.click(fn=delete_image, inputs=[gr.State(filename), gr.State(storage)], outputs=[upload_status, gallery_html, analytics_html])
        like_btn.click(fn=like_image, inputs=[gr.State(filename), gr.State(storage)], outputs=[upload_status, gallery_html, analytics_html])
        hidden_components.extend([delete_btn, like_btn])

    # Event Handlers
    upload_btn.click(
        fn=save_image,
        inputs=[image_input, store_option, tags_input],
        outputs=[upload_status, preview_output, gallery_html, analytics_html, image_input, tags_input]
    )
    image_input.change(
        fn=generate_preview_html,
        inputs=image_input,
        outputs=preview_output
    )

gradio_app.pwa = True  # Enable PWA

# Mount Gradio at /gradio
app = gr.mount_gradio_app(app, gradio_app, path="/gradio")