#!/usr/bin/env python3
import base64
import json
import os
import re
import sys
import tempfile
import urllib.request
from datetime import datetime
from gradio_client import Client, handle_file

import contextlib
import io

    
def get_itunes_top10():
    """Fetches top 10 tracks, metadata, and audio preview URLs directly from iTunes RSS."""
    url = "https://itunes.apple.com/us/rss/topsongs/limit=10/json"
    req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})

    try:
        with urllib.request.urlopen(req) as resp:
            data = json.loads(resp.read().decode("utf-8"))

        entries = data["feed"]["entry"]
        entries = entries if isinstance(entries, list) else [entries]

        songs = []
        for idx, entry in enumerate(entries[:10], 1):
            title = entry["im:name"]["label"]
            artist = entry["im:artist"]["label"]
            genre = entry["category"]["attributes"]["label"]

            # Extract 30s audio preview URL
            links = entry.get("link", [])
            preview_url = None
            if isinstance(links, list):
                for l in links:
                    attrs = l.get("attributes", {})
                    if attrs.get("title") == "Preview" or attrs.get("type", "").startswith("audio"):
                        preview_url = attrs.get("href")
                        break
            elif isinstance(links, dict):
                preview_url = links.get("attributes", {}).get("href")

            songs.append({
                "rank": idx,
                "title": title,
                "artist": artist,
                "genre": genre,
                "preview_url": preview_url,
            })
        return songs
    except Exception as e:
        print(f"Error fetching iTunes top tracks: {e}", file=sys.stderr)
        return []


def extract_embedded_chart(res, max_width=220):
    """
    Extracts plot data from Gradio's response (file path, raw SVG, or dict)
    and formats it into a self-contained inline element for Markdown tables.
    """
    # 1. Unpack list/tuple wrappers returned by gradio_client
    if isinstance(res, (tuple, list)):
        for item in res:
            if isinstance(item, (dict, str)):
                res = item
                break

    plot_data = None
    if isinstance(res, dict):
        plot_data = res.get("plot") or res.get("value") or res.get("name") or res.get("path")
    elif isinstance(res, str):
        plot_data = res

    if not plot_data or not isinstance(plot_data, str):
        return "No plot payload"

    # 2. Handle temporary file paths created on disk by gradio_client
    if os.path.isfile(plot_data):
        try:
            if plot_data.lower().endswith(".svg"):
                with open(plot_data, "r", encoding="utf-8") as f:
                    svg_content = f.read()
                svg_clean = re.sub(r"<\?xml.*?\?>", "", svg_content, flags=re.DOTALL).strip()
                svg_single_line = re.sub(r"\s+", " ", svg_clean)
                return f'<div style="max-width: {max_width}px; max-height: 150px; display: inline-block;">{svg_single_line}</div>'
            else:
                # Base64-encode PNG/JPEG files directly into an inline <img> tag
                with open(plot_data, "rb") as f:
                    encoded = base64.b64encode(f.read()).decode("utf-8")
                ext = os.path.splitext(plot_data)[1].lstrip(".").lower() or "png"
                return f'<img src="data:image/{ext};base64,{encoded}" width="{max_width}" alt="Classification Plot" />'
        except Exception as e:
            return f"Error reading chart file: {e}"

    # 3. Handle raw in-memory SVG string
    if "<svg" in plot_data:
        svg_clean = re.sub(r"<\?xml.*?\?>", "", plot_data, flags=re.DOTALL).strip()
        svg_single_line = re.sub(r"\s+", " ", svg_clean)
        return f'<div style="max-width: {max_width}px; max-height: 150px; display: inline-block;">{svg_single_line}</div>'

    # 4. Handle raw Base64 data URI string
    if plot_data.startswith("data:image"):
        return f'<img src="{plot_data}" width="{max_width}" alt="Classification Plot" />'

    return f"Unsupported payload: {plot_data[:40]}..."


def main():
    print("Fetching iTunes Top 10 tracks...", file=sys.stderr)
    songs = get_itunes_top10()

    if not songs:
        print("Error: Could not retrieve iTunes top tracks.", file=sys.stderr)
        sys.exit(1)

    print("Connecting to Hugging Face Space (dkappe/AISong)...", file=sys.stderr)
    # Mute stdout while connecting to the Space
    with contextlib.redirect_stdout(io.StringIO()):
    	client = Client("dkappe/AISong")
    
    date_str = datetime.now().strftime("%Y-%m-%d")

    results = []
    for song in songs:
        rank = song["rank"]
        title = song["title"]
        artist = song["artist"]
        genre = song["genre"]
        preview_url = song["preview_url"]

        if not preview_url:
            print(f"[{rank}/10] '{title}' - No audio preview available.", file=sys.stderr)
            results.append((rank, title, artist, genre, "No Preview Available"))
            continue

        print(f"[{rank}/10] Analyzing '{title}' by {artist}...", file=sys.stderr)
        with tempfile.NamedTemporaryFile(suffix=".m4a", delete=False) as tmp:
            tmp_path = tmp.name

        try:
            urllib.request.urlretrieve(preview_url, tmp_path)
            res = client.predict(handle_file(tmp_path), api_name="/predict")

            chart_html = extract_embedded_chart(res, max_width=220)
            results.append((rank, title, artist, genre, chart_html))
        except Exception as e:
            print(f"Error evaluating '{title}': {e}", file=sys.stderr)
            results.append((rank, title, artist, genre, f"Error: {e}"))
        finally:
            if os.path.exists(tmp_path):
                os.remove(tmp_path)

    # Output Markdown
    print("# iTunes Weekly Top 10 Audio Classification Report")
    print(f"**Date:** {date_str}\n")
    print("Classifier Model: [`dkappe/AISong`](https://huggingface.co/spaces/dkappe/AISong)\n")
    print("| Rank | Song Name | Artist | Genre | Classification Plot |")
    print("| :---: | :--- | :--- | :--- | :---: |")
    for r, t, a, g, p in results:
        clean_title = t.replace("|", "/")
        clean_artist = a.replace("|", "/")
        print(f"| {r} | {clean_title} | {clean_artist} | {g} | {p} |")


if __name__ == "__main__":
    main()