import gradio as gr
import tempfile, os, subprocess, shutil
from typing import List, Tuple
from mutagen.id3 import ID3, ID3NoHeaderError, GEOB, TSSE
from mutagen.mp3 import MP3

TARGET_SR = 44100
TARGET_CH = 1
TARGET_BR = "128k"

def standardize_one(in_path: str, tsse_text: str, geob_bytes: bytes, geob_mime: str, geob_desc: str) -> str:
    base = os.path.splitext(os.path.basename(in_path))[0]
    out_path = os.path.join(tempfile.gettempdir(), f"{base}_standardized.mp3")

    # 1) Re-encode audio with FFmpeg and enforce container/tag options
    cmd = [
        "ffmpeg", "-y",
        "-i", in_path,
        "-ar", str(TARGET_SR),
        "-ac", str(TARGET_CH),
        "-c:a", "libmp3lame",
        "-b:a", TARGET_BR,
        "-write_xing", "0",
        "-id3v2_version", "4",
        "-write_id3v1", "0",
        "-metadata", f"TSSE={tsse_text or 'HF-Space'}",
        out_path
    ]
    subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    # 2) Ensure ID3v2.4 and add GEOB + TSSE with Mutagen
    try:
        tags = ID3(out_path)
    except ID3NoHeaderError:
        tags = ID3()

    # Enforce TSSE (software/encoder)
    tags.setall("TSSE", [TSSE(encoding=3, text=[tsse_text or "HF-Space"])])
    # Always add a GEOB frame (either from upload or a tiny placeholder)
    if geob_bytes is None or len(geob_bytes) == 0:
        geob_bytes = b"HF-Space GEOB placeholder"
        geob_mime = geob_mime or "application/octet-stream"
        geob_desc = geob_desc or "asset"
    tags.add(GEOB(encoding=3, mime=geob_mime or "application/octet-stream", desc=geob_desc or "asset", data=geob_bytes))

    # Save as ID3v2.4 explicitly
    tags.save(out_path, v2_version=4)

    # Sanity check via Mutagen
    _ = MP3(out_path)  # will raise if corrupt
    return out_path

def standardize_many(files: List[gr.File], tsse_text: str, geob_file: gr.File, geob_mime: str, geob_desc: str):
    outputs = []
    geob_bytes = None
    if geob_file is not None and hasattr(geob_file, "name") and geob_file.name:
        with open(geob_file.name, "rb") as f:
            geob_bytes = f.read()

    for f in files or []:
        out = standardize_one(f.name, tsse_text, geob_bytes, geob_mime, geob_desc)
        outputs.append(out)

    # Optionally zip if multiple
    if len(outputs) > 1:
        zip_path = os.path.join(tempfile.gettempdir(), "standardized_mp3s.zip")
        if os.path.exists(zip_path):
            os.remove(zip_path)
        with tempfile.TemporaryDirectory() as tmpd:
            paths = []
            for p in outputs:
                dst = os.path.join(tmpd, os.path.basename(p))
                shutil.copy2(p, dst)
                paths.append(dst)
            shutil.make_archive(zip_path[:-4], 'zip', tmpd)
        return zip_path
    elif len(outputs) == 1:
        return outputs[0]
    else:
        return None

with gr.Blocks(title="MP3 Standardizer (44.1k/mono/128k CBR, no Xing, ID3v2.4, TSSE+GEOB)") as demo:
    gr.Markdown("Upload MP3s to standardize them to 44.1 kHz, mono, 128 kbps CBR, no Xing/Info, ID3v2.4 only, plus TSSE and GEOB tags.")

    with gr.Row():
        files = gr.File(label="MP3 files", file_types=[".mp3"], file_count="multiple")
    with gr.Accordion("Metadata options", open=False):
        tsse = gr.Textbox(label="TSSE (encoder/software)", value="HF-Space", placeholder="e.g., Lavf60.16.100")
        geob = gr.File(label="GEOB payload (optional)", file_types=[".bin", ".txt", ".json", ".png", ".jpg", ".mp3", ".mp4"], file_count="single")
        geob_mime = gr.Textbox(label="GEOB MIME type", value="application/octet-stream")
        geob_desc = gr.Textbox(label="GEOB description", value="asset")

    run = gr.Button("Standardize")
    result = gr.File(label="Download standardized file(s)")
    run.click(standardize_many, inputs=[files, tsse, geob, geob_mime, geob_desc], outputs=[result])

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