#!/usr/bin/env python3
from typing import Any, Literal
import requests
import tqdm
from mutagen.flac import FLAC, Picture
import sys
from dataclasses import dataclass
import datetime
from PIL import Image
from io import BytesIO

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


@dataclass
class Track:
    title: str
    artist: str
    nr: int
    id: str
    album_name: str
    release_date: datetime.date
    cover: Image.Image | None


@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

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

            image_data = Image.open(image_content)

        if "tracks" in data and "track_ids" in data:
            tracks = [
                Track(
                    nr=tnr + 1,
                    artist=data["artist"]["name"],
                    id=tid,
                    title=track["title"],
                    album_name=data["title"],
                    release_date=release_date,
                    cover=image_data,
                )
                for tnr, (tid, track) in enumerate(
                    zip(data["track_ids"], data["tracks"]["items"])
                )
            ]
        return cls(
            artist=data["artist"]["name"],
            title=data["title"],
            tracks=tracks,
            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,
        )


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) -> 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(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 = FLAC(filename)
            audio["title"] = track.title
            audio["artist"] = track.artist
            audio["tracknumber"] = str(track.nr)
            audio["album"] = track.album_name
            audio["date"] = track.release_date.isoformat()
            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) -> None:
        artist = album.artist
        album_title = album.title
        tracks = album.tracks if album.tracks is not None else []
        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)

    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)


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()