#!/usr/bin/env python3
import os
from typing import Any, Literal
import requests
import tqdm

import music_tag
import sys
from dataclasses import dataclass
import datetime

QOBUZ_API_BASE = "https://us.qobuz.squid.wtf/api"
QUALITY = 27


@dataclass
class Track:
    album: "Album"
    title: str
    nr: int
    id: str


@dataclass
class ImageLink:
    small: str | None
    thumbnail: str | None
    large: str | None
    back: str | None

    @classmethod
    def from_dict(
        cls,
        data: dict[
            Literal["small"]
            | Literal["thumbnail"]
            | Literal["large"]
            | Literal["back"],
            str | None,
        ],
    ) -> "ImageLink":
        return ImageLink(**data)


@dataclass
class Album:
    artist: str
    title: str
    tracks: list[Track] | None
    maximum_bit_depth: int
    maximum_sampling_rate: int
    id: str
    image: ImageLink
    release_date: datetime.date
    cover_data: bytes | None

    @classmethod
    def from_dict(cls, data: dict[Any, Any]) -> "Album":
        release_date = datetime.date.fromtimestamp(data["released_at"])
        image = ImageLink.from_dict(data["image"])
        image_content = None
        if image.large is not None:
            # Fetch image
            print(f"Fetching image from {image.large}")
            image_content = requests.get(image.large).content

        album = cls(
            artist=data["artist"]["name"],
            title=data["title"],
            tracks=None,
            maximum_bit_depth=data["maximum_bit_depth"],
            maximum_sampling_rate=data["maximum_sampling_rate"],
            id=data["id"],
            image=ImageLink.from_dict(data["image"]),
            release_date=release_date,
            cover_data=image_content,
        )

        if "tracks" in data and "track_ids" in data:
            tracks = [
                Track(
                    nr=tnr + 1,
                    id=tid,
                    title=track["title"],
                    album=album,
                )
                for tnr, (tid, track) in enumerate(
                    zip(data["track_ids"], data["tracks"]["items"])
                )
            ]
            album.tracks = tracks
        return album


class Qobuz:
    def __init__(self, api_base=QOBUZ_API_BASE, quality: int = QUALITY) -> None:
        self.api_base = api_base
        self.quality = quality

    def get_download_url(self, track_id: str) -> str:
        url = f"{QOBUZ_API_BASE}/download-music"
        result = requests.get(
            url, params={"track_id": track_id, "quality": self.quality}
        )
        if result.status_code == 200:
            result_json = result.json()
            if result_json["success"]:
                return result_json["data"]["url"]
        raise RuntimeError(f"Failed to get download URL, result: {result}")

    def download_track(self, track: Track, filename: str, dest: str) -> None:
        download_url = self.get_download_url(track.id)
        response = requests.get(download_url, stream=True)
        if response.status_code == 200:
            total_size = int(response.headers.get("content-length", 0))
            with tqdm.tqdm(
                total=total_size, unit="B", unit_scale=True, desc=filename
            ) as pbar:
                with open(os.path.join(dest, filename), "wb") as file:
                    for chunk in response.iter_content(chunk_size=8192):
                        file.write(chunk)
                        pbar.update(len(chunk))
            # Set metadata using mutagen
            audio = music_tag.load_file(filename)
            audio["tracktitle"] = track.title
            audio["artist"] = track.album.artist
            audio["tracknumber"] = str(track.nr)
            audio["album"] = track.album.title
            audio["year"] = track.album.release_date.year
            audio["artwork"] = track.album.cover_data
            audio.save()

            print(f"Downloaded {filename}")
        else:
            raise RuntimeError("Failed to download track")

    def get_album(self, album_id: str) -> Album:
        url = f"{self.api_base}/get-album"
        result = requests.get(url, params={"album_id": album_id})
        if result.status_code == 200:
            data = result.json()["data"]
            return Album.from_dict(data)
        raise RuntimeError(f"Failed to get album, result: {result}")

    def download_album(self, album: Album, dest: str) -> None:
        artist = album.artist
        album_title = album.title
        tracks = album.tracks if album.tracks is not None else []
        os.makedirs(dest, exist_ok=True)
        print(f"Downloading album: {artist} - {album_title}")
        for nr, track in enumerate(tracks):
            print(f"Downloading track #{nr + 1:02d} {track.title}...")
            filename = f"{artist} - {album_title} - {nr + 1:02d} {track.title}.flac"
            self.download_track(track, filename, dest)

    def search_album(self, query: str) -> Album:
        print(f'Searching for "{query}"')
        url = f"{self.api_base}/get-music"
        result = requests.get(url, params={"q": query, "offset": 0})
        if result.status_code == 200:
            first_album_id = result.json()["data"]["albums"]["items"][0]["id"]
            return self.get_album(first_album_id)
        raise RuntimeError("No results")

    def search_and_download(self, query: str) -> None:
        album = self.search_album(query)
        self.download_album(album, f"{album.artist} - {album.title}")


def main() -> None:
    # trackid = sys.argv[1]
    # filename = sys.argv[2]
    query = sys.argv[1]
    qb = Qobuz()
    print(qb.search_and_download(query))


if __name__ == "__main__":
    main()