from pathlib import Path

import gradio as gr
import pi_heif
import spaces
import requests
import base64
import io
import os
import json
from huggingface_hub import hf_hub_download
from PIL import Image

css = """
footer {display: none !important;}
"""

TITLE = """
<p>
    Erase unwanted watermarks automatically with AI, no manual editing needed.
</p><p>
    For premium-quality results,
    <a href="https://nowatermark.cloud">try NoWatermark API</a>
    — it's free to test!
</p>
"""

LIKE_BUTTON = """
<div style="text-align: left; margin: 5px 0 0 0;">
  <a href="https://huggingface.co/spaces/abdul9999/NoWatermark" target="_blank"
     style="text-decoration: none;">
    <button style="
      background-color: #ff4081;
      color: white;
      border: none;
      padding: 4px 10px;
      font-size: 13px;
      border-radius: 6px;
      cursor: pointer;
    ">
      ❤️ Like
    </button>
  </a>
</div>
"""

# HEADERS = os.getenv("HEADERS")
API_HOST = os.getenv("API_HOST")

def remove_watermark(img_bytes):
    """
    headers = json.loads(HEADERS)

    files = {
        "original_preview_image": ("1.webp", img_bytes, "image/webp"),
    }

    data = {
        "zoom_factor": 1,
        "predict_mode": "old",
    }
    
    resp = requests.post(
        API_HOST,
        headers=headers,
        files=files,
        data=data,
    )
    """
    files = {"file": ("input_image.webp", img_bytes, "image/webp")}
    
    resp = requests.post(
        API_HOST,
        files=files,
    )
    
    if resp.status_code != 200:
        print(f"Error: {resp.status_code} - {resp.text}")
        return None

    result_img = Image.open(io.BytesIO(resp.content)).convert("RGB")
    return result_img

def process(input_image):
    """
    input_image: PIL.Image from Gradio
    returns: PIL.Image (dewatermarked)
    """
    # Convert PIL image to bytes (WebP, since your API expects webp)
    img_bytes_io = io.BytesIO()
    input_image.save(img_bytes_io, format="WEBP")
    img_bytes = img_bytes_io.getvalue()

    # Call your dewatermark function
    result_img = remove_watermark(img_bytes)
    if result_img is None:
        return [input_image, input_image]
        # return None  # or return input_image as fallback

    # Convert bytes back to PIL image
    return [input_image, result_img]
    
with gr.Blocks(css=css) as demo:
    gr.HTML(LIKE_BUTTON)
    gr.HTML(TITLE)

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil", label="Input Image")
            run_button = gr.ClearButton(components=None, value="Remove Watermark")
        with gr.Column():
            output_slider = gr.ImageSlider(label="Before / After", max_height=1200, show_fullscreen_button=False)
            run_button.add(output_slider)

    run_button.click(fn=process, inputs=input_image, outputs=output_slider, api_name=False)

    gr.Examples(
        examples=[
            "examples/1.jpg",
            "examples/2.jpg",
            "examples/3.jpg",
            "examples/4.jpg",
            "examples/5.jpg",
            "examples/6.jpg",
            "examples/7.jpg",
            "examples/8.jpg",
            "examples/9.jpg",
            "examples/10.jpg",
        ],
        inputs=[input_image],
        outputs=output_slider,
        fn=process,
        cache_examples=True,
        cache_mode="lazy",
        run_on_click=False,
    )

demo.launch(share=False, ssr_mode=False, show_api=False)