import gradio as gr
import requests
import zipfile
import io
import os
import warnings
from pathlib import Path
from PIL import Image
import json
from typing import List, Tuple

# Suppress insecure request warnings
warnings.filterwarnings('ignore', message='Unverified HTTPS request')

# Try to import huggingface_hub and wd_tagger
try:
    from huggingface_hub import login, logout, HfApi
    HAS_HF_HUB = True
except ImportError:
    HAS_HF_HUB = False
    print("Warning: huggingface_hub not installed. HF dataset features disabled.")

try:
    from wd_tagger import (
        get_loaded_models,
        load,
        tag
    )
    HAS_WD_TAGGER = True
except ImportError:
    HAS_WD_TAGGER = False
    print("Warning: wd_tagger not installed. Tagging features disabled.")


# Global state for selected images
selected_images_state = {"images": [], "file_names": []}


def download_and_extract_images(zip_url, hf_token=None, verify_ssl=False):
    """
    Download a deflate zip file from URL and extract images.
    
    Args:
        zip_url: URL to the zip file or HuggingFace dataset
        hf_token: Huggingface authentication token
        verify_ssl: Whether to verify SSL certificates
        
    Returns:
        Tuple of (images list, file names list, status message)
    """
    try:
        # Handle HuggingFace dataset URLs
        if 'huggingface.co' in zip_url or zip_url.startswith('datasets/'):
            if not HAS_HF_HUB:
                return None, [], "❌ Error: huggingface_hub not installed"
            
            if hf_token:
                try:
                    login(token=hf_token)
                    print("✅ Logged into Huggingface")
                except Exception as e:
                    return None, [], f"❌ Error: Failed to login to HF - {str(e)}"
            
            try:
                # Parse dataset name
                if zip_url.startswith('datasets/'):
                    dataset_name = zip_url.replace('datasets/', '')
                else:
                    dataset_name = zip_url.split('/')[-1]
                
                print(f"Loading dataset: {dataset_name}")
                
                # Load dataset
                from datasets import load_dataset
                dataset = load_dataset(dataset_name)
                
                images = []
                file_names = []
                
                # Extract images from dataset
                for split_name, split in dataset.items():
                    for idx, example in enumerate(split):
                        if 'image' in example:
                            img = example['image']
                            if isinstance(img, Image.Image):
                                images.append(img)
                                file_names.append(f"{split_name}_{idx}.png")
                
                if not images:
                    return None, [], "❌ No images found in dataset"
                
                return images, file_names, f"✅ Loaded {len(images)} images from HF dataset"
                
            except Exception as e:
                return None, [], f"❌ Error loading dataset: {str(e)}"
        
        # Handle regular zip URLs
        print(f"Downloading from: {zip_url}")
        
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
        }
        
        response = requests.get(
            zip_url, 
            verify=verify_ssl, 
            allow_redirects=True, 
            timeout=30,
            headers=headers
        )
        response.raise_for_status()
        
        print(f"Response status: {response.status_code}")
        print(f"Content length: {len(response.content)} bytes")
        print(f"Content type: {response.headers.get('content-type', 'unknown')}")
        
        # Check if content is actually a zip file
        if not response.content.startswith(b'PK'):
            try:
                import gzip
                decompressed = gzip.decompress(response.content)
                if decompressed.startswith(b'PK'):
                    print("Content was gzip-compressed, decompressed successfully")
                    response.content = decompressed
                else:
                    return None, [], f"❌ Error: Downloaded file is not a zip file.\n\nContent-Type: {response.headers.get('content-type')}\n\nTip: Use a direct .zip link"
            except Exception as e:
                return None, [], f"❌ Error: Downloaded file is not a zip file.\n\nContent-Type: {response.headers.get('content-type')}\n\nTip: Use a direct .zip link"
        
        zip_buffer = io.BytesIO(response.content)
        
        images = []
        file_names = []
        supported_formats = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.webp'}
        
        try:
            with zipfile.ZipFile(zip_buffer, 'r') as zip_file:
                file_list = zip_file.namelist()
                print(f"Found {len(file_list)} files in zip")
                
                image_files = [f for f in file_list if Path(f).suffix.lower() in supported_formats]
                print(f"Found {len(image_files)} image files")
                
                for file_name in image_files:
                    try:
                        image_data = zip_file.read(file_name)
                        image = Image.open(io.BytesIO(image_data))
                        images.append(image)
                        file_names.append(file_name)
                        print(f"Loaded: {file_name}")
                    except Exception as e:
                        print(f"Error loading {file_name}: {e}")
                        continue
        
        except zipfile.BadZipFile as e:
            print(f"BadZipFile error: {e}")
            return None, [], f"❌ Error: Invalid or corrupted zip file"
        except Exception as e:
            print(f"Zip extraction error: {e}")
            return None, [], f"❌ Error: Failed to extract zip - {str(e)}"
        
        if not images:
            return None, [], "❌ No images found in the zip file"
        
        print(f"Successfully extracted {len(images)} images")
        return images, file_names, f"✅ Loaded {len(images)} images"
        
    except requests.exceptions.RequestException as e:
        return None, [], f"❌ Download error: {str(e)}"
    except Exception as e:
        return None, [], f"❌ Error: {str(e)}"


def resize_image(image: Image.Image, target_height: int) -> Image.Image:
    """Resize image maintaining aspect ratio based on target height"""
    if target_height <= 0:
        return image
    
    ratio = target_height / image.height
    new_width = int(image.width * ratio)
    return image.resize((new_width, target_height), Image.Resampling.LANCZOS)


def tag_images_wd(selected_indices: List[int], images_list: List[Image.Image], file_names_list: List[str]) -> Tuple[str, dict]:
    """Tag selected images using WD-Tagger"""
    if not HAS_WD_TAGGER:
        return "❌ Error: wd_tagger not installed", {}
    
    if not selected_indices:
        return "❌ Error: No images selected", {}
    
    try:
        print("Loading WD-Tagger model...")
        model_name = list(get_loaded_models())[0] if get_loaded_models() else None
        
        if not model_name:
            load("wd14-vit")
            model_name = "wd14-vit"
        
        captions = {}
        
        for idx in selected_indices:
            if idx < len(images_list):
                image = images_list[idx]
                file_name = file_names_list[idx]
                
                print(f"Tagging: {file_name}")
                results, ratings = tag(image, model_name)
                
                # Combine tags with confidence scores
                tags_str = ", ".join([f"{tag} ({conf:.2f})" for tag, conf in results.items()])
                captions[file_name] = tags_str
        
        return f"✅ Tagged {len(captions)} images", captions
        
    except Exception as e:
        print(f"Error tagging: {e}")
        return f"❌ Error tagging images: {str(e)}", {}


def export_selected(selected_indices: List[int], images_list: List[Image.Image], file_names_list: List[str], captions_dict: dict, target_height: int):
    """Export selected images and captions as zip"""
    if not selected_indices:
        return None, "❌ Error: No images selected"
    
    try:
        if target_height <= 0:
            return None, "❌ Error: Invalid height value"
        
        zip_buffer = io.BytesIO()
        
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
            for idx in selected_indices:
                if idx < len(images_list):
                    image = images_list[idx]
                    file_name = file_names_list[idx]
                    
                    # Resize image
                    resized = resize_image(image, target_height)
                    
                    # Save image to zip
                    img_buffer = io.BytesIO()
                    resized.save(img_buffer, format='PNG')
                    zip_file.writestr(f"images/{Path(file_name).stem}.png", img_buffer.getvalue())
                    
                    # Save caption if available
                    if file_name in captions_dict:
                        caption_name = f"captions/{Path(file_name).stem}.txt"
                        zip_file.writestr(caption_name, captions_dict[file_name])
        
        zip_buffer.seek(0)
        return zip_buffer, f"✅ Exported {len(selected_indices)} images with captions"
        
    except Exception as e:
        print(f"Export error: {e}")
        return None, f"❌ Error exporting: {str(e)}"


def create_gallery_with_checkboxes(images_list: List[Image.Image], file_names_list: List[str], target_height: int):
    """Create gallery display with resized images"""
    if not images_list:
        return [], "No images loaded"
    
    resized_images = []
    for img in images_list:
        resized = resize_image(img, target_height)
        resized_images.append(resized)
    
    return resized_images, f"Displaying {len(resized_images)} images (resized to height: {target_height}px)"


# Create Gradio interface
with gr.Blocks(title="Zip Image Gallery Pro") as demo:
    gr.Markdown("# 📸 Zip Image Gallery Pro")
    gr.Markdown("Download, tag, and export images with captions")
    
    # State to store images and file names
    images_state = gr.State([])
    file_names_state = gr.State([])
    captions_state = gr.State({})
    
    with gr.Tabs():
        # Tab 1: Download
        with gr.Tab("📥 Download"):
            gr.Markdown("### Load Images from Zip or HuggingFace Dataset")
            
            with gr.Row():
                with gr.Column(scale=3):
                    url_input = gr.Textbox(
                        label="Zip File URL or HF Dataset",
                        placeholder="Enter .zip URL or datasets/username/dataset-name",
                        lines=1
                    )
                with gr.Column(scale=1):
                    hf_token_input = gr.Textbox(
                        label="HF Token (optional)",
                        placeholder="Your HF token",
                        type="password",
                        lines=1
                    )
            
            download_btn = gr.Button("Download & Extract", variant="primary", scale=1)
            status_text = gr.Textbox(label="Status", interactive=False)
        
        # Tab 2: Selection & Tagging
        with gr.Tab("🏷️ Select & Tag"):
            gr.Markdown("### Select Images and Tag with WD-Tagger")
            
            with gr.Row():
                with gr.Column(scale=2):
                    height_input = gr.Slider(
                        label="Image Height (pixels)",
                        minimum=256,
                        maximum=2048,
                        value=1024,
                        step=64
                    )
                
                with gr.Column(scale=1):
                    refresh_gallery_btn = gr.Button("Refresh Gallery", scale=1)
            
            image_gallery = gr.Gallery(
                label="Images (Click to select)",
                show_label=True,
                elem_id="gallery",
                columns=3,
                rows=2,
                object_fit="scale-down",
                height="auto"
            )
            
            gallery_status = gr.Textbox(label="Gallery Status", interactive=False)
            
            with gr.Row():
                tag_btn = gr.Button("🏷️ Tag Selected Images (WD-Tagger)", variant="primary")
                tag_status = gr.Textbox(label="Tag Status", interactive=False)
            
            captions_display = gr.Textbox(
                label="Captions (JSON format)",
                interactive=False,
                lines=10,
                max_lines=20
            )
        
        # Tab 3: Export
        with gr.Tab("💾 Export"):
            gr.Markdown("### Export Selected Images with Captions")
            
            with gr.Row():
                with gr.Column():
                    export_status = gr.Textbox(label="Export Status", interactive=False)
                with gr.Column():
                    export_btn = gr.Button("📦 Export as ZIP", variant="primary", size="lg")
            
            export_file = gr.File(label="Download ZIP")
            
            gr.Markdown("**Instructions:**\n1. Load images from a zip file or HF dataset\n2. Select images in the gallery\n3. Optionally tag images with WD-Tagger\n4. Export selected images with captions as ZIP")
    
    # Event handlers
    def download_click(url, hf_token):
        images, file_names, status = download_and_extract_images(url, hf_token)
        return images, file_names, status, [], "No images loaded", {}
    
    download_btn.click(
        fn=download_click,
        inputs=[url_input, hf_token_input],
        outputs=[images_state, file_names_state, status_text, image_gallery, gallery_status, captions_state]
    )
    
    def refresh_gallery(images_list, height):
        gallery_imgs, gallery_status_txt = create_gallery_with_checkboxes(images_list, [], height)
        return gallery_imgs, gallery_status_txt
    
    refresh_gallery_btn.click(
        fn=refresh_gallery,
        inputs=[images_state, height_input],
        outputs=[image_gallery, gallery_status]
    )
    
    def tag_click(selected_indices, images_list, file_names_list):
        if not selected_indices:
            return "", "❌ No images selected"
        
        tag_status_msg, captions = tag_images_wd(selected_indices, images_list, file_names_list)
        captions_json = json.dumps(captions, indent=2)
        return captions_json, tag_status_msg
    
    tag_btn.click(
        fn=tag_click,
        inputs=[image_gallery, images_state, file_names_state],
        outputs=[captions_display, tag_status]
    )
    
    def export_click(selected_indices, images_list, file_names_list, captions_dict, height):
        file_data, export_msg = export_selected(selected_indices, images_list, file_names_list, captions_dict, height)
        return file_data, export_msg
    
    export_btn.click(
        fn=export_click,
        inputs=[image_gallery, images_state, file_names_state, captions_state, height_input],
        outputs=[export_file, export_status]
    )


if __name__ == "__main__":
    demo.launch(share=False)