# app.py — SeparateTracks Gradio application
# Entry point: python app.py  (runs on http://localhost:7860)
# MCP endpoint: http://localhost:7860/gradio_api/mcp/sse
import shutil
import sys
from importlib import import_module
from pathlib import Path
from urllib.parse import parse_qs, urlparse

import gradio as gr

from modules.yt_audio_get_tracks import (
    download_audio,
    get_title,
    sanitize_job_id,
    separate_tracks,
)
from modules.file_utils import make_gradio_file_url


audio_gallery_module = import_module("modules.AudioGallery")
audio_gallery_head = f"<script>{getattr(audio_gallery_module, 'GALLERY_JS', '')}</script>"


SEPARATED_DIR = Path("separated").resolve()
gr.set_static_paths(paths=["separated/", SEPARATED_DIR.as_posix()])


def _extract_video_id(video_input: str) -> str:
    candidate = video_input.strip()
    if not candidate:
        return ""

    is_raw_id = "://" not in candidate and all(
        marker not in candidate for marker in ("/", "?", "&")
    )
    if "://" not in candidate and (
        "youtube.com" in candidate or "youtu.be" in candidate
    ):
        candidate = f"https://{candidate}"

    if is_raw_id:
        return candidate

    parsed = urlparse(candidate)
    host = parsed.netloc.lower()

    if host.endswith("youtu.be"):
        return parsed.path.strip("/").split("/")[0]

    if "youtube.com" in host:
        video_id = parse_qs(parsed.query).get("v", [""])[0]
        if video_id:
            return video_id

        path_parts = [part for part in parsed.path.split("/") if part]
        for prefix in ("shorts", "embed", "live", "v"):
            if prefix in path_parts:
                prefix_index = path_parts.index(prefix)
                if prefix_index + 1 < len(path_parts):
                    return path_parts[prefix_index + 1]

    return ""


def _build_audio_gallery(paths) -> str:
    audio_urls = [make_gradio_file_url(path) for path in paths]
    return audio_gallery_module.AudioGallery._build_html(
        audio_urls=audio_urls,
        labels=audio_gallery_module.AudioGallery.DEFAULT_LABELS,
        columns=3,
    )


def _prepare_uploaded_audio(uploaded_audio: str) -> tuple[str, str]:
    source_path = Path(uploaded_audio)
    suffix = source_path.suffix.lower()
    if suffix not in {".wav", ".mp3"}:
        raise ValueError("Please upload a .wav or .mp3 file.")

    job_id = sanitize_job_id(source_path.stem)
    target_path = SEPARATED_DIR / f"{job_id}{suffix}"
    shutil.copy2(source_path, target_path)
    return str(target_path), job_id


def _resolve_youtube_job_id(video_input: str) -> tuple[str, str, str]:
    video_id = _extract_video_id(video_input)
    if not video_id:
        return "", "", ""

    url = f"https://www.youtube.com/watch?v={video_id}"
    youtube_title = get_title(url)
    return video_id, youtube_title, sanitize_job_id(youtube_title or video_id)
# ---------------------------------------------------------------------------
# AudioGallery CSS — injected inline so the component is self-contained
# ---------------------------------------------------------------------------
_CSS = """
#versions {
    margin-top: 1em;
    width: 100%;
    text-align: center;
}
"""

# ---------------------------------------------------------------------------
# Version footer (graceful fallback if torch/cuda not available)
# ---------------------------------------------------------------------------
def _footer_html():
    try:
        from modules.version_info import versions_html
        return versions_html()
    except Exception:
        python_ver = ".".join(str(x) for x in sys.version_info[:3])
        return f"python: {python_ver} &bull; gradio: {gr.__version__}"


# ---------------------------------------------------------------------------
# Core processing function (also exposed as MCP tool)
# ---------------------------------------------------------------------------
def _process_video_impl(video_id: str, progress=None):
    progress_messages = []

    def on_progress(message):
        progress_messages.append(message)

    video_id, youtube_title, job_id = _resolve_youtube_job_id(video_id)
    if not video_id:
        return (
            "<p style='color:red;'>Please enter a YouTube video ID or URL.</p>",
            "No video ID provided.",
        )

    try:
        on_progress(f"YouTube title: {youtube_title}")
        if progress is not None:
            progress(0.0, desc="Preparing request")
        url = f"https://www.youtube.com/watch?v={video_id}"
        if progress is not None:
            progress(0.15, desc="Downloading audio")
        wav = download_audio(url, job_id, progress_callback=on_progress)
        if progress is not None:
            progress(0.45, desc="Separating tracks")
        drums, vocals, guitar, bass, other, piano, music = separate_tracks(
            wav,
            job_id,
            progress_callback=on_progress,
        )
        if progress is not None:
            progress(0.9, desc="Building audio gallery")
    except Exception as exc:
        status = "\n".join(progress_messages) if progress_messages else "Starting..."
        return f"<p style='color:red;'>Error: {exc}</p>", f"{status}\nError: {exc}"

    paths = [drums, vocals, guitar, bass, other, piano, music]
    status = "\n".join(progress_messages + ["Done."])
    if progress is not None:
        progress(1.0, desc="Done")
    return (
        _build_audio_gallery(paths),
        status,
    )


def process_video(video_id: str, progress=gr.Progress(track_tqdm=True)) -> str:
    """Download audio from a YouTube video and separate it into instrument stems.

    Uses Demucs htdemucs_6s to produce drums, vocals, guitar, bass, piano,
    other, and a combined music track. Results are displayed as an audio gallery.

    Args:
        video_id: YouTube video ID or URL (e.g. dQw4w9WgXcQ).

    Returns:
        HTML string containing the AudioGallery with all separated stems.
    """
    video_id, _youtube_title, job_id = _resolve_youtube_job_id(video_id)
    if not video_id:
        return "<p style='color:red;'>Please enter a YouTube video ID or URL.</p>"

    try:
        url = f"https://www.youtube.com/watch?v={video_id}"
        wav = download_audio(url, job_id)
        drums, vocals, guitar, bass, other, piano, music = separate_tracks(wav, job_id)
    except Exception as exc:
        return f"<p style='color:red;'>Error: {exc}</p>"

    paths = [drums, vocals, guitar, bass, other, piano, music]
    return _build_audio_gallery(paths)


def process_video_with_progress(
    video_id: str,
    uploaded_audio: str | None,
    cookies_upload: str | None,
    progress=gr.Progress(track_tqdm=True),
):
    status_lines = []

    def on_progress(message):
        status_lines.append(message)

    try:
        if cookies_upload is not None:            
            shutil.copy(cookies_upload, "modules/cookies.txt")
        if uploaded_audio:
            progress(0.05, desc="Preparing uploaded audio")
            yield "", "Preparing uploaded audio..."
            audio_path, job_id = _prepare_uploaded_audio(uploaded_audio)
            status_lines.append("Using uploaded audio file.")
        else:
            video_id, youtube_title, job_id = _resolve_youtube_job_id(video_id)
            if not video_id:
                yield (
                    "<p style='color:red;'>Please enter a YouTube video ID or URL, or upload an audio file.</p>",
                    "No video ID, URL, or audio file provided.",
                )
                return

            status_lines.append(f"YouTube title: {youtube_title}")
            url = f"https://www.youtube.com/watch?v={video_id}"
            progress(0.05, desc="Downloading audio")
            yield "", "\n".join(status_lines + ["Downloading audio from YouTube..."])
            audio_path = download_audio(url, job_id, progress_callback=on_progress)

        progress(0.4, desc="Separating tracks")
        yield "", "\n".join(status_lines)

        drums, vocals, guitar, bass, other, piano, music = separate_tracks(
            audio_path, job_id, progress_callback=on_progress
        )
        progress(0.9, desc="Building gallery")
        yield "", "\n".join(status_lines)

    except Exception as exc:
        yield (
            f"<p style='color:red;'>Error: {exc}</p>",
            "\n".join(status_lines) + f"\nError: {exc}",
        )
        return

    paths = [drums, vocals, guitar, bass, other, piano, music]
    status_lines.append("Done.")
    progress(1.0, desc="Done")
    yield (
        _build_audio_gallery(paths),
        "\n".join(status_lines),
    )


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------
with gr.Blocks(title="SeparateTracks") as demo:
    gr.Markdown(
        "## \U0001f3bc SeparateTracks\n"
        "Enter a YouTube video URL or ID, or upload a WAV/MP3 file, to "
        "separate the audio into instrument stems "
        "using [Demucs htdemucs\\_6s](https://github.com/adefossez/demucs)."
    )


    with gr.Row():
        video_id_input = gr.Textbox(
            label="YouTube Video ID or URL",
            placeholder="dQw4w9WgXcQ or https://www.youtube.com/watch?v=dQw4w9WgXcQ",
            scale=4,
        )
        run_btn = gr.Button("Separate Tracks", variant="primary", scale=1)
    with gr.Row():
        cookies_upload = gr.File(
            label="Upload Chrome cookies.txt (Netscape format)",
            file_types=[".txt"],
            type="filepath",
        )
        gr.Markdown("""
            **How to get cookies.txt:**
            1. Install [Get cookies.txt LOCALLY](https://chromewebstore.google.com/detail/get-cookiestxt-locally/cclelndahbckbenkjhflpdbgdldlbecc)
            2. Log into YouTube in Chrome
            3. Click extension → Export cookies
            4. Upload the file here
            """)
        gr.Markdown(
            "\n\nFor details about yt-dlp's extractor behavior, see: https://github.com/yt-dlp/yt-dlp-wiki/blob/master/Extractors.md\n\n"
            "One way to provide cookies safely is through a private browsing/incognito window:\n\n"
            "1. Open a new private browsing/incognito window and log into YouTube.\n"
            "2. chrome://extensions and specify Allow in Incognito.\n"
            "2. In the same window and same tab from step 1, navigate to https://www.youtube.com/robots.txt (this should be the only private/incognito browsing tab open).\n"
            "3. Export youtube.com cookies from the browser, then close the private browsing/incognito window so that the session is never opened in the browser again.\n"
        )
        # https://github.com/yt-dlp/yt-dlp-wiki/
    upload_input = gr.File(
        label="Audio File Override (.wav or .mp3)",
        file_types=[".wav", ".mp3"],
        type="filepath",
    )

    progress_output = gr.Textbox(label="Progress", interactive=False, lines=6)
    audio_output = gr.HTML(label="Separated Tracks")
    gr.HTML(value=_footer_html(), elem_id="versions", elem_classes="version-info")    

    run_btn.click(
        fn=process_video_with_progress,
        inputs=[video_id_input, upload_input, cookies_upload],
        outputs=[audio_output, progress_output],
    )

if __name__ == "__main__":
    demo.launch(
        mcp_server=True,
        server_name="0.0.0.0",
        server_port=7860,
        allowed_paths=[SEPARATED_DIR.as_posix(), "separated/", ".separated/"],
        favicon_path="separated/favicon.ico",
        css=_CSS,
        head=audio_gallery_head,
    )