import datetime as dt
import hashlib
import re
import time
import typing
from collections.abc import Callable
from dataclasses import dataclass
from enum import IntEnum, auto
from pathlib import Path
from typing import Optional, Union

import mutagen
from mutagen.flac import FLAC, Picture
from mutagen.id3._frames import (
    APIC,
    TALB,
    TCON,
    TDRC,
    TIT2,
    TPE1,
    TPE2,
    TPOS,
    TRCK,
    USLT,
    WOAF,
)
from mutagen.id3._specs import ID3TimeStamp, PictureType
from mutagen.mp3 import MP3
from mutagen.mp4 import MP4, MP4Cover
from strenum import LowercaseStrEnum
from yandex_music import (
    Album,
    Client,
    Track,
    YandexMusicModel,
)
from yandex_music.exceptions import NetworkError

from ymd import api
from ymd.api import (
    ApiTrackQuality,
    Container,
    CustomDownloadInfo,
    get_download_info,
)
from ymd.mime_utils import MimeType, guess_mime_type

UNSAFE_PATH_CLEAR_RE = re.compile(r"[/\\]+")
SAFE_PATH_CLEAR_RE = re.compile(r"([^\w\-\'() ]|^\s+|\s+$)")

DEFAULT_PATH_PATTERN = Path("#album-artist", "#album", "#number-#title")
DEFAULT_COVER_RESOLUTION = 400

MIN_COMPATIBILITY_LEVEL = 0
MAX_COMPATIBILITY_LEVEL = 1

AUDIO_FILE_SUFFIXES = {".mp3", ".flac", ".m4a"}
TEMPORARY_FILE_NAME_TEMPLATE = ".yandex-music-downloader.{}.tmp"
MAX_FILE_NAME_LENGTH_WITHOUT_SUFFIX = 255 - max(
    len(suffix) for suffix in AUDIO_FILE_SUFFIXES
)


class CoreTrackQuality(IntEnum):
    LOW = 0
    NORMAL = auto()
    LOSSLESS = auto()


class LyricsFormat(LowercaseStrEnum):
    NONE = auto()
    TEXT = auto()
    LRC = auto()


CONTAINER_MUTAGEN_MAPPING: dict[Container, type[mutagen.FileType]] = {  # type: ignore
    Container.MP3: MP3,
    Container.FLAC: FLAC,
    Container.MP4: MP4,
}


@dataclass
class DownloadableTrack:
    download_info: CustomDownloadInfo
    path: Path
    track: Track


@dataclass
class AlbumCover:
    data: bytes
    mime_type: MimeType


def init_client(
    token: str, timeout: int, max_try_count: int, retry_delay: int
) -> Client:
    assert timeout > 0
    assert max_try_count >= 0
    assert retry_delay >= 0

    client = Client(token)
    client.request.set_timeout(timeout)

    original_wrapper = client.request._request_wrapper

    def retry_wrapper(*args, **kwargs):
        try_count = 0
        while True:
            try:
                return original_wrapper(*args, **kwargs)
            except NetworkError as error:
                if max_try_count == 0 or try_count < max_try_count:
                    try_count += 1
                    time.sleep(retry_delay)
                    continue
                raise error

    client.request._request_wrapper = retry_wrapper
    return client.init()


def full_title(obj: YandexMusicModel) -> str:
    """Return title with optional version; be tolerant to missing keys."""
    title = getattr(obj, "title", None)
    if title is None:
        try:
            title = obj["title"]
        except Exception:  # noqa: BLE001
            title = None
    if not title:
        return ""

    version = getattr(obj, "version", None)
    if version is None:
        try:
            version = obj["version"]
        except Exception:  # noqa: BLE001
            version = None
    if version:
        title = f"{title} ({version})"
    return title


def sanitize_path_component(text: str, unsafe_path: bool) -> str:
    """Return path-safe component."""
    if unsafe_path:
        return UNSAFE_PATH_CLEAR_RE.sub("_", text)
    text = text.strip()
    text = SAFE_PATH_CLEAR_RE.sub("_", text)
    text = re.sub(r"\s+", "_", text)
    text = re.sub(r"_+", "_", text)
    return text


def choose_first(*values: Optional[Union[str, int]]) -> str:
    for candidate in values:
        if candidate is None:
            continue
        text = str(candidate).strip()
        if text:
            return text
    return ""


def prepare_base_path(
    path_pattern: Path,
    track: Track,
    unsafe_path: bool = False,
    fallback_number: Optional[int] = None,
    fallback_total: Optional[int] = None,
    override_title: Optional[str] = None,
    override_album: Optional[str] = None,
    override_album_artist: Optional[str] = None,
    override_track_artist: Optional[str] = None,
    override_track_number: Optional[int] = None,
) -> Path:
    path_str = str(path_pattern)
    album = None
    album_artist = None
    track_artist = None
    track_position = None
    if albums := track.albums:
        album = albums[0]
        if artists := album.artists:
            album_artist = artists[0]
    # позиция трека корректнее брать из самого трека
    track_position = getattr(track, "track_position", None) or track_position
    if artists := track.artists:
        track_artist = artists[0]

    track_id = getattr(track, "id", None)
    track_title = choose_first(
        override_title,
        full_title(track),
        f"track_{track_id}" if track_id is not None else None,
        f"track_{fallback_number}" if fallback_number is not None else None,
        "track",
    )
    album_title = choose_first(
        override_album,
        full_title(album) if album else None,
        str(album.id) if album and album.id is not None else None,
        "album",
    )
    album_artist_name = choose_first(
        override_album_artist,
        album_artist.name if album_artist else None,
        "artist",
    )
    track_artist_name = choose_first(
        override_track_artist,
        track_artist.name if track_artist else None,
        album_artist_name,
        "artist",
    )
    pad_width = None
    if track_position and album and album.track_count:
        pad_width = max(2, len(str(album.track_count)))
    elif fallback_total:
        pad_width = max(2, len(str(fallback_total)))
    elif fallback_number:
        pad_width = max(2, len(str(fallback_number)))

    number = (
        override_track_number
        if override_track_number is not None
        else (track_position.index if track_position else fallback_number)
    )
    number_padded = (
        str(number).zfill(pad_width) if number is not None and pad_width is not None else number
    )

    number_value = number_padded if number_padded is not None else number
    repl_dict: dict[str, Union[str, int, None]] = {
        "#number-padded": number_padded,
        "#album-artist": album_artist_name,
        "#track-artist": track_artist_name,
        "#artist-id": track_artist.id if track_artist else None,
        "#album-id": album.id if album else None,
        "#track-id": track.id,
        "#number": number_value,
        "#title": track_title,
        "#album": album_title,
        "#year": album.year if album else None,
    }
    for placeholder, replacement in repl_dict.items():
        if replacement is None:
            replacement = ""
        replacement = sanitize_path_component(str(replacement), unsafe_path)
        if not replacement:
            replacement = sanitize_path_component("untitled", unsafe_path)
        path_str = path_str.replace(placeholder, replacement)
    path = Path(path_str)
    if not unsafe_path:
        path = Path(
            *(
                re.sub(r"_+", "_", part.replace(" ", "_"))
                for part in path.parts
            )
        )
    trimmed_parts = [
        part
        if len(part) <= MAX_FILE_NAME_LENGTH_WITHOUT_SUFFIX
        else part[:MAX_FILE_NAME_LENGTH_WITHOUT_SUFFIX]
        for part in path.parts
    ]
    return Path(*trimmed_parts)


def set_tags(
    path: Path,
    track: Track,
    container: Container,
    lyrics: Optional[str],
    album_cover: Optional[AlbumCover],
    compatibility_level: int,
    source_tags: Optional[dict[str, Optional[Union[str, int]]]] = None,
) -> None:
    file_type = CONTAINER_MUTAGEN_MAPPING.get(container)
    if file_type is None:
        raise ValueError(f"Unknown container: {container}")

    source_tags = source_tags or {}

    def as_list(value: Optional[Union[str, list[str], tuple[str, ...]]]) -> list[str]:
        if value is None:
            return []
        if isinstance(value, (list, tuple, set)):
            return [str(v).strip() for v in value if str(v).strip()]
        text = str(value).strip()
        return [text] if text else []

    tag = file_type(path)
    album = track.albums[0] if track.albums else Album()
    album_title = choose_first(
        source_tags.get("album"),
        full_title(album),
        str(album.id) if getattr(album, "id", None) is not None else None,
    )
    track_title = choose_first(
        source_tags.get("title"),
        full_title(track),
        f"track_{track.id}" if getattr(track, "id", None) is not None else None,
        "track",
    )
    track_artists = [a.name for a in track.artists if a.name]
    if not track_artists:
        track_artists = as_list(source_tags.get("track_artist"))
    album_artists = [a.name for a in album.artists if a.name]
    if not album_artists:
        album_artists = as_list(source_tags.get("album_artist"))
    if not track_artists:
        track_artists = album_artists
    genre = None
    if album.genre:
        genre = album.genre
    track_number = None
    disc_number = None
    position = getattr(track, "track_position", None) or album.track_position
    if position:
        track_number = position.index
        disc_number = position.volume
    if track_number is None and (tag_track := source_tags.get("track_number")):
        try:
            track_number = int(tag_track)
        except (TypeError, ValueError):
            track_number = None
    iso8601_release_date = None
    release_year: Optional[str] = None
    if album.release_date is not None:
        iso8601_release_date = dt.datetime.fromisoformat(album.release_date).astimezone(
            dt.timezone.utc
        )
        release_year = str(iso8601_release_date.year)
        iso8601_release_date = iso8601_release_date.strftime("%Y-%m-%d %H:%M:%S")
    if year := album.year:
        release_year = str(year)
    track_url = f"https://music.yandex.ru/album/{album.id}/track/{track.id}"

    if isinstance(tag, MP3):
        tag["TIT2"] = TIT2(encoding=3, text=track_title)
        tag["TALB"] = TALB(encoding=3, text=album_title)
        tag["TPE1"] = TPE1(encoding=3, text=track_artists)
        tag["TPE2"] = TPE2(encoding=3, text=album_artists)

        if tdrc_text := iso8601_release_date or release_year:
            tag["TDRC"] = TDRC(encoding=3, text=[ID3TimeStamp(tdrc_text)])
        if track_number:
            tag["TRCK"] = TRCK(encoding=3, text=str(track_number))
        if disc_number:
            tag["TPOS"] = TPOS(encoding=3, text=str(disc_number))
        if genre:
            tag["TCON"] = TCON(encoding=3, text=genre)

        if lyrics:
            tag["USLT"] = USLT(encoding=3, text=lyrics)
        if album_cover:
            tag["APIC"] = APIC(
                encoding=3,
                mime=album_cover.mime_type.value,
                type=3,
                data=album_cover.data,
            )

        tag["WOAF"] = WOAF(
            encoding=3,
            text=track_url,
        )
    elif isinstance(tag, MP4):
        tag["\xa9nam"] = track_title
        tag["\xa9alb"] = album_title
        artists_value = track_artists
        album_artists_value = album_artists
        if compatibility_level == 1:
            artists_value = "; ".join(track_artists)
            album_artists_value = "; ".join(album_artists)
        tag["\xa9ART"] = artists_value
        tag["aART"] = album_artists_value

        if iso8601_release_date is not None:
            tag["\xa9day"] = iso8601_release_date
        elif release_year is not None:
            tag["\xa9day"] = release_year
        if track_number:
            tag["trkn"] = [(track_number, 0)]
        if disc_number:
            tag["disk"] = [(disc_number, 0)]
        if genre:
            tag["\xa9gen"] = genre

        if lyrics:
            tag["\xa9lyr"] = lyrics
        if album_cover:
            mime_mp4_dict = {
                MimeType.JPEG: MP4Cover.FORMAT_JPEG,
                MimeType.PNG: MP4Cover.FORMAT_PNG,
            }
            mp4_image_format = mime_mp4_dict.get(album_cover.mime_type)
            if mp4_image_format is None:
                raise RuntimeError("Unsupported cover type")
            tag["covr"] = [MP4Cover(album_cover.data, imageformat=mp4_image_format)]
        tag["\xa9cmt"] = track_url
    elif isinstance(tag, FLAC):
        tag["title"] = track_title
        tag["album"] = album_title
        tag["artist"] = track_artists
        tag["albumartist"] = album_artists

        if date_text := iso8601_release_date or release_year:
            tag["date"] = date_text
        if track_number:
            tag["tracknumber"] = str(track_number)
        if disc_number:
            tag["discnumber"] = str(disc_number)
        if genre:
            tag["genre"] = genre

        if lyrics:
            tag["lyrics"] = lyrics
        if album_cover is not None:
            pic = Picture()
            pic.type = PictureType.COVER_FRONT
            pic.data = album_cover.data
            pic.mime = album_cover.mime_type.value
            tag.add_picture(pic)
        tag["comment"] = track_url
    else:
        raise RuntimeError("Unknown file format")

    tag.save()


def extract_basic_tags(path: Path, container: Container) -> dict[str, Optional[Union[str, int]]]:
    """Read minimal tags to be used for renaming."""
    title = album = album_artist = track_artist = None
    track_number: Optional[int] = None
    tag = CONTAINER_MUTAGEN_MAPPING[container](path)

    def first_text(value) -> Optional[str]:
        if value is None:
            return None
        if isinstance(value, list):
            if not value:
                return None
            value = value[0]
        if hasattr(value, "text"):
            value = value.text
            if isinstance(value, list):
                value = value[0] if value else None
        if isinstance(value, bytes):
            try:
                value = value.decode("utf-8", "ignore")
            except Exception:
                return None
        return str(value) if value else None

    if isinstance(tag, MP3):
        title = first_text(tag.get("TIT2"))
        album = first_text(tag.get("TALB"))
        track_artist = first_text(tag.get("TPE1"))
        album_artist = first_text(tag.get("TPE2"))
        if trck := first_text(tag.get("TRCK")):
            try:
                track_number = int(str(trck).split("/")[0])
            except ValueError:
                track_number = None
    elif isinstance(tag, MP4):
        title = first_text(tag.tags.get("\xa9nam") if tag.tags else None)
        album = first_text(tag.tags.get("\xa9alb") if tag.tags else None)
        track_artist = first_text(tag.tags.get("\xa9ART") if tag.tags else None)
        album_artist = first_text(tag.tags.get("aART") if tag.tags else None)
        if trkn := tag.tags.get("trkn") if tag.tags else None:
            try:
                track_number = int(trkn[0][0])
            except Exception:
                track_number = None
    elif isinstance(tag, FLAC):
        def get_flac(key: str) -> Optional[str]:
            values = tag.get(key)
            if not values:
                return None
            return values[0]

        title = get_flac("title")
        album = get_flac("album")
        artists_val = get_flac("artist")
        if artists_val:
            track_artist = artists_val
        album_artist = get_flac("albumartist")
        if trck := get_flac("tracknumber"):
            try:
                track_number = int(str(trck).split("/")[0])
            except ValueError:
                track_number = None

    return {
        "title": title,
        "album": album,
        "album_artist": album_artist,
        "track_artist": track_artist,
        "track_number": track_number,
    }


def retarget_with_tags(
    path_pattern: Path,
    track: Track,
    container: Container,
    current_path: Path,
    unsafe_path: bool,
    fallback_number: Optional[int],
    fallback_total: Optional[int],
    override_tags: Optional[dict[str, Optional[Union[str, int]]]] = None,
) -> Path:
    tags: dict[str, Optional[Union[str, int]]] = {}
    if override_tags:
        tags.update(override_tags)
    try:
        extracted = extract_basic_tags(current_path, container)
    except Exception:
        extracted = {}
    for key, value in extracted.items():
        if key not in tags or not tags[key]:
            tags[key] = value
    new_base = prepare_base_path(
        path_pattern,
        track,
        unsafe_path=unsafe_path,
        fallback_number=fallback_number,
        fallback_total=fallback_total,
        override_title=tags.get("title"),
        override_album=tags.get("album"),
        override_album_artist=tags.get("album_artist"),
        override_track_artist=tags.get("track_artist"),
        override_track_number=tags.get("track_number"),
    )
    new_target = new_base.with_suffix(current_path.suffix)
    try:
        same = new_target.resolve() == current_path.resolve()
    except FileNotFoundError:
        same = False
    if same:
        return current_path
    new_target.parent.mkdir(parents=True, exist_ok=True)
    current_path.replace(new_target)
    return new_target


def download_track(
    track_info: DownloadableTrack,
    cover_resolution: int = DEFAULT_COVER_RESOLUTION,
    lyrics_format: LyricsFormat = LyricsFormat.NONE,
    embed_cover: bool = False,
    covers_cache: Optional[dict[int, AlbumCover]] = None,
    compatibility_level: int = 1,
    path_pattern: Optional[Path] = None,
    unsafe_path: bool = False,
    fallback_number: Optional[int] = None,
    fallback_total: Optional[int] = None,
    rename_using_tags: bool = False,
):
    if embed_cover and covers_cache is None:
        raise RuntimeError("covers_cache isn't provided")
    covers_cache = typing.cast(dict[int, AlbumCover], covers_cache)
    target_path = track_info.path
    track = track_info.track
    client = typing.cast(Client, track.client)
    assert client

    text_lyrics = None
    if lyrics_format != LyricsFormat.NONE and (lyrics_info := track.lyrics_info):
        if lyrics_format == LyricsFormat.LRC and lyrics_info.has_available_sync_lyrics:
            lrc_path = target_path.with_suffix(".lrc")
            if not lrc_path.is_file() and (
                track_lyrics := track.get_lyrics(format_="LRC")
            ):
                lyrics = track_lyrics.fetch_lyrics()
                write_via_temporary_file(lyrics.encode("utf-8"), lrc_path)
        elif lyrics_info.has_available_text_lyrics:
            if track_lyrics := track.get_lyrics(format_="TEXT"):
                text_lyrics = track_lyrics.fetch_lyrics()

    cover = None
    if track.cover_uri is not None:
        if cover_resolution == -1:
            cover_size = "orig"
        else:
            cover_size = f"{cover_resolution}x{cover_resolution}"
        cover_bytes = track.download_cover_bytes(size=cover_size)
        mime_type = guess_mime_type(cover_bytes)
        if mime_type is None:
            raise RuntimeError("Unknown cover mime type")
        album_cover = AlbumCover(data=cover_bytes, mime_type=mime_type)
        if embed_cover:
            album = track.albums[0] if track.albums else Album()
            if album.id and (cached_cover := covers_cache.get(album.id)):
                cover = cached_cover
            else:
                if album.id:
                    cover = covers_cache[album.id] = album_cover
        else:
            mime_suffix_dict = {MimeType.JPEG: ".jpg", MimeType.PNG: ".png"}
            file_suffix = mime_suffix_dict.get(album_cover.mime_type)
            if file_suffix is None:
                raise RuntimeError("Unknown mime type")
            cover_path = target_path.parent / ("cover" + file_suffix)
            if not cover_path.is_file():
                write_via_temporary_file(album_cover.data, cover_path)

    download_info = track_info.download_info
    track_data = api.download_track(client, download_info)

    source_tags: dict[str, Optional[Union[str, int]]] = {}

    def tag_hook(tmp_path: Path) -> None:
        nonlocal source_tags
        try:
            source_tags = extract_basic_tags(tmp_path, download_info.file_format.container)
        except Exception:
            source_tags = {}
        set_tags(
            tmp_path,
            track,
            download_info.file_format.container,
            text_lyrics,
            cover,
            compatibility_level,
            source_tags=source_tags,
        )

    write_via_temporary_file(
        track_data,
        target_path,
        temporary_file_hook=tag_hook,
    )

    if rename_using_tags and path_pattern is not None:
        try:
            renamed_path = retarget_with_tags(
                path_pattern=path_pattern,
                track=track,
                container=download_info.file_format.container,
                current_path=target_path,
                unsafe_path=unsafe_path,
                fallback_number=fallback_number,
                fallback_total=fallback_total,
                override_tags=source_tags if source_tags else None,
            )
            target_path = renamed_path
        except Exception:
            # Renaming is non-critical; ignore failures to keep download intact
            pass


def to_downloadable_track(
    track: Track, quality: CoreTrackQuality, base_path: Path
) -> DownloadableTrack:
    api_quality = ApiTrackQuality.NORMAL
    if quality == CoreTrackQuality.LOW:
        api_quality = ApiTrackQuality.LOW
    elif quality == CoreTrackQuality.NORMAL:
        api_quality = ApiTrackQuality.NORMAL
    elif quality == CoreTrackQuality.LOSSLESS:
        api_quality = ApiTrackQuality.LOSSLESS

    download_info = get_download_info(track, api_quality)
    container = download_info.file_format.container

    if container == Container.MP3:
        suffix = ".mp3"
    elif container == Container.MP4:
        suffix = ".m4a"
    elif container == Container.FLAC:
        suffix = ".flac"
    else:
        raise RuntimeError("Unknown codec")

    target_path = str(base_path) + suffix
    return DownloadableTrack(
        download_info=download_info,
        track=track,
        path=Path(target_path),
    )


def write_via_temporary_file(
    data: bytes,
    target_path: Path,
    temporary_file_hook: Optional[Callable[[Path], None]] = None,
) -> Path:
    target_name = hashlib.sha256(target_path.name.encode()).hexdigest()
    temporary_file = target_path.parent / (
        TEMPORARY_FILE_NAME_TEMPLATE.format(target_name)
    )
    try:
        temporary_file.write_bytes(data)
        if temporary_file_hook is not None:
            temporary_file_hook(temporary_file)
    except InterruptedError as e:
        temporary_file.unlink()
        raise e
    temporary_file.rename(target_path)
    return target_path