"""
Pixagram Pixel Art Converter
With Commercial FaceID Support

Gradio web interface for converting images to pixel art
while preserving face identity.

All AI components use commercially-permissive licenses:
- SDXL: Stability AI License
- AuraFace: Commercial OK (fal.ai)
- OpenCV: BSD License
- MediaPipe: Apache 2.0

Copyright (c) 2024 Pixagram SA
"""
import gradio as gr
import spaces
import torch
from model import ModelHandler
from generator import Generator
from config import Config
from utils import visualize_face_detection, has_face

# ============================================================
# 1. Initialize Models Globally (in RAM)
# ============================================================
print("=" * 60)
print("Initializing Pixagram Application (with FaceID)")
print("=" * 60)

handler = ModelHandler()
handler.load_models()
gen = Generator(handler)

print("\n[READY] Application initialized successfully!\n")


# ============================================================
# 2. Define GPU-enabled Inference Function
# ============================================================
@spaces.GPU(duration=30)
def process_img(
    image, 
    prompt,
    negative_prompt,
    cfg_scale, 
    steps, 
    img_strength, 
    depth_strength,
    # FaceID parameters
    use_faceid,
    faceid_strength,
    face_padding,
    # Other
    seed
):
    """
    Process image and generate pixel art with optional face identity preservation.
    """
    if image is None:
        raise gr.Error("Please upload an image first.")
    
    try:
        print("\n" + "=" * 60)
        print("Starting Generation")
        print("=" * 60)
        
        # Log FaceID status (detection happens inside generator)
        if use_faceid:
            print(f"[FaceID] Enabled (strength: {faceid_strength})")
        else:
            print("[FaceID] Disabled by user")
        
        result = gen.predict(
            image, 
            user_prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=cfg_scale,
            num_inference_steps=steps,
            img2img_strength=img_strength,
            depth_strength=depth_strength,
            seed=seed,
            # FaceID
            use_faceid=use_faceid,
            faceid_strength=faceid_strength,
            face_padding=face_padding
        )
        
        print("=" * 60)
        print("Generation Complete!")
        print("=" * 60 + "\n")
        
        return result
        
    except Exception as e:
        print(f"[ERROR] Generation failed: {e}")
        import traceback
        traceback.print_exc()
        raise gr.Error(f"An error occurred: {str(e)}")


def check_face_preview(image):
    """
    Check if face is detected and show preview with bounding box.
    """
    if image is None:
        return None, "No image uploaded"
    
    face_detected = has_face(image)
    
    if face_detected:
        preview = visualize_face_detection(image)
        status = "✅ Face detected - identity will be preserved"
    else:
        preview = image
        status = "ℹ️ No face detected - standard generation"
    
    return preview, status


# ============================================================
# 3. Build Gradio Interface
# ============================================================
css = """
.face-status-positive { color: #22c55e; font-weight: bold; }
.face-status-neutral { color: #6b7280; }
.commercial-badge { 
    background: linear-gradient(135deg, #10b981, #059669);
    color: white;
    padding: 4px 12px;
    border-radius: 4px;
    font-size: 12px;
    font-weight: bold;
}
"""

with gr.Blocks(
    title="Pixagram - Pixel Art with FaceID", 
) as demo:
    
    gr.Markdown(
        """
        # 🎮 Pixagram - Image to Pixel Art
        ### With Face Identity Preservation
        
        Upload any image to convert it into retro pixel art style.
        **Faces are automatically detected and preserved** during conversion.
        
        <span class="commercial-badge">✓ Commercial License</span>
        """
    )
    
    with gr.Row():
        # =====================
        # LEFT COLUMN - Inputs
        # =====================
        with gr.Column(scale=2):
            input_img = gr.Image(
                type="pil", 
                label="Input Image",
                height=400
            )
            
            # Face detection status
            with gr.Row():
                face_preview = gr.Image(
                    type="pil",
                    label="Face Detection Preview",
                    height=150,
                    visible=False
                )
                face_status = gr.Markdown(
                    value="Upload an image to check for faces",
                    elem_classes=["face-status-neutral"]
                )
            
            prompt = gr.Textbox(
                label="Prompt (Optional)", 
                placeholder="Leave empty for auto-captioning...",
                info="Pixel art trigger words are added automatically."
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt (Optional)", 
                placeholder="e.g., blurry, text, watermark...",
                value=Config.DEFAULT_NEGATIVE_PROMPT
            )
            
            # FaceID Settings (prominent)
            with gr.Accordion("🎭 Face Identity Settings", open=True):
                gr.Markdown(
                    """
                    When enabled, detected faces will have their identity preserved
                    during pixel art conversion. Works with the LoRA style!
                    """
                )
                
                use_faceid = gr.Checkbox(
                    label="Enable Face Identity Preservation",
                    value=Config.AUTO_FACEID,
                    info="Automatically detect and preserve face identity"
                )
                
                faceid_strength = gr.Slider(
                    minimum=0.0,
                    maximum=1.5,
                    step=0.05,
                    value=Config.DEFAULT_FACEID_STRENGTH,
                    label="Face Identity Strength",
                    info="How strongly to preserve the face (0.7-1.0 recommended)"
                )
                
                face_padding = gr.Slider(
                    minimum=0.1,
                    maximum=0.6,
                    step=0.05,
                    value=Config.FACE_DETECTION_PADDING,
                    label="Face Detection Padding",
                    info="Padding around detected face for better context"
                )
            
            # Advanced Settings
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                seed = gr.Number(
                    label="Seed", 
                    value=-1, 
                    info="-1 for random", 
                    precision=0
                )
                
                cfg_scale = gr.Slider(
                    minimum=1.0, 
                    maximum=5.0, 
                    step=0.1, 
                    value=Config.CGF_SCALE, 
                    label="CFG Scale",
                    info="Higher = more prompt adherence"
                )
                
                steps = gr.Slider(
                    minimum=4, 
                    maximum=20, 
                    step=1, 
                    value=Config.STEPS_NUMBER, 
                    label="Inference Steps",
                    info="More steps = better quality (LCM: 4-10 is enough)"
                )
                
                img_strength = gr.Slider(
                    minimum=0.1, 
                    maximum=1.0, 
                    step=0.05, 
                    value=Config.IMG_STRENGTH, 
                    label="Image Transformation Strength",
                    info="How much to change from original"
                )
                
                depth_strength = gr.Slider(
                    minimum=0.0, 
                    maximum=1.0, 
                    step=0.05, 
                    value=Config.DEPTH_STRENGTH, 
                    label="Depth Control Strength",
                    info="How much to follow the depth structure"
                )
            
            run_btn = gr.Button(
                "🎨 Generate Pixel Art", 
                variant="primary",
                size="lg"
            )
        
        # =====================
        # RIGHT COLUMN - Output
        # =====================
        with gr.Column(scale=1):
            output_img = gr.Image(
                label="Pixel Art Result",
                height=500
            )
            
            gr.Markdown(
                """
                ### 📋 License Information
                
                This tool uses **commercially-licensed** components:
                
                | Component | License |
                |-----------|---------|
                | Face Detection | BSD / Apache 2.0 |
                | Face Encoding | Commercial (AuraFace) |
                | Style Model | Your LoRA |
                | Base Model | Stability AI |
                
                ✅ **Safe for commercial use**
                """
            )
    
    # =====================
    # Examples
    # =====================
    gr.Markdown("### 📸 Examples")
    gr.Examples(
        examples=[
            ["A portrait of a person", 0.85, 0.65],
            ["A character in fantasy setting", 0.90, 0.55],
            ["A person in cyberpunk style", 0.80, 0.70],
        ],
        inputs=[prompt, faceid_strength, img_strength],
        label="Try these prompts"
    )
    
    # =====================
    # Event Handlers
    # =====================
    
    # Check face when image is uploaded
    input_img.change(
        fn=check_face_preview,
        inputs=[input_img],
        outputs=[face_preview, face_status]
    )
    
    # Show/hide face preview based on checkbox
    use_faceid.change(
        fn=lambda x: gr.update(visible=x),
        inputs=[use_faceid],
        outputs=[face_preview]
    )
    
    # Main generation
    all_inputs = [
        input_img, 
        prompt, 
        negative_prompt,
        cfg_scale, 
        steps, 
        img_strength,
        depth_strength,
        use_faceid,
        faceid_strength,
        face_padding,
        seed
    ]
    
    run_btn.click(
        fn=process_img, 
        inputs=all_inputs, 
        outputs=[output_img]
    )


# ============================================================
# 4. Launch the App
# ============================================================
if __name__ == "__main__":
    demo.queue(max_size=20)
    demo.launch(
        share=False,
        show_error=True,
        theme=gr.themes.Soft(),
        css=css,
        ssr_mode=False
    )