205 lines
7.4 KiB
Python
205 lines
7.4 KiB
Python
"""JSON-backed queue store with file locking for thread safety."""
|
|
|
|
import json
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from threading import Lock
|
|
|
|
from models import QueueItem
|
|
|
|
|
|
class QueueStore:
|
|
"""Persistent queue backed by a JSON file."""
|
|
|
|
def __init__(self, store_path: str = None):
|
|
if store_path is None:
|
|
store_path = str(Path.home() / ".config" / "youtube_cli" / "queue.json")
|
|
self.store_path = store_path
|
|
self._lock = Lock()
|
|
self._ensure_file()
|
|
|
|
def _ensure_file(self):
|
|
"""Create the store file if it doesn't exist."""
|
|
store_dir = Path(self.store_path).parent
|
|
store_dir.mkdir(parents=True, exist_ok=True)
|
|
if not Path(self.store_path).exists():
|
|
with open(self.store_path, "w") as f:
|
|
json.dump({}, f)
|
|
|
|
def _load(self) -> dict:
|
|
"""Load queue data from file."""
|
|
try:
|
|
with open(self.store_path, "r") as f:
|
|
return json.load(f)
|
|
except (json.JSONDecodeError, ValueError):
|
|
with open(self.store_path, "w") as f:
|
|
json.dump({}, f)
|
|
return {}
|
|
|
|
def _save(self, data: dict):
|
|
"""Save queue data to file."""
|
|
with open(self.store_path, "w") as f:
|
|
json.dump(data, f, indent=2)
|
|
|
|
def add_item(self, item: QueueItem) -> QueueItem:
|
|
"""Add an item to the queue."""
|
|
with self._lock:
|
|
data = self._load()
|
|
data[item.id] = item.to_dict()
|
|
self._save(data)
|
|
return item
|
|
|
|
def get_all(self) -> list:
|
|
"""Get all queue items."""
|
|
with self._lock:
|
|
data = self._load()
|
|
items = []
|
|
for item_id, item_data in data.items():
|
|
item = self._dict_to_item(item_data)
|
|
items.append(item)
|
|
return items
|
|
|
|
def get_item(self, queue_id: str) -> QueueItem:
|
|
"""Get a specific queue item."""
|
|
with self._lock:
|
|
data = self._load()
|
|
item_data = data.get(queue_id)
|
|
if item_data:
|
|
return self._dict_to_item(item_data)
|
|
return None
|
|
|
|
def update_item(self, queue_id: str, updates: dict) -> QueueItem:
|
|
"""Update fields of a queue item."""
|
|
with self._lock:
|
|
data = self._load()
|
|
if queue_id not in data:
|
|
return None
|
|
data[queue_id].update(updates)
|
|
self._save(data)
|
|
item = self._dict_to_item(data[queue_id])
|
|
return item
|
|
|
|
def update_progress(self, queue_id: str, progress: float, speed: str = None, eta: str = None):
|
|
"""Update download progress for a queue item."""
|
|
with self._lock:
|
|
data = self._load()
|
|
if queue_id in data:
|
|
data[queue_id]["progress"] = progress
|
|
if speed:
|
|
data[queue_id]["speed"] = speed
|
|
if eta:
|
|
data[queue_id]["eta"] = eta
|
|
self._save(data)
|
|
|
|
def update_status(self, queue_id: str, status: str, error_message: str = None,
|
|
download_path: str = None, file_size: str = None):
|
|
"""Update download status for a queue item."""
|
|
with self._lock:
|
|
data = self._load()
|
|
if queue_id in data:
|
|
data[queue_id]["status"] = status
|
|
if status in ("completed", "failed"):
|
|
data[queue_id]["completedAt"] = datetime.now(timezone.utc).isoformat()
|
|
if error_message:
|
|
data[queue_id]["errorMessage"] = error_message
|
|
if download_path:
|
|
data[queue_id]["downloadPath"] = download_path
|
|
if file_size:
|
|
data[queue_id]["fileSize"] = file_size
|
|
if status == "completed":
|
|
data[queue_id]["progress"] = 100.0
|
|
self._save(data)
|
|
|
|
def remove_item(self, queue_id: str) -> bool:
|
|
"""Remove an item from the queue."""
|
|
with self._lock:
|
|
data = self._load()
|
|
if queue_id in data:
|
|
del data[queue_id]
|
|
self._save(data)
|
|
return True
|
|
return False
|
|
|
|
def clear_completed(self) -> int:
|
|
"""Clear all completed items. Returns count of removed items."""
|
|
with self._lock:
|
|
data = self._load()
|
|
completed_ids = [qid for qid, item in data.items() if item["status"] == "completed"]
|
|
for qid in completed_ids:
|
|
del data[qid]
|
|
self._save(data)
|
|
return len(completed_ids)
|
|
|
|
def clear_failed(self) -> int:
|
|
"""Clear all failed items. Returns count of removed items."""
|
|
with self._lock:
|
|
data = self._load()
|
|
failed_ids = [qid for qid, item in data.items() if item["status"] == "failed"]
|
|
for qid in failed_ids:
|
|
del data[qid]
|
|
self._save(data)
|
|
return len(failed_ids)
|
|
|
|
def clear_all(self) -> int:
|
|
"""Clear all items from the queue. Returns count of removed items."""
|
|
with self._lock:
|
|
count = len(self._load())
|
|
self._save({})
|
|
return count
|
|
|
|
def get_stats(self) -> dict:
|
|
"""Get queue statistics."""
|
|
items = self.get_all()
|
|
return {
|
|
"total": len(items),
|
|
"pending": sum(1 for i in items if i.status == "pending"),
|
|
"downloading": sum(1 for i in items if i.status == "downloading"),
|
|
"completed": sum(1 for i in items if i.status == "completed"),
|
|
"failed": sum(1 for i in items if i.status == "failed"),
|
|
"cancelled": sum(1 for i in items if i.status == "cancelled"),
|
|
}
|
|
|
|
def reorder_item(self, queue_id: str, direction: str) -> bool:
|
|
"""Reorder a queue item (up/down). Returns True if reordered."""
|
|
with self._lock:
|
|
data = self._load()
|
|
ids = list(data.keys())
|
|
if queue_id not in ids:
|
|
return False
|
|
idx = ids.index(queue_id)
|
|
if direction == "up" and idx > 0:
|
|
ids[idx], ids[idx - 1] = ids[idx - 1], ids[idx]
|
|
elif direction == "down" and idx < len(ids) - 1:
|
|
ids[idx], ids[idx + 1] = ids[idx + 1], ids[idx]
|
|
else:
|
|
return False
|
|
# Rebuild dict in new order
|
|
new_data = {}
|
|
for kid in ids:
|
|
new_data[kid] = data[kid]
|
|
self._save(new_data)
|
|
return True
|
|
|
|
def _dict_to_item(self, data: dict) -> QueueItem:
|
|
"""Convert a dictionary to a QueueItem."""
|
|
return QueueItem(
|
|
id=data["id"],
|
|
video_id=data.get("videoId", ""),
|
|
title=data.get("title", ""),
|
|
url=data.get("url", ""),
|
|
thumbnail=data.get("thumbnail", ""),
|
|
status=data.get("status", "pending"),
|
|
progress=data.get("progress", 0.0),
|
|
category=data.get("category", ""),
|
|
network_folder=data.get("network_folder"),
|
|
added_at=data.get("addedAt", data.get("created_at", datetime.now(timezone.utc).isoformat())),
|
|
completed_at=data.get("completedAt"),
|
|
error_message=data.get("errorMessage", data.get("message")),
|
|
download_path=data.get("downloadPath"),
|
|
file_size=data.get("fileSize"),
|
|
speed=data.get("speed"),
|
|
eta=data.get("eta"),
|
|
item_type=data.get("type", "video"),
|
|
quality=data.get("quality"),
|
|
)
|