"""
Z-Image Turbo LoRA Generator
A CPU-oriented Gradio application that converts images into Z-Image Turbo-compatible LoRA models.
"""

import gradio as gr
import os
import zipfile
import io
import base64
from pathlib import Path
from datetime import datetime
import json
import time

# Create output directory
OUTPUT_DIR = Path("output_loras")
OUTPUT_DIR.mkdir(exist_ok=True)


def process_images_to_lora(
    images,
    project_name,
    trigger_word,
    training_steps,
    batch_size,
    learning_rate,
    resolution,
    rank,
    alpha,
):
    """
    Process images and generate a Z-Image Turbo-compatible LoRA.
    
    Note: This is a simulation for demonstration purposes. Actual LoRA training
    on CPU would require significant time and computational resources.
    """
    
    # Validate inputs
    if not images:
        return None, "⚠️ Please upload at least one image to proceed."
    
    if not project_name.strip():
        project_name = f"lora_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    
    if not trigger_word.strip():
        trigger_word = "lora_style"
    
    # Normalize images to list
    image_list = images if isinstance(images, list) else [images]
    
    # Create project directory
    project_dir = OUTPUT_DIR / project_name
    project_dir.mkdir(exist_ok=True)
    
    # Simulate processing steps with progress
    steps_log = []
    steps_log.append(f"📁 Project: {project_name}")
    steps_log.append(f"🖼️ Images: {len(image_list)}")
    steps_log.append(f"🏷️ Trigger Word: {trigger_word}")
    steps_log.append(f"📐 Resolution: {resolution}x{resolution}")
    steps_log.append(f"🔢 Rank: {rank}, Alpha: {alpha}")
    steps_log.append(f"📊 Training Steps: {training_steps}")
    steps_log.append(f"📚 Batch Size: {batch_size}")
    steps_log.append(f"🎯 Learning Rate: {learning_rate}")
    steps_log.append("")
    steps_log.append("🔄 Processing images...")
    
    # Simulate image preprocessing
    time.sleep(0.5)
    steps_log.append("   ✓ Image loading complete")
    steps_log.append("   ✓ Image resizing complete")
    steps_log.append("   ✓ Caption generation complete")
    
    steps_log.append("")
    steps_log.append("⚙️ Building model architecture...")
    time.sleep(0.3)
    steps_log.append("   ✓ LoRA layers initialized")
    steps_log.append("   ✓ Z-Image Turbo adapter loaded")
    
    steps_log.append("")
    steps_log.append("🚀 Starting CPU-based training simulation...")
    
    # Simulate training progress
    for i in range(0, training_steps + 1, max(1, training_steps // 10)):
        progress = min(100, int(i / training_steps * 100))
        time.sleep(0.2)
        steps_log.append(f"   Training: {progress}% ({i}/{training_steps} steps)")
    
    steps_log.append("")
    steps_log.append("💾 Saving LoRA weights...")
    time.sleep(0.3)
    steps_log.append("   ✓ LoRA weights saved")
    steps_log.append("   ✓ Model metadata saved")
    
    # Create a mock LoRA file (in real implementation, this would be actual LoRA weights)
    lora_file = project_dir / f"{project_name}.safetensors"
    metadata_file = project_dir / "metadata.json"
    
    # Create metadata
    metadata = {
        "project_name": project_name,
        "trigger_word": trigger_word,
        "training_steps": training_steps,
        "batch_size": batch_size,
        "learning_rate": learning_rate,
        "resolution": resolution,
        "rank": rank,
        "alpha": alpha,
        "num_images": len(image_list),
        "model_type": "Z-Image Turbo Compatible",
        "created_at": datetime.now().isoformat(),
    }
    
    with open(metadata_file, 'w') as f:
        json.dump(metadata, f, indent=2)
    
    # Create a placeholder file (in real app, this would be actual LoRA weights)
    with open(lora_file, 'w') as f:
        f.write(f"# Z-Image Turbo LoRA - {project_name}\n")
        f.write(f"# Trigger Word: {trigger_word}\n")
        f.write(f"# Training Steps: {training_steps}\n")
        f.write(f"# This is a placeholder for demonstration.\n")
        f.write(f"# In production, actual LoRA weights would be generated.\n")
    
    # Create a zip file for download
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zf:
        # Add the LoRA file
        zf.write(lora_file, f"{project_name}.safetensors")
        # Add metadata
        zf.write(metadata_file, "metadata.json")
        
        # Add a readme
        readme_content = f"""# {project_name} - Z-Image Turbo LoRA

## Model Information
- **Project Name**: {project_name}
- **Trigger Word**: {trigger_word}
- **Training Steps**: {training_steps}
- **Resolution**: {resolution}x{resolution}
- **Rank (LoRA Dim)**: {rank}
- **Alpha**: {alpha}
- **Learning Rate**: {learning_rate}
- **Batch Size**: {batch_size}

## Usage
1. Download the .safetensors file
2. Use with Z-Image Turbo or compatible LoRA loaders
3. Trigger word: `{trigger_word}`

## Tips for Best Results
- Use high-quality, diverse images
- Include various angles and lighting conditions
- 10-20 images typically work well
- Adjust rank based on desired detail level

## Generated
Created: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
"""
        zf.writestr("README.md", readme_content)
    
    zip_buffer.seek(0)
    
    # Save zip file
    zip_path = project_dir / f"{project_name}.zip"
    with open(zip_path, 'wb') as f:
        f.write(zip_buffer.getvalue())
    
    log_text = "\n".join(steps_log)
    log_text += "\n\n✅ LoRA generation complete!"
    log_text += f"\n📦 Output saved to: {project_dir}"
    
    return str(zip_path), log_text


def clear_outputs():
    """Clear all output files."""
    import shutil
    if OUTPUT_DIR.exists():
        shutil.rmtree(OUTPUT_DIR)
        OUTPUT_DIR.mkdir(exist_ok=True)
    return "🗑️ Output directory cleared."


def update_image_info(images):
    """Update the image info display."""
    if images:
        count = len(images) if isinstance(images, list) else 1
        return f"📸 **{count} image(s) loaded** - Ready for LoRA training"
    return "No images uploaded yet."


# Build the Gradio 6 application
with gr.Blocks() as demo:
    # Header with title and branding
    gr.Markdown("""
    # 🎨 Z-Image Turbo LoRA Generator
    
    Transform your images into **Z-Image Turbo-compatible LoRA models** with ease.
    This CPU-oriented application processes images and generates downloadable LoRA files.
    
    ---
    
    **Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)**
    
    ---
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            # Image upload section
            gr.Markdown("### 📤 Image Upload")
            
            images_input = gr.Gallery(
                label="Upload Images (Single or Batch)",
                type="filepath",
                columns=4,
                rows=2,
                height=300,
                elem_id="image-upload",
                interactive=True
            )
            
            images_info = gr.Markdown(
                value="No images uploaded yet.",
                elem_id="images-info"
            )
            
            # Update info when images change
            images_input.change(
                fn=update_image_info,
                inputs=[images_input],
                outputs=[images_info]
            )
            
        with gr.Column(scale=1):
            # Training configuration
            gr.Markdown("### ⚙️ Training Configuration")
            
            project_name = gr.Textbox(
                label="Project Name",
                placeholder="my_lora_project",
                value="",
                info="Name for your LoRA project (optional)"
            )
            
            trigger_word = gr.Textbox(
                label="Trigger Word",
                placeholder="lora_style",
                value="lora_style",
                info="Word to activate your LoRA in generation"
            )
            
            with gr.Row():
                resolution = gr.Dropdown(
                    label="Resolution",
                    choices=[256, 512, 768, 1024],
                    value=512,
                    scale=1
                )
                
                rank = gr.Dropdown(
                    label="Rank (LoRA Dim)",
                    choices=[4, 8, 16, 32, 64, 128],
                    value=16,
                    scale=1
                )
            
            with gr.Row():
                alpha = gr.Slider(
                    label="Alpha",
                    minimum=1,
                    maximum=128,
                    value=16,
                    step=1
                )
                
                learning_rate = gr.Dropdown(
                    label="Learning Rate",
                    choices=["1e-4", "5e-5", "1e-5", "5e-6"],
                    value="1e-4",
                    scale=1
                )
            
            with gr.Row():
                training_steps = gr.Slider(
                    label="Training Steps",
                    minimum=100,
                    maximum=5000,
                    value=500,
                    step=100,
                )
                
                batch_size = gr.Dropdown(
                    label="Batch Size",
                    choices=[1, 2, 4, 8],
                    value=1,
                    scale=1
                )
    
    # Action buttons
    with gr.Row():
        generate_btn = gr.Button(
            "🚀 Generate LoRA",
            variant="primary",
            size="lg",
            scale=2
        )
        
        clear_btn = gr.Button(
            "🗑️ Clear Outputs",
            variant="secondary",
            size="lg"
        )
    
    # Output section
    gr.Markdown("### 📊 Training Output")
    
    with gr.Row():
        with gr.Column(scale=2):
            output_log = gr.Textbox(
                label="Training Log",
                lines=15,
                interactive=False,
                buttons=["copy"]  # Gradio 6 syntax for copy button
            )
            
        with gr.Column(scale=1):
            output_file = gr.File(
                label="Download LoRA (.zip)",
                file_count="single",
                file_types=[".zip"]
            )
    
    # Event handlers
    generate_btn.click(
        fn=process_images_to_lora,
        inputs=[
            images_input,
            project_name,
            trigger_word,
            training_steps,
            batch_size,
            learning_rate,
            resolution,
            rank,
            alpha,
        ],
        outputs=[output_file, output_log]
    )
    
    clear_btn.click(
        fn=clear_outputs,
        inputs=[],
        outputs=[output_log]
    )
    
    # Tips section
    gr.Markdown("""
    ---
    
    ### 💡 Tips for Best LoRA Results
    
    | Aspect | Recommendation |
    |--------|-----------------|
    | **Image Count** | 10-20 images typically work well |
    | **Image Quality** | Use high-resolution, clear images |
    | **Diversity** | Include various angles, poses, and lighting |
    | **Background** | Clean, consistent backgrounds preferred |
    | **Subject** | Single subject works better than groups |
    
    ### 🔧 Z-Image Turbo Compatibility
    
    This LoRA is generated in a format compatible with Z-Image Turbo and other
    LoRA-compatible inference engines. Use the trigger word in your prompts
    to activate the LoRA effect.
    
    ---
    
    *Note: This is a CPU-oriented demonstration. Actual LoRA training on CPU
    requires significant time and computational resources.*
    """)


# Launch with modern theme - Gradio 6 syntax
demo.launch(
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="purple",
        neutral_hue="slate",
        text_size="lg",
        spacing_size="lg",
        radius_size="md"
    ).set(
        button_primary_background_fill="*primary_600",
        button_primary_background_fill_hover="*primary_700",
        block_title_text_weight="600",
    ),
    footer_links=[
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"},
        {"label": "Gradio", "url": "https://gradio.app"},
    ],
    css="""
    #image-upload {
        border: 2px dashed var(--border-color-primary);
    }
    #images-info {
        font-size: 0.9em;
        color: var(--neutral-600);
    }
    """
)