import gradio as gr
from gradio_client import Client, handle_file
import os
import smtplib
from email.message import EmailMessage

BACKEND_SPACE = "CleanSong-AI/Main-tool-backend-main"

# --- Email sender (example with SMTP) ---
SMTP_SERVER = os.getenv("SMTP_SERVER")      # e.g., smtp.gmail.com
SMTP_PORT = int(os.getenv("SMTP_PORT", 587))
SMTP_USER = os.getenv("SMTP_USER")
SMTP_PASS = os.getenv("SMTP_PASS")

def send_email(to_email, file_path):
    msg = EmailMessage()
    msg['Subject'] = "Your CleanSong AI track is ready!"
    msg['From'] = SMTP_USER
    msg['To'] = to_email
    msg.set_content("Hi! Your cleaned track is ready. See the attachment!")

    with open(file_path, 'rb') as f:
        file_data = f.read()
        file_name = os.path.basename(file_path)

    msg.add_attachment(file_data, maintype='audio', subtype='mpeg', filename=file_name)

    with smtplib.SMTP(SMTP_SERVER, SMTP_PORT) as server:
        server.starttls()
        server.login(SMTP_USER, SMTP_PASS)
        server.send_message(msg)

# --- Main processing function ---
def clean_song(file_path, email=None):
    if file_path is None:
        return None

    # Send file to backend
    client = Client(BACKEND_SPACE)
    result = client.predict(
        handle_file(file_path),
        api_name="/predict"
    )

    # Normalize HF return
    if isinstance(result, (list, tuple)):
        result = result[0]
    if isinstance(result, dict) and "name" in result:
        result = result["name"]

    # Send email if provided
    if email:
        try:
            send_email(email, result)
        except Exception as e:
            print("Failed to send email:", e)

    return result

# ------------------------------------------------------------
# UI
# ------------------------------------------------------------
iface = gr.Interface(
    fn=clean_song,
    inputs=[
        gr.Audio(type="filepath", label="Upload Song"),
        gr.Textbox(label="Email (optional)", placeholder="Enter your email to get notified")
    ],
    outputs=gr.Audio(type="filepath", label="Cleaned Song"),
    title="CleanSong AI - Prototype Frontend",
    description="Upload a song and get a cleaned version. Optionally enter your email to receive the cleaned track."
)

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