import os
import io
import zipfile
import sys
import traceback
import time
import tempfile
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

# --- DEBUG SETUP ---
sys.stdout.flush()
print("=== APP STARTING UP ===")
print(f"Current Working Directory: {os.getcwd()}")
print(f"Python Version: {sys.version}")

# CRITICAL FIX: Catch import errors
try:
    import cv2
    import numpy as np
    from PIL import Image
    from rembg import remove, new_session
    from huggingface_hub import InferenceClient
    import gradio as gr
    print("All libraries imported successfully.")
except Exception as e:
    print("CRITICAL IMPORT ERROR:", e)
    print(traceback.format_exc())
    sys.exit(1)

# Get your HF Token
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    print("WARNING: HF_TOKEN environment variable is NOT SET. AI Backgrounds will fail.")
else:
    print("HF_TOKEN found.")
client = InferenceClient(token=HF_TOKEN)

# Preload the rembg model
print("Loading rembg model... (this takes ~20 seconds)")
try:
    session = new_session() 
    print("rembg model loaded successfully!")
except Exception as e:
    print("ERROR loading rembg model:", e)
    print(traceback.format_exc())
    session = None

def enhance_lighting(pil_img):
    try:
        cv_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
        lab = cv2.cvtColor(cv_img, cv2.COLOR_BGR2LAB)
        l, a, b = cv2.split(lab)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        l = clahe.apply(l)
        lab = cv2.merge([l, a, b])
        cv_img = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
        return Image.fromarray(cv2.cvtColor(cv_img, cv2.COLOR_BGR2RGB))
    except Exception as e:
        print(f"Lighting fix failed: {e}")
        return pil_img

def fix_aspect_ratio(pil_img, ratio_str):
    w, h = pil_img.size
    ratio_map = {"1:1": 1.0, "4:5": 0.8, "16:9": 1.777}
    target_r = ratio_map[ratio_str]
    current_r = w / h
    if current_r > target_r:
        new_w = int(h * target_r)
        pil_img = pil_img.crop(((w - new_w)//2, 0, (w + new_w)//2, h))
    else:
        new_h = int(w / target_r)
        pil_img = pil_img.crop((0, (h - new_h)//2, w, (h + new_h)//2))
    return pil_img

def process_batch(files, ratio, rm_bg, prompt):
    processed = []
    
    # Check if files is None or empty
    if not files:
        return "❌ Error: No files uploaded. Please upload at least one image."
    
    # Loop through all uploaded files
    for idx, file in enumerate(files):
        try:
            print(f"--- Processing image {idx+1} of {len(files)} ---")
            start_time = time.time()
            
            # 1. Open the image
            img = Image.open(file).convert("RGB")
            print(f"Image opened. Size: {img.size}")
            
            # 2. Fix Lighting
            img = enhance_lighting(img)
            
            # 3. Fix Aspect Ratio
            img = fix_aspect_ratio(img, ratio)
            
            # 4. Remove Background
            if rm_bg:
                if session is None:
                    raise Exception("rembg model failed to load during startup. Check Space Logs.")
                print("Running rembg...")
                img = remove(img, session=session)
                print("rembg completed.")
            
            # 5. Generate AI Background
            if prompt and rm_bg:
                print(f"Generating AI background for prompt: {prompt}")
                # --- FIXED MODEL HERE ---
                bg_img = client.text_to_image(
                    f"Product photography background of {prompt}, elegant, soft studio light, photorealistic, 8k",
                    model="black-forest-labs/FLUX.1-dev"  # Updated model
                )
                bg_img = bg_img.resize(img.size).convert("RGBA")
                final = bg_img.copy()
                final.paste(img, (0, 0), img)
                img = final
                print("AI background composited.")
            
            processed.append(img)
            print(f"Image {idx+1} finished in {time.time() - start_time:.2f} seconds.")
            
        except Exception as e:
            # --- FIXED ERROR HANDLER ---
            error_msg = f"❌ ERROR on image {idx+1} ('{file.name}'):\n{str(e)}"
            print(error_msg)
            print(traceback.format_exc())
            # Return the error as a string directly to the UI
            return error_msg
    
    # --- FINAL RETURN LOGIC ---
    if not processed:
        return None
    
    # Create a temporary directory
    temp_dir = tempfile.mkdtemp()
    zip_path = os.path.join(temp_dir, "jewelry_processed.zip")
    
    # Write the ZIP file to disk with PNG compression
    with zipfile.ZipFile(zip_path, "w") as zf:
        for i, img in enumerate(processed):
            buff = io.BytesIO()
            
            # PNG compression
            if img.width > 3000 or img.height > 3000:
                img.thumbnail((3000, 3000), Image.LANCZOS)
            
            img.save(buff, format="PNG", compress_level=9)
            
            buff.seek(0)
            zf.writestr(f"jewelry_processed_{i+1}.png", buff.getvalue())
    
    # Return the absolute file path string
    return zip_path

# Gradio UI
with gr.Blocks(title="Jewelry Batch Processor") as demo:
    gr.Markdown("## ✨ Free Jewelry Batch Processor")
    
    with gr.Row():
        files = gr.Files(label="Upload Jewelry Images", file_count="multiple")
        ratio = gr.Dropdown(choices=["1:1", "4:5", "16:9"], value="1:1", label="Aspect Ratio")
    
    with gr.Row():
        rm_bg = gr.Checkbox(label="Remove Background", value=True)
        prompt = gr.Textbox(label="AI Background Prompt (leave empty to skip)")
    
    btn = gr.Button("🚀 Process Batch", variant="primary")
    output = gr.File(label="Download Processed ZIP")
    
    btn.click(process_batch, inputs=[files, ratio, rm_bg, prompt], outputs=output)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)