#!/usr/bin/env python3
"""Download one completed song from the public Music API (Python 3 standard library).

Example:
  python3 examples/download_song.py --query "song title artist" --quality hires
  python3 examples/download_song.py --query "song title" --id TRACK_ID --quality hires

The API's /audio response determines the saved file type, not task.format or
byte_size. In particular, a cached ALAC task may serve a lossless FLAC file.
Resolve returned audio URLs against music.op.wiki without rewriting their host.
Same-origin deployments use that host for API and media; existing
music.zr.wiki media URLs remain accepted for compatibility. Custom deployments
can use --media-base-url. All requests are anonymous.
"""

from __future__ import annotations

import argparse
from http.client import HTTPException
import json
import math
import os
from pathlib import Path
import re
import sys
import tempfile
import time
from urllib.error import HTTPError, URLError
from urllib.parse import urljoin, urlsplit
from urllib.request import HTTPRedirectHandler, Request, build_opener


class DownloadError(Exception):
    """An actionable API, network, or output-file error."""


def origin(url: str) -> tuple[str, str, int]:
    try:
        parsed = urlsplit(url)
        if parsed.scheme not in ("http", "https") or not parsed.hostname:
            raise ValueError("expected an HTTP(S) URL")
        if parsed.username or parsed.password:
            raise ValueError("credentials in URLs are not supported")
        return parsed.scheme.lower(), parsed.hostname.lower(), parsed.port or (443 if parsed.scheme == "https" else 80)
    except ValueError as exc:
        raise DownloadError(f"Invalid API URL: {exc}") from exc


USER_AGENT = "MusicAPI-Downloader/1.1 (+https://music.op.wiki/docs/api/)"
PRODUCTION_API = "https://music.op.wiki"
LEGACY_PRODUCTION_MEDIA = "https://music.zr.wiki"


class TrustedOriginRedirects(HTTPRedirectHandler):
    def __init__(self, allowed_origins):
        super().__init__()
        self.allowed_origins = allowed_origins

    def redirect_request(self, request, fp, code, msg, headers, newurl):
        if origin(newurl) not in self.allowed_origins:
            raise DownloadError("Refusing a redirect to another origin")
        return super().redirect_request(request, fp, code, msg, headers, newurl)


class Client:
    def __init__(self, base_url: str, media_base_url: str | None = None):
        self.base_url = base_url.rstrip("/") + "/"
        self.allowed_origin = origin(self.base_url)
        self.media_origins = {self.allowed_origin}
        # This allowlist accepts old returned links; it never selects or rewrites
        # the media host. Relative resources resolve against the API base URL.
        if self.allowed_origin == origin(PRODUCTION_API):
            self.media_origins.add(origin(LEGACY_PRODUCTION_MEDIA))
        if media_base_url is not None:
            self.media_origins.add(origin(media_base_url))
        # Keep control API redirects separate from anonymous media redirects.
        # No cookie jar, browser session, or authorization headers are used.
        self.opener = build_opener(TrustedOriginRedirects({self.allowed_origin}))
        self.media_opener = build_opener(TrustedOriginRedirects(self.media_origins))

    def url(self, path: str, *, media: bool = False) -> str:
        if not isinstance(path, str) or not path.strip():
            raise DownloadError("API response is missing a valid resource URL")
        result = urljoin(self.base_url, path)
        allowed = self.media_origins if media else {self.allowed_origin}
        if origin(result) not in allowed:
            raise DownloadError("Refusing a resource URL on another origin")
        return result

    def open(self, path: str, *, payload=None, timeout: float = 30, media: bool = False):
        if media and payload is not None:
            raise DownloadError("Media requests must be anonymous GET requests")
        body = None if payload is None else json.dumps(payload).encode("utf-8")
        headers = {"Accept": "application/json" if payload is not None else "*/*",
                   "User-Agent": USER_AGENT, "Accept-Encoding": "identity"}
        if body is not None:
            headers["Content-Type"] = "application/json"
        request = Request(self.url(path, media=media), data=body, headers=headers)
        try:
            opener = self.media_opener if media else self.opener
            return opener.open(request, timeout=timeout)
        except HTTPError as exc:
            try:
                detail = exc.read(8192).decode("utf-8", errors="replace")
                try:
                    parsed = json.loads(detail)
                    if isinstance(parsed, dict):
                        detail = str(parsed.get("detail") or parsed.get("error") or parsed.get("message") or detail)
                except ValueError:
                    pass
                detail = re.sub(r"\s+", " ", detail).strip()[:500]
            finally:
                exc.close()
            raise DownloadError(f"HTTP {exc.code}: {detail or exc.reason}") from exc
        except (URLError, HTTPException, TimeoutError, OSError) as exc:
            raise DownloadError(f"Network request failed: {exc}") from exc

    def json(self, path: str, *, payload=None, timeout: float = 30) -> dict:
        try:
            with self.open(path, payload=payload, timeout=timeout) as response:
                raw = response.read(1024 * 1024 + 1)
        except (HTTPException, OSError) as exc:
            raise DownloadError(f"Could not read API response: {exc}") from exc
        if len(raw) > 1024 * 1024:
            raise DownloadError("API JSON response is unexpectedly large")
        try:
            result = json.loads(raw)
        except (ValueError, UnicodeError) as exc:
            raise DownloadError("API returned invalid JSON") from exc
        if not isinstance(result, dict):
            raise DownloadError("API returned an unexpected JSON response")
        return result


EXTENSIONS = {
    "audio/flac": ".flac", "audio/x-flac": ".flac",
    "audio/mp4": ".m4a", "audio/x-m4a": ".m4a", "video/mp4": ".m4a",
    "audio/mpeg": ".mp3", "audio/aac": ".aac",
    "audio/wav": ".wav", "audio/x-wav": ".wav", "audio/ogg": ".ogg",
    "audio/webm": ".webm", "application/octet-stream": ".bin",
}


def response_filename(headers) -> str:
    # HTTPMessage handles filename*=UTF-8''... as well as ordinary filename=.
    filename = headers.get_filename()
    if filename:
        filename = re.split(r"[/\\]", filename)[-1]
        filename = re.sub(r'[\x00-\x1f\x7f<>:"|?*]', "_", filename).strip(" .")
        if filename:
            return filename[:240]
    return "song" + EXTENSIONS.get(headers.get_content_type().lower(), ".bin")


def check_output(path: Path) -> None:
    if os.path.lexists(path):
        raise DownloadError(f"Output already exists (not overwritten): {path}")
    if not path.parent.is_dir():
        raise DownloadError(f"Output directory does not exist: {path.parent}")


def save_audio(client: Client, audio_url: str, output: Path | None, *, read_timeout: float = 120) -> Path:
    temporary = None
    try:
        with client.open(audio_url, timeout=read_timeout, media=True) as response:
            if response.status != 200:
                raise DownloadError(f"Expected a complete audio response (HTTP 200), received HTTP {response.status}")
            encoding = response.headers.get("Content-Encoding", "identity").strip().lower()
            if encoding not in ("", "identity"):
                raise DownloadError(f"Audio response unexpectedly uses Content-Encoding: {encoding}")
            content_type = response.headers.get_content_type().lower()
            if not (content_type.startswith("audio/") or content_type in ("video/mp4", "application/octet-stream")):
                raise DownloadError(f"Audio endpoint returned unexpected Content-Type: {content_type}")
            target = output if output is not None else Path(response_filename(response.headers))
            check_output(target)
            length = response.headers.get("Content-Length")
            try:
                expected = None if length is None else int(length)
                if expected is not None and expected < 0:
                    raise ValueError()
            except ValueError as exc:
                raise DownloadError("Audio response has an invalid Content-Length") from exc
            count = 0
            with tempfile.NamedTemporaryFile(dir=target.parent, prefix=".music-", suffix=".part", delete=False) as handle:
                temporary = Path(handle.name)
                while True:
                    chunk = response.read(1024 * 1024)
                    if not chunk:
                        break
                    handle.write(chunk)
                    count += len(chunk)
                handle.flush()
                os.fsync(handle.fileno())
            if expected is not None and count != expected:
                raise DownloadError(f"Incomplete audio transfer: expected {expected} bytes, received {count}")
            if count == 0:
                raise DownloadError("Audio response was empty")
        # Same-directory hard linking publishes the completed file atomically,
        # and unlike replace/rename it cannot overwrite an existing user file.
        os.link(temporary, target)
        return target
    except FileExistsError as exc:
        raise DownloadError("Output was created by another process; it was not overwritten") from exc
    except (OSError, URLError, HTTPException) as exc:
        raise DownloadError(f"Could not save audio: {exc}") from exc
    finally:
        if temporary is not None:
            temporary.unlink(missing_ok=True)


def download_song(*, query: str, platform: str = "applemusic", selected_id: str | None = None,
                  quality: str = "high", output: Path | None = None,
                  base_url: str = PRODUCTION_API, media_base_url: str | None = None,
                  timeout: float = 600,
                  read_timeout: float = 120,
                  poll_interval: float = 2) -> Path:
    if not query.strip():
        raise DownloadError("--query must not be empty")
    if not math.isfinite(timeout) or timeout <= 0:
        raise DownloadError("--timeout must be finite and greater than zero")
    if not math.isfinite(read_timeout) or read_timeout <= 0:
        raise DownloadError("--read-timeout must be finite and greater than zero")
    if output is not None:
        output = Path(output)
        check_output(output)  # Refuse an existing explicit output before creating any API work.
    client = Client(base_url, media_base_url)
    deadline = time.monotonic() + timeout

    def remaining() -> float:
        left = deadline - time.monotonic()
        if left <= 0:
            raise DownloadError("Timed out waiting for the song to finish preparing; the server task may continue")
        return min(30, left)

    payload = {"platform": platform, "query": query.strip(), "quality": quality, "include_cover": False}
    if selected_id:
        payload["selected_id"] = selected_id
    created = client.json("/api/v1/playbacks", payload=payload, timeout=remaining())
    task = created.get("playback")
    status_url = created.get("status_url")
    last_status = None
    while True:
        if not isinstance(task, dict):
            raise DownloadError("API response is missing the playback task")
        status = task.get("status")
        if status != last_status:
            print(f"Song status: {status}", file=sys.stderr)
            last_status = status
        if status == "completed":
            return save_audio(client, task.get("audio_url") or created.get("audio_url"), output,
                              read_timeout=read_timeout)
        if status == "failed":
            raise DownloadError(f"Song preparation failed: {task.get('error') or 'unknown server error'}")
        if status not in ("queued", "processing", "streaming"):
            raise DownloadError(f"Unexpected playback status: {status!r}")
        # streaming only means playback can start; a complete download must wait.
        delay = min(poll_interval, remaining())
        time.sleep(delay)
        task = client.json(status_url, timeout=remaining())


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--platform", default="applemusic", help="Music platform (default: applemusic)")
    parser.add_argument("--query", required=True, help="Song URL, title/search text, or ID (required)")
    parser.add_argument("--id", dest="selected_id", help="Exact track ID chosen from search results")
    parser.add_argument("--quality", choices=("standard", "high", "lossless", "hires"), default="high")
    parser.add_argument("--output", type=Path, help="Output file; existing files are never overwritten")
    parser.add_argument("--base-url", default="https://music.op.wiki")
    parser.add_argument("--media-base-url", help="Optional trusted media origin for a custom deployment; "
                        "same-origin needs no option, and production accepts legacy music.zr.wiki links")
    parser.add_argument("--timeout", type=float, default=600, help="Preparation/polling deadline only, in seconds (default: 600)")
    parser.add_argument("--read-timeout", type=float, default=120,
                        help="Audio request socket timeout in seconds, including first response/conversion wait (default: 120)")
    args = parser.parse_args(argv)
    try:
        saved = download_song(**vars(args))
    except (DownloadError, OSError, ValueError) as exc:
        print(f"Download failed: {exc}", file=sys.stderr)
        return 1
    except KeyboardInterrupt:
        print("Download cancelled", file=sys.stderr)
        return 130
    print(f"Saved: {saved.resolve()}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())