import os
import spaces
import gradio as gr
import io, base64, random
from gradio_client import Client
from transformers import pipeline
from PIL import Image, ImageDraw, ImageFont

HF_TOKEN = os.getenv("hf_token")
client = Client("mrfakename/Z-Image-Turbo")


# Use GPU if available in the Space
DEVICE = 0 if os.environ.get("CUDA_VISIBLE_DEVICES") not in (None, "", "-1") else -1

nsfw_detector = pipeline(
    "image-classification",
    model="Falconsai/nsfw_image_detection",
    device=DEVICE,
)

def get_image(image_prompt:str):
  image_result = client.predict(
    prompt=image_prompt,
    height=512,
    width=512,
    num_inference_steps=9,
    seed=42,
    randomize_seed=True,
    api_name="/generate_image"
  )
  try:
      image_result_path = image_result[0]
      # Define a path for the converted PNG image
      png_path = image_result_path.replace('.webp', '.png')

      # Open the webp image and save it as png
      with Image.open(image_result_path) as img:
          img.save(png_path, 'PNG')

      return {"status": 200, "image_path": png_path}

  except Exception as e:
      return {"status": 400, "image_path": None, "error": str(e)}

def make_image(text, w=768, h=512):
    bg = (random.randint(150,230), random.randint(150,230), random.randint(150,230))
    img = Image.new("RGB", (w, h), bg)
    draw = ImageDraw.Draw(img)

    try:
        font = ImageFont.truetype("DejaVuSans.ttf", 28)
    except:
        font = ImageFont.load_default()

    draw.text((40, 40), text, fill=(0,0,0), font=font)
    return img

def pil_to_markdown_image(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
    return f"![image](data:image/png;base64,{b64})"

def is_image_nsfw(img: Image.Image, threshold=0.6) -> bool:
    preds = nsfw_detector(img)
    # Example output: [{'label': 'nsfw', 'score': 0.93}, {'label': 'normal', 'score': 0.07}]
    for p in preds:
        if p["label"].lower() in ["nsfw", "porn", "sexy", "naked", "girl"] and p["score"] >= threshold:
            return True
    return False

@spaces.GPU
def respond(image_prompt, history):
    img_dict = get_image(image_prompt)
    if img_dict["status"] == 200 and not is_image_nsfw(img_dict["image_path"]):
        img_path = img_dict["image_path"]
        with Image.open(img_path) as img:
            img = img.convert("RGBA")  # safe for transparency
            md_img = pil_to_markdown_image(img)
            return md_img  # pure markdown with embedded binary image
    else:
        return pil_to_markdown_image(make_image(image_prompt))

demo = gr.ChatInterface(
    fn=respond,
    save_history=True,
    type="messages", # Recommended for newer Gradio versions
    title="Chat → Image (Binary in Markdown)",
    description="Each message returns an image embedded as base64 in Markdown."
)

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