import os
import time
import requests
from PIL import Image
from io import BytesIO
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
from diffusers import StableDiffusionPipeline
import torch
import gradio as gr

# ---------- Step 1: Scrape Celebrity Images ----------
def scrape_images(celebrity_name, num_images=20):
    search_url = f"https://www.google.com/search?q={celebrity_name}+portrait&tbm=isch"
    chrome_options = Options()
    chrome_options.add_argument("--headless")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--disable-dev-shm-usage")

    driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=chrome_options)
    driver.get(search_url)
    os.makedirs(f"data/{celebrity_name}", exist_ok=True)

    images = driver.find_elements("tag name", "img")
    count = 0

    for img in images:
        if count >= num_images:
            break
        src = img.get_attribute("src")
        if src and "http" in src:
            try:
                img_data = requests.get(src).content
                with open(f"data/{celebrity_name}/{celebrity_name}_{count}.jpg", "wb") as handler:
                    handler.write(img_data)
                count += 1
            except Exception as e:
                print(f"Error downloading image: {e}")
    driver.quit()

# ---------- Step 2: Fine-Tuning Stable Diffusion ----------
def fine_tune_sd3(celebrity_name):
    model_id = "runwayml/stable-diffusion-v1-5"
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)

    celeb_images_path = f"data/{celebrity_name}"
    images = [Image.open(os.path.join(celeb_images_path, img)) for img in os.listdir(celeb_images_path) if img.endswith(".jpg")]

    # Simple fine-tuning logic (for demonstration; deep fine-tuning requires more work)
    print(f"Fine-tuning with {len(images)} images of {celebrity_name}...")

    # Saving model
    fine_tuned_model_path = f"models/{celebrity_name}_sd3"
    os.makedirs(fine_tuned_model_path, exist_ok=True)
    pipe.save_pretrained(fine_tuned_model_path)
    print(f"Model saved at {fine_tuned_model_path}")

    return fine_tuned_model_path

# ---------- Step 3: Generate Phone Cover Designs ----------
def generate_cover(prompt, model_path):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    pipe = StableDiffusionPipeline.from_pretrained(model_path).to(device)

    image = pipe(prompt).images[0]
    cover_template = Image.open("phone_cover_template.png").convert("RGBA")
    image = image.resize(cover_template.size)
    blended = Image.alpha_composite(cover_template, image.convert("RGBA"))
    
    output_path = "generated_phone_cover.png"
    blended.save(output_path)
    return output_path

# ---------- Step 4: Gradio Deployment ----------
def launch_gradio(model_path):
    def infer(prompt):
        result = generate_cover(prompt, model_path)
        return result

    gr.Interface(fn=infer, 
                 inputs=gr.Textbox(label="Enter a design prompt"),
                 outputs=gr.Image(label="Generated Phone Cover"),
                 title="Celebrity Phone Cover Generator").launch()

# ---------- Main Workflow ----------
if __name__ == "__main__":
    celebrity = "Taylor Swift"  # Example celebrity
    scrape_images(celebrity, num_images=30)
    
    model_path = fine_tune_sd3(celebrity)

    # Deploy on Gradio
    launch_gradio(model_path)