lintunes/lintunes/smart.py

"""Smart playlists: criteria model, evaluation, and iTunes import.

A smart playlist's membership is *derived* from a `SmartCriteria` (a tree of
rules) rather than hand-curated. This module owns:

* the data model (`SmartRule` / `SmartGroup` / `SmartLimit` / `SmartCriteria`),
  recursive so nested rule groups round-trip through the playlist JSON;
* `FIELD_REGISTRY`, the single source of truth shared by the evaluator and the
  editor dialog (field -> Track attribute, value type, allowed operators);
* `evaluate()`, which walks the tree against the library and returns the
  ordered list of matching track ids; and
* `parse_itunes_smart()`, which converts the iTunes "Smart Info"/"Smart
  Criteria" binary blobs into a `SmartCriteria` (via the vendored parser in
  `lintunes.itunes_smart`).

The evaluator reads tracks purely by attribute name (`getattr`), so this module
imports nothing from `lintunes.models` -- keeping it free of import cycles with
`playlist.py`.
"""

from __future__ import annotations

import logging
import random as _random
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from enum import Enum
from typing import Optional

logger = logging.getLogger(__name__)


# --------------------------------------------------------------------------- #
# Field metadata
# --------------------------------------------------------------------------- #

class FieldType(Enum):
    STRING = "string"
    INT = "int"
    DURATION = "duration"   # milliseconds, like INT but rendered as time
    DATE = "date"
    BOOL = "bool"
    RATING = "rating"       # stored 0-100, edited as 0-5 stars


# Operator vocabularies per field type.
STRING_OPS = ("is", "is_not", "contains", "not_contains", "starts_with", "ends_with")
NUMERIC_OPS = ("is", "is_not", "greater", "less", "in_range")
DATE_OPS = ("after", "before", "in_range", "not_in_range", "in_last", "not_in_last")
BOOL_OPS = ("is",)


@dataclass(frozen=True)
class FieldMeta:
    key: str
    label: str
    track_attr: str
    type: FieldType
    operators: tuple

    @property
    def ops(self) -> tuple:
        return self.operators


def _ops_for(ftype: FieldType) -> tuple:
    if ftype is FieldType.STRING:
        return STRING_OPS
    if ftype in (FieldType.INT, FieldType.DURATION, FieldType.RATING):
        return NUMERIC_OPS
    if ftype is FieldType.DATE:
        return DATE_OPS
    return BOOL_OPS


def _meta(key, label, attr, ftype) -> FieldMeta:
    return FieldMeta(key, label, attr, ftype, _ops_for(ftype))


# The fields the editor offers and the evaluator understands. Only fields the
# LinTunes Track actually backs are present (the library is music-only).
FIELD_REGISTRY: dict[str, FieldMeta] = {m.key: m for m in [
    _meta("name", "Name", "name", FieldType.STRING),
    _meta("artist", "Artist", "artist", FieldType.STRING),
    _meta("album_artist", "Album Artist", "album_artist", FieldType.STRING),
    _meta("album", "Album", "album", FieldType.STRING),
    _meta("genre", "Genre", "genre", FieldType.STRING),
    _meta("composer", "Composer", "composer", FieldType.STRING),
    _meta("grouping", "Grouping", "grouping", FieldType.STRING),
    _meta("comments", "Comments", "comments", FieldType.STRING),
    _meta("kind", "Kind", "kind", FieldType.STRING),
    _meta("sort_name", "Sort Name", "sort_name", FieldType.STRING),
    _meta("sort_album", "Sort Album", "sort_album", FieldType.STRING),
    _meta("year", "Year", "year", FieldType.INT),
    _meta("bpm", "BPM", "bpm", FieldType.INT),
    _meta("bit_rate", "Bit Rate", "bit_rate", FieldType.INT),
    _meta("sample_rate", "Sample Rate", "sample_rate", FieldType.INT),
    _meta("size", "Size", "size", FieldType.INT),
    _meta("disc_number", "Disc Number", "disc_number", FieldType.INT),
    _meta("track_number", "Track Number", "track_number", FieldType.INT),
    _meta("play_count", "Plays", "play_count", FieldType.INT),
    _meta("skip_count", "Skips", "skip_count", FieldType.INT),
    _meta("rating", "Rating", "rating", FieldType.RATING),
    _meta("total_time", "Time", "total_time", FieldType.DURATION),
    _meta("compilation", "Compilation", "compilation", FieldType.BOOL),
    _meta("date_added", "Date Added", "date_added", FieldType.DATE),
    _meta("date_modified", "Date Modified", "date_modified", FieldType.DATE),
    _meta("last_played", "Last Played", "play_date_utc", FieldType.DATE),
    _meta("last_skipped", "Last Skipped", "skip_date", FieldType.DATE),
]}


# Limit "selected by" methods -> (Track attribute to sort on, reverse?).
# "random" is special-cased. Used by the evaluator to pick which N items.
SELECTION_SORT: dict[str, tuple] = {
    "random": (None, False),
    "name": ("sort_name", False),
    "album": ("sort_album", False),
    "artist": ("sort_artist", False),
    "genre": ("genre", False),
    "highest_rating": ("rating", True),
    "lowest_rating": ("rating", False),
    "most_recently_played": ("play_date_utc", True),
    "least_recently_played": ("play_date_utc", False),
    "most_often_played": ("play_count", True),
    "least_often_played": ("play_count", False),
    "most_recently_added": ("date_added", True),
    "least_recently_added": ("date_added", False),
}

_UNIT_SECONDS = {"days": 86400, "weeks": 604800, "months": 2628000}


# --------------------------------------------------------------------------- #
# Schema
# --------------------------------------------------------------------------- #

@dataclass
class SmartRule:
    field: str = "artist"
    operator: str = "contains"
    value: object = ""
    value2: object = None          # upper bound for in_range/not_in_range
    unit: Optional[str] = None     # "days"|"weeks"|"months" for in_last/not_in_last

    def to_dict(self) -> dict:
        d = {"kind": "rule", "field": self.field,
             "operator": self.operator, "value": self.value}
        if self.value2 is not None:
            d["value2"] = self.value2
        if self.unit is not None:
            d["unit"] = self.unit
        return d

    @classmethod
    def from_dict(cls, d: dict) -> "SmartRule":
        return cls(field=d.get("field", ""), operator=d.get("operator", ""),
                   value=d.get("value", ""), value2=d.get("value2"),
                   unit=d.get("unit"))


@dataclass
class SmartGroup:
    conjunction: str = "all"  # "all" (AND) | "any" (OR)
    children: list = field(default_factory=list)  # SmartRule | SmartGroup

    def to_dict(self) -> dict:
        return {"kind": "group", "conjunction": self.conjunction,
                "children": [c.to_dict() for c in self.children]}

    @classmethod
    def from_dict(cls, d: dict) -> "SmartGroup":
        return cls(conjunction=d.get("conjunction", "all"),
                   children=[_node_from_dict(c) for c in d.get("children", [])])


def _node_from_dict(d: dict):
    if d.get("kind") == "group" or "children" in d:
        return SmartGroup.from_dict(d)
    return SmartRule.from_dict(d)


@dataclass
class SmartLimit:
    enabled: bool = False
    count: int = 25
    by: str = "items"          # items|minutes|hours|mb|gb
    selection: str = "random"  # key into SELECTION_SORT

    def to_dict(self) -> dict:
        return {"enabled": self.enabled, "count": self.count,
                "by": self.by, "selection": self.selection}

    @classmethod
    def from_dict(cls, d: dict) -> "SmartLimit":
        return cls(enabled=d.get("enabled", False), count=d.get("count", 25),
                   by=d.get("by", "items"), selection=d.get("selection", "random"))


@dataclass
class SmartCriteria:
    root: SmartGroup = field(default_factory=SmartGroup)
    limit: SmartLimit = field(default_factory=SmartLimit)
    live_update: bool = True
    match_only_checked: bool = False
    # True when the iTunes blob couldn't be fully converted; the manager then
    # keeps the imported track snapshot instead of recomputing membership.
    unsupported: bool = False
    # True when the rule tree contains a nested group; the editor is read-only
    # for these in Phase 1 (nested-group editing UI is Phase 2).
    has_nested: bool = False

    def to_dict(self) -> dict:
        return {
            "root": self.root.to_dict(),
            "limit": self.limit.to_dict(),
            "live_update": self.live_update,
            "match_only_checked": self.match_only_checked,
            "unsupported": self.unsupported,
            "has_nested": self.has_nested,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "SmartCriteria":
        root_d = d.get("root")
        return cls(
            root=SmartGroup.from_dict(root_d) if root_d else SmartGroup(),
            limit=SmartLimit.from_dict(d.get("limit", {})),
            live_update=d.get("live_update", True),
            match_only_checked=d.get("match_only_checked", False),
            unsupported=d.get("unsupported", False),
            has_nested=d.get("has_nested", False),
        )

    def references_attr(self, attr: str) -> bool:
        """True if any rule matches on the given Track attribute."""
        return _group_references(self.root, attr)

    def referenced_attrs(self) -> set:
        out: set = set()
        _collect_attrs(self.root, out)
        return out


def _group_references(group: SmartGroup, attr: str) -> bool:
    for child in group.children:
        if isinstance(child, SmartGroup):
            if _group_references(child, attr):
                return True
        else:
            meta = FIELD_REGISTRY.get(child.field)
            if meta and meta.track_attr == attr:
                return True
    return False


def _collect_attrs(group: SmartGroup, out: set) -> None:
    for child in group.children:
        if isinstance(child, SmartGroup):
            _collect_attrs(child, out)
        else:
            meta = FIELD_REGISTRY.get(child.field)
            if meta:
                out.add(meta.track_attr)


# --------------------------------------------------------------------------- #
# Evaluation
# --------------------------------------------------------------------------- #

def _parse_dt(value) -> Optional[datetime]:
    """Parse a Track ISO date string into a naive-UTC datetime."""
    if not value:
        return None
    if isinstance(value, datetime):
        dt = value
    else:
        try:
            dt = datetime.fromisoformat(str(value))
        except (ValueError, TypeError):
            return None
    if dt.tzinfo is not None:
        dt = dt.replace(tzinfo=None)
    return dt


def _match_string(op, track_val, qval) -> bool:
    t = ("" if track_val is None else str(track_val)).casefold()
    q = ("" if qval is None else str(qval)).casefold()
    if op == "is":
        return t == q
    if op == "is_not":
        return t != q
    if op == "contains":
        return q in t
    if op == "not_contains":
        return q not in t
    if op == "starts_with":
        return t.startswith(q)
    if op == "ends_with":
        return t.endswith(q)
    return False


def _match_numeric(op, track_val, qval, qval2) -> bool:
    try:
        t = float(track_val or 0)
        q = float(qval or 0)
    except (ValueError, TypeError):
        return False
    if op == "is":
        return t == q
    if op == "is_not":
        return t != q
    if op == "greater":
        return t > q
    if op == "less":
        return t < q
    if op == "in_range":
        try:
            q2 = float(qval2 or 0)
        except (ValueError, TypeError):
            return False
        lo, hi = (q, q2) if q <= q2 else (q2, q)
        return lo <= t <= hi
    return False


def _match_date(op, track_val, rule: "SmartRule", now: datetime) -> bool:
    dt = _parse_dt(track_val)
    if op in ("after", "before", "in_range"):
        # iTunes treats a null date (never played/skipped/etc.) as the distant
        # past: "before X" matches it; "after"/"in range" don't.
        d = dt if dt is not None else datetime.min
        if op == "in_range":
            lo, hi = _parse_dt(rule.value), _parse_dt(rule.value2)
            if lo is None or hi is None:
                return False
            if lo > hi:
                lo, hi = hi, lo
            return lo <= d <= hi
        q = _parse_dt(rule.value)
        if q is None:
            return False
        return d > q if op == "after" else d < q
    if op == "not_in_range":
        if dt is None:
            return True
        lo, hi = _parse_dt(rule.value), _parse_dt(rule.value2)
        if lo is None or hi is None:
            return True
        if lo > hi:
            lo, hi = hi, lo
        return not (lo <= dt <= hi)
    if op in ("in_last", "not_in_last"):
        seconds = _rule_window_seconds(rule)
        if op == "in_last":
            return dt is not None and (now - dt) <= timedelta(seconds=seconds)
        return dt is None or (now - dt) > timedelta(seconds=seconds)
    return False


def _rule_window_seconds(rule: "SmartRule") -> float:
    try:
        count = float(rule.value or 0)
    except (ValueError, TypeError):
        count = 0
    return count * _UNIT_SECONDS.get(rule.unit or "days", 86400)


def _match_rule(rule: SmartRule, track, now: datetime) -> bool:
    meta = FIELD_REGISTRY.get(rule.field)
    if meta is None:
        return False
    track_val = getattr(track, meta.track_attr, None)
    if meta.type is FieldType.STRING:
        return _match_string(rule.operator, track_val, rule.value)
    if meta.type in (FieldType.INT, FieldType.DURATION, FieldType.RATING):
        return _match_numeric(rule.operator, track_val, rule.value, rule.value2)
    if meta.type is FieldType.DATE:
        return _match_date(rule.operator, track_val, rule, now)
    if meta.type is FieldType.BOOL:
        return bool(track_val) == bool(rule.value)
    return False


def _match_group(group: SmartGroup, track, now: datetime) -> bool:
    if not group.children:
        return True
    results = (
        _match_group(c, track, now) if isinstance(c, SmartGroup)
        else _match_rule(c, track, now)
        for c in group.children
    )
    return all(results) if group.conjunction == "all" else any(results)


def _selection_key(selection: str, pid_seed: int):
    attr, reverse = SELECTION_SORT.get(selection, (None, False))

    if attr is None:  # random -> deterministic shuffle (stable between runs)
        rng = _random.Random(pid_seed)

        def keyfn(track):
            return rng.random()
        return keyfn, False

    def keyfn(track):
        val = getattr(track, attr, None)
        if attr in ("play_date_utc", "date_added"):
            dt = _parse_dt(val)
            # Tracks never played/added sort oldest.
            return dt or datetime.min
        return val if val is not None else ""
    return keyfn, reverse


def _apply_limit(matches: list, limit: SmartLimit, seed: int) -> list:
    keyfn, reverse = _selection_key(limit.selection, seed)
    ordered = sorted(matches, key=keyfn, reverse=reverse)
    if limit.by == "items":
        return ordered[:max(0, limit.count)]
    # Cumulative budget by time or size.
    if limit.by in ("minutes", "hours"):
        budget = limit.count * (60000 if limit.by == "minutes" else 3600000)
        attr = "total_time"
    else:  # mb, gb
        budget = limit.count * (1024 ** 2 if limit.by == "mb" else 1024 ** 3)
        attr = "size"
    out, total = [], 0
    for track in ordered:
        amount = getattr(track, attr, 0) or 0
        if out and total + amount > budget:
            break
        out.append(track)
        total += amount
    return out


def evaluate(criteria: SmartCriteria, tracks, now: Optional[datetime] = None,
             seed: int = 0) -> list:
    """Return the ordered list of track ids matching `criteria`.

    `tracks` is any iterable of Track objects. `now` (naive UTC) is injectable
    for deterministic date tests. `seed` makes "limit ... selected by random"
    stable between recomputes (pass the playlist's persistent id hash).
    """
    if now is None:
        now = datetime.now(timezone.utc).replace(tzinfo=None)
    matched = [t for t in tracks if _match_group(criteria.root, t, now)]
    if criteria.limit.enabled and criteria.limit.count > 0:
        matched = _apply_limit(matched, criteria.limit, seed)
    return sorted(t.track_id for t in matched)


# --------------------------------------------------------------------------- #
# iTunes import
# --------------------------------------------------------------------------- #

class _UnsupportedCriteria(Exception):
    """Raised mid-conversion when a rule can't be represented in our model."""


# cvzi field enum name -> our registry key.
_ITUNES_FIELD_TO_KEY = {
    "Name": "name", "Artist": "artist", "AlbumArtist": "album_artist",
    "Album": "album", "Genre": "genre", "Composer": "composer",
    "Grouping": "grouping", "Comments": "comments", "Kind": "kind",
    "SortName": "sort_name", "SortAlbum": "sort_album",
    "Year": "year", "BPM": "bpm", "BitRate": "bit_rate",
    "SampleRate": "sample_rate", "Size": "size", "DiskNumber": "disc_number",
    "TrackNumber": "track_number", "Plays": "play_count", "Skips": "skip_count",
    "Rating": "rating", "Duration": "total_time", "Compilation": "compilation",
    "DateAdded": "date_added", "DateModified": "date_modified",
    "LastPlayed": "last_played", "LastSkipped": "last_skipped",
}

_ITUNES_STRING_OP = {
    "like": "contains", "not like": "not_contains", "is": "is",
    "is not": "is_not", "starts with": "starts_with", "ends with": "ends_with",
}
_ITUNES_INT_OP = {
    "is": "is", "is not": "is_not", "greater than": "greater",
    "less than": "less", "between": "in_range",
}
_ITUNES_DATE_OP = {
    "is after": "after", "is before": "before", "is in the range": "in_range",
    "is not in the range": "not_in_range", "is in the last": "in_last",
    "is not in the last": "not_in_last",
}


def _unix_to_iso(ts) -> str:
    return datetime.fromtimestamp(int(ts), timezone.utc).replace(tzinfo=None).isoformat()


def _convert_leaf(leaf: dict):
    """Convert a cvzi fullTree leaf dict into a SmartRule (or None to drop)."""
    fname = leaf.get("field", "")
    ltype = leaf.get("type", "")
    op = leaf.get("operator", "")
    value = leaf.get("value")

    # Fields with no Track backing.
    if fname == "Love":
        return None  # user does not use "loved"; drop the rule entirely
    if fname == "MediaKind":
        if op == "is" and value in ("Music", "Music Video"):
            return None  # implicit/no-op filter in a music-only library
        raise _UnsupportedCriteria("MediaKind rule")

    key = _ITUNES_FIELD_TO_KEY.get(fname)
    if key is None or key not in FIELD_REGISTRY:
        raise _UnsupportedCriteria("unmapped field %r" % fname)
    meta = FIELD_REGISTRY[key]

    if meta.type is FieldType.STRING:
        our_op = _ITUNES_STRING_OP.get(op)
        if our_op is None:
            raise _UnsupportedCriteria("string op %r" % op)
        return SmartRule(field=key, operator=our_op, value=str(value))

    if meta.type is FieldType.BOOL:  # compilation
        truth = bool(value)
        if op == "is not":
            truth = not truth
        return SmartRule(field=key, operator="is", value=truth)

    if meta.type in (FieldType.INT, FieldType.DURATION, FieldType.RATING):
        our_op = _ITUNES_INT_OP.get(op)
        if our_op is None:
            raise _UnsupportedCriteria("int op %r" % op)
        scale = 20 if meta.type is FieldType.RATING else 1
        if our_op == "in_range" and isinstance(value, (tuple, list)):
            return SmartRule(field=key, operator="in_range",
                             value=int(value[0]) * scale,
                             value2=int(value[1]) * scale)
        return SmartRule(field=key, operator=our_op, value=int(value) * scale)

    if meta.type is FieldType.DATE:
        our_op = _ITUNES_DATE_OP.get(op)
        if our_op is None:
            raise _UnsupportedCriteria("date op %r" % op)
        if our_op in ("in_last", "not_in_last"):
            count, unit = _parse_relative_date(leaf.get("value_date"), value)
            return SmartRule(field=key, operator=our_op, value=count, unit=unit)
        if our_op in ("in_range", "not_in_range") and isinstance(value, (tuple, list)):
            return SmartRule(field=key, operator=our_op,
                             value=_unix_to_iso(value[0]),
                             value2=_unix_to_iso(value[1]))
        return SmartRule(field=key, operator=our_op, value=_unix_to_iso(value))

    raise _UnsupportedCriteria("unhandled type %r" % meta.type)


def _parse_relative_date(value_date, raw_seconds):
    """cvzi gives value_date like '12 months'; fall back to raw seconds."""
    if isinstance(value_date, str):
        parts = value_date.split()
        if len(parts) == 2 and parts[1] in _UNIT_SECONDS:
            try:
                return int(parts[0]), parts[1]
            except ValueError:
                pass
    # Fallback: convert raw seconds to whole days.
    try:
        return max(1, int(int(raw_seconds) // 86400)), "days"
    except (ValueError, TypeError):
        return 1, "days"


def _convert_group(node: dict) -> SmartGroup:
    op = "and" if "and" in node else "or" if "or" in node else None
    if op is None:
        return SmartGroup("all", [])
    children = []
    for child in node[op]:
        if isinstance(child, dict) and ("and" in child or "or" in child):
            children.append(_convert_group(child))
        else:
            rule = _convert_leaf(child)
            if rule is not None:
                children.append(rule)
    return SmartGroup("all" if op == "and" else "any", children)


def _simplify(group: SmartGroup) -> SmartGroup:
    """Drop empty groups (left behind by dropped MediaKind/Love rules) and
    unwrap a group whose only child is itself a group. iTunes wraps every
    playlist in an implicit "MediaKind is Music" group; stripping that keeps
    ordinary playlists flat (and therefore editable)."""
    new_children = []
    for child in group.children:
        if isinstance(child, SmartGroup):
            child = _simplify(child)
            if not child.children:
                continue
        new_children.append(child)
    group.children = new_children
    if len(group.children) == 1 and isinstance(group.children[0], SmartGroup):
        return group.children[0]
    return group


def _contains_group(group: SmartGroup) -> bool:
    return any(isinstance(c, SmartGroup) for c in group.children)


def parse_itunes_smart(info: Optional[bytes],
                       criteria: Optional[bytes]) -> SmartCriteria:
    """Decode the iTunes "Smart Info"/"Smart Criteria" blobs into a
    SmartCriteria. Never raises: on any problem it returns a criteria flagged
    `unsupported=True`, so the importer falls back to the track snapshot.
    """
    if not info or not criteria:
        return SmartCriteria(unsupported=True)
    info, criteria = bytes(info), bytes(criteria)
    # Real iTunes criteria blobs always start with the "SLst" magic. Anything
    # else is junk we shouldn't trust (and would otherwise parse to an empty
    # "match everything" rule set), so fall back to the snapshot.
    if not criteria.startswith(b"SLst") or len(info) < 14:
        return SmartCriteria(unsupported=True)
    try:
        from lintunes.itunes_smart.parse import SmartPlaylistParser
        parser = SmartPlaylistParser()
        parser.data(info, criteria)
        parser.parse()
    except Exception as exc:  # noqa: BLE001 - import must never abort
        logger.warning("smart criteria parse failed: %r", exc)
        return SmartCriteria(unsupported=True)

    limit = _read_limit(bytes(info))
    result = SmartCriteria(limit=limit,
                           live_update=bool(info[0]),
                           match_only_checked=bool(info[12]))

    if getattr(parser, "ignore", "").strip():
        # The parser hit a field/case it couldn't decode -> incomplete tree.
        result.unsupported = True
        return result

    try:
        root = _simplify(_convert_group(parser.fullTreeRoot or {}))
    except _UnsupportedCriteria as exc:
        logger.info("smart criteria not fully supported: %s", exc)
        result.unsupported = True
        return result
    result.root = root
    result.has_nested = _contains_group(root)
    return result


_LIMIT_METHOD = {1: "minutes", 2: "mb", 3: "items", 4: "hours", 5: "gb"}
# (iTunes selection-method byte, "is least" sign) -> our selection key.
_SELECTION_BASE = {
    0x02: "random", 0x05: "name", 0x06: "album", 0x07: "artist", 0x09: "genre",
    0x1c: "highest_rating", 0x01: "lowest_rating",
}
_SELECTION_SIGNED = {
    0x1a: ("most_recently_played", "least_recently_played"),
    0x19: ("most_often_played", "least_often_played"),
    0x15: ("most_recently_added", "least_recently_added"),
}


def _read_limit(info: bytes) -> SmartLimit:
    if len(info) < 14 or info[2] != 1:
        return SmartLimit(enabled=False)
    count = int.from_bytes(info[8:12], "big")
    by = _LIMIT_METHOD.get(info[3], "items")
    method = info[7]
    if method in _SELECTION_SIGNED:
        # cvzi: sign == 1 (i.e. the sign byte is 0) selects the "least" variant.
        least = (info[13] == 0)
        selection = _SELECTION_SIGNED[method][1 if least else 0]
    else:
        selection = _SELECTION_BASE.get(method, "random")
    return SmartLimit(enabled=True, count=count, by=by, selection=selection)