import re, requests, uuid, os
import gradio as gr

DOWNLOAD_FOLDER = "downloads"
os.makedirs(DOWNLOAD_FOLDER, exist_ok=True)

def normalize_udio_input(input_text):
    """Accept song page URL, embed URL, or iframe snippet, return embed URL"""
    match = re.search(r'src=["\'](https://www\.udio\.com/embed/[^"\']+)["\']', input_text)
    if match:
        return match.group(1)
    if input_text.startswith("https://www.udio.com/embed/"):
        return input_text.split()[0]
    song_match = re.search(r'https://www\.udio\.com/songs/([A-Za-z0-9]+)', input_text)
    if song_match:
        song_id = song_match.group(1)
        return f"https://www.udio.com/embed/{song_id}"
    return None

def get_mp3(embed_url):
    """Extract MP3 URL from Udio embed"""
    if not embed_url:
        return "❌ Invalid or missing URL", None, None

    try:
        r = requests.get(embed_url)
        html = r.text
        mp3_match = re.search(r'https://storage\.googleapis\.com/.*?\.mp3', html)
        if mp3_match:
            mp3_url = mp3_match.group(0)
            return f"✅ Found MP3: {mp3_url}", embed_url, mp3_url
        else:
            return "⚠️ MP3 link not found on embed page.", embed_url, None
    except Exception as e:
        return f"❌ Error fetching embed page: {e}", embed_url, None

def download_mp3(mp3_url):
    """Download MP3 and return local path"""
    if not mp3_url:
        return None
    try:
        filename = f"{uuid.uuid4()}.mp3"
        filepath = os.path.join(DOWNLOAD_FOLDER, filename)
        r = requests.get(mp3_url)
        with open(filepath, "wb") as f:
            f.write(r.content)
        return filepath
    except Exception:
        return None

def process_input(user_input):
    embed_url = normalize_udio_input(user_input)
    msg, embed, mp3 = get_mp3(embed_url)
    
    embed_html = f'<iframe src="{embed}" width="400" height="200" frameborder="0"></iframe>' if embed else ""
    mp3_html = f'<a href="{mp3}" target="_blank">{mp3}</a>' if mp3 else ""
    local_file = download_mp3(mp3) if mp3 else None

    return msg, embed_html, mp3_html, local_file

demo = gr.Interface(
    fn=process_input,
    inputs=gr.Textbox(label="🎶 Udio Song / Embed / Snippet", placeholder="Paste Udio URL or embed iframe here"),
    outputs=[
        gr.Textbox(label="Status / Message"),
        gr.HTML(label="Embed Preview"),
        gr.HTML(label="Direct MP3 Link"),
        gr.File(label="Download MP3")
    ],
    title="🎵 Udio Fetcher + Downloader",
    description="Paste any Udio song link or embed code and get the playable embed, direct MP3 link, and downloadable file."
)

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