from src.Library import *

import json
from datetime import datetime
from pathlib import Path
from typing import List, Optional, Any

def _dt_to_str(dt: Optional[datetime]) -> Optional[str]:
    return dt.isoformat() if dt else None


def _str_to_dt(s: Optional[str]) -> Optional[datetime]:
    return datetime.fromisoformat(s) if s else None


def _path_to_str(p: Optional[Path]) -> Optional[str]:
    return str(p) if p else None


def _str_to_path(s: Optional[str]) -> Optional[Path]:
    return Path(s) if s else None


class LibrarySaver:
    def __init__(self, library: Library):
        self.library = library

    def save(self, path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as f:
            json.dump(self.to_dict(), f, indent=2, ensure_ascii=False)

    def to_dict(self) -> dict[str, Any]:
        return {
            "tracks": [self._track_to_dict(t) for t in self.library.tracks],
            "playlists": [self._playlist_to_dict(p) for p in self.library.playlists],
            "tracklists": [self._tracklist_to_dict(tl) for tl in self.library.tracklists],
            "albums": [self._album_to_dict(a) for a in self.library.albums],
        }

    def _track_to_dict(self, t: Track) -> dict[str, Any]:
        return {
            "id": t.id,
            "title": t.title,
            "artists": t.artists,
            "album": t.album,
            "duration_ms": t.duration_ms,
            "added_at": _dt_to_str(t.added_at),
            "spotify": self._spotify_to_dict(t.spotify),
            "local": self._local_to_dict(t.local),
        }

    def _playlist_to_dict(self, p: Playlist) -> dict[str, Any]:
        return {
            "name": p.name,
            "tracks": [t.id for t in p.tracks],
            "created_at": _dt_to_str(p.created_at),
            "description": p.description,
            "spotify": self._spotify_to_dict(p.spotify),
        }

    def _tracklist_to_dict(self, tl: TrackList) -> dict[str, Any]:
        return {"tracks": [t.id for t in tl.tracks]}

    def _album_to_dict(self, a: Album) -> dict[str, Any]:
        return {
            "title": a.title,
            "release_date": a.release_date,
            "tracks": [t.id for t in a.tracks],
            "spotify": self._spotify_to_dict(a.spotify),
        }

    def _local_to_dict(self, local):
        if not local:
            return None
        return {"path": local.path}

    def _spotify_to_dict(self, s: Optional[SpotifyData]) -> Optional[dict[str, Any]]:
        if not s:
            return None
        return {
            "id": s.id,
            "uri": s.uri,
            "url": s.url,
            "href": s.href,
            "file_path": _path_to_str(s.file_path),
            "popularity": s.popularity,
            "is_playable": s.is_playable,
            "is_local": s.is_local,
            "external_ids": s.external_ids,
            "artists": [{"id": a.id, "name": a.name, "uri": a.uri, "url": a.url} for a in s.artists],
            "album": {
                "id": s.album.id,
                "name": s.album.name,
                "release_date": s.album.release_date,
                "total_tracks": s.album.total_tracks,
                "uri": s.album.uri,
                "url": s.album.url,
            } if s.album else None,
        }


class LibraryLoader:
    def __init__(self):
        # Map of track id -> Track object
        self.track_map: dict[str, Track] = {}

    def load(self, path: Path) -> Library:
        with path.open("r", encoding="utf-8") as f:
            data = json.load(f)

        # Load tracks first
        tracks = []
        for t in data.get("tracks", []):
            track = self._track_from_dict(t)
            tracks.append(track)
            self.track_map[track.id] = track

        # Load playlists using track_map
        playlists = []
        for p in data.get("playlists", []):
            pl_tracks = []
            for tid in p.get("tracks", []):
                if tid not in self.track_map:
                    raise KeyError(f"Track ID '{tid}' not found in track map")
                pl_tracks.append(self.track_map[tid])

            playlists.append(Playlist(
                name=p["name"],
                tracks=pl_tracks,
                created_at=_str_to_dt(p.get("created_at")),
                description=p.get("description"),
                spotify=self._spotify_from_dict(p.get("spotify")),
            ))


        # Load tracklists
        tracklists = []
        for tl in data.get("tracklists", []):
            tl_tracks = []
            for tid in tl.get("tracks", []):
                if tid not in self.track_map:
                    raise KeyError(f"Track ID '{tid}' not found in track map (tracklist)")
                tl_tracks.append(self.track_map[tid])
            tracklists.append(TrackList(tracks=tl_tracks))

        # Load albums
        albums = []
        for a in data.get("albums", []):
            alb_tracks = []
            for tid in a.get("tracks", []):
                if tid not in self.track_map:
                    raise KeyError(f"Track ID '{tid}' not found in track map (album '{a.get('title')}')")
                alb_tracks.append(self.track_map[tid])
            albums.append(Album(
                title=a["title"],
                release_date=a.get("release_date"),
                tracks=alb_tracks,
                spotify=self._spotify_from_dict(a.get("spotify")),
            ))

        return Library(playlists=playlists, tracks=tracks, tracklists=tracklists, albums=albums)

    def _track_from_dict(self, d: dict[str, Any]) -> Track:
        return Track(
            id=d["id"],
            title=d["title"],
            artists=d["artists"],
            album=d.get("album"),
            duration_ms=d.get("duration_ms"),
            added_at=_str_to_dt(d.get("added_at")),
            spotify=self._spotify_from_dict(d.get("spotify")),
            local=self._local_from_dict(d.get("local")),
        )

    def _local_from_dict(self, d):
        if not d:
            return None
        return LocalData(
            path=d["path"],
        )

    def _spotify_from_dict(self, d: Optional[dict[str, Any]]) -> Optional[SpotifyData]:
        if not d:
            return None
        return SpotifyData(
            id=d["id"],
            uri=d["uri"],
            url=d["url"],
            href=d["href"],
            file_path=_str_to_path(d.get("file_path")),
            popularity=d.get("popularity"),
            is_playable=d.get("is_playable"),
            is_local=d.get("is_local"),
            external_ids=d.get("external_ids"),
            artists=[SpotifyArtist(**a) for a in d.get("artists", [])],
            album=SpotifyAlbum(**d["album"]) if d.get("album") else None,
        )