Medios-Macina/Provider/soulseek.py

from __future__ import annotations

import asyncio
import contextlib
import io
import logging
import os
import re
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

from ProviderCore.base import SearchProvider, SearchResult
from SYS.logger import log, debug


_SOULSEEK_NOISE_SUBSTRINGS = (
	"search reply ticket does not match any search request",
	"failed to receive transfer ticket on file connection",
	"aioslsk.exceptions.ConnectionReadError",
)


def _configure_aioslsk_logging() -> None:
	"""Reduce aioslsk internal log noise.

	Some aioslsk components emit non-fatal warnings/errors during high churn
	(search + download + disconnect). We keep our own debug output, but push
	aioslsk to ERROR and stop propagation so it doesn't spam the CLI.
	"""
	for name in (
		"aioslsk",
		"aioslsk.network",
		"aioslsk.search",
		"aioslsk.transfer",
		"aioslsk.transfer.manager",
	):
		logger = logging.getLogger(name)
		logger.setLevel(logging.ERROR)
		logger.propagate = False


class _LineFilterStream(io.TextIOBase):
	"""A minimal stream wrapper that filters known noisy lines.

	It also suppresses entire traceback blocks when they contain known non-fatal
	aioslsk noise (e.g. ConnectionReadError during peer init).
	"""

	def __init__(self, underlying: Any, suppress_substrings: tuple[str, ...]):
		super().__init__()
		self._underlying = underlying
		self._suppress = suppress_substrings
		self._buf = ""
		self._in_tb = False
		self._tb_lines: list[str] = []
		self._tb_suppress = False

	def writable(self) -> bool:  # pragma: no cover
		return True

	def _should_suppress_line(self, line: str) -> bool:
		return any(sub in line for sub in self._suppress)

	def _flush_tb(self) -> None:
		if not self._tb_lines:
			return
		if not self._tb_suppress:
			for l in self._tb_lines:
				try:
					self._underlying.write(l + "\n")
				except Exception:
					pass
		self._tb_lines = []
		self._tb_suppress = False
		self._in_tb = False

	def write(self, s: str) -> int:
		self._buf += str(s)
		while "\n" in self._buf:
			line, self._buf = self._buf.split("\n", 1)
			self._handle_line(line)
		return len(s)

	def _handle_line(self, line: str) -> None:
		# Start capturing tracebacks so we can suppress the whole block if it matches.
		if not self._in_tb and line.startswith("Traceback (most recent call last):"):
			self._in_tb = True
			self._tb_lines = [line]
			self._tb_suppress = False
			return

		if self._in_tb:
			self._tb_lines.append(line)
			if self._should_suppress_line(line):
				self._tb_suppress = True
			# End traceback block on blank line.
			if line.strip() == "":
				self._flush_tb()
			return

		# Non-traceback line
		if self._should_suppress_line(line):
			return
		try:
			self._underlying.write(line + "\n")
		except Exception:
			pass

	def flush(self) -> None:
		# Flush any pending traceback block.
		if self._in_tb:
			# If the traceback ends without a trailing blank line, decide here.
			self._flush_tb()
		if self._buf:
			line = self._buf
			self._buf = ""
			if not self._should_suppress_line(line):
				try:
					self._underlying.write(line)
				except Exception:
					pass
		try:
			self._underlying.flush()
		except Exception:
			pass


@contextlib.contextmanager
def _suppress_aioslsk_noise() -> Any:
	"""Temporarily suppress known aioslsk noise printed to stdout/stderr.

	Opt out by setting DOWNLOW_SOULSEEK_VERBOSE=1.
	"""
	if os.environ.get("DOWNLOW_SOULSEEK_VERBOSE"):
		yield
		return

	_configure_aioslsk_logging()
	old_out, old_err = sys.stdout, sys.stderr
	sys.stdout = _LineFilterStream(old_out, _SOULSEEK_NOISE_SUBSTRINGS)
	sys.stderr = _LineFilterStream(old_err, _SOULSEEK_NOISE_SUBSTRINGS)
	try:
		yield
	finally:
		try:
			sys.stdout.flush()
			sys.stderr.flush()
		except Exception:
			pass
		sys.stdout, sys.stderr = old_out, old_err


class Soulseek(SearchProvider):
	"""Search provider for Soulseek P2P network."""

	MUSIC_EXTENSIONS = {
		".flac",
		".mp3",
		".m4a",
		".aac",
		".ogg",
		".opus",
		".wav",
		".alac",
		".wma",
		".ape",
		".aiff",
		".dsf",
		".dff",
		".wv",
		".tta",
		".tak",
		".ac3",
		".dts",
	}

	# NOTE: These defaults preserve existing behavior.
	USERNAME = "asjhkjljhkjfdsd334"
	PASSWORD = "khhhg"
	DOWNLOAD_DIR = "./downloads"
	MAX_WAIT_TRANSFER = 1200

	def download(self, result: SearchResult, output_dir: Path) -> Optional[Path]:
		"""Download file from Soulseek."""

		try:
			full_metadata = result.full_metadata or {}
			username = full_metadata.get("username")
			filename = full_metadata.get("filename") or result.path

			if not username or not filename:
				log(f"[soulseek] Missing metadata for download: {result.title}", file=sys.stderr)
				return None

			# This cmdlet stack is synchronous; use asyncio.run for clarity.
			return asyncio.run(
				download_soulseek_file(
					username=username,
					filename=filename,
					output_dir=output_dir,
					timeout=self.MAX_WAIT_TRANSFER,
				)
			)

		except RuntimeError:
			# If we're already inside an event loop (e.g., TUI), fall back to a
			# dedicated loop in this thread.
			loop = asyncio.new_event_loop()
			try:
				asyncio.set_event_loop(loop)
				return loop.run_until_complete(
					download_soulseek_file(
						username=username,
						filename=filename,
						output_dir=output_dir,
						timeout=self.MAX_WAIT_TRANSFER,
					)
				)
			finally:
				try:
					loop.close()
				except Exception:
					pass

		except Exception as exc:
			log(f"[soulseek] Download error: {exc}", file=sys.stderr)
			return None

	async def perform_search(self, query: str, timeout: float = 9.0, limit: int = 50) -> List[Dict[str, Any]]:
		"""Perform async Soulseek search."""

		from aioslsk.client import SoulSeekClient
		from aioslsk.settings import CredentialsSettings, Settings

		os.makedirs(self.DOWNLOAD_DIR, exist_ok=True)

		settings = Settings(credentials=CredentialsSettings(username=self.USERNAME, password=self.PASSWORD))
		client = SoulSeekClient(settings)

		with _suppress_aioslsk_noise():
			try:
				await client.start()
				await client.login()
			except Exception as exc:
				log(f"[soulseek] Login failed: {type(exc).__name__}: {exc}", file=sys.stderr)
				return []

			try:
				search_request = await client.searches.search(query)
				await self._collect_results(search_request, timeout=timeout)
				return self._flatten_results(search_request)[:limit]
			except Exception as exc:
				log(f"[soulseek] Search error: {type(exc).__name__}: {exc}", file=sys.stderr)
				return []
			finally:
				# Best-effort: try to cancel/close the search request before stopping
				# the client to reduce stray reply spam.
				try:
					if "search_request" in locals() and search_request is not None:
						cancel = getattr(search_request, "cancel", None)
						if callable(cancel):
							maybe = cancel()
							if asyncio.iscoroutine(maybe):
								await maybe
				except Exception:
					pass
				try:
					await client.stop()
				except Exception:
					pass

	def _flatten_results(self, search_request: Any) -> List[dict]:
		flat: List[dict] = []
		for result in getattr(search_request, "results", []):
			username = getattr(result, "username", "?")

			for file_data in getattr(result, "shared_items", []):
				flat.append(
					{
						"file": file_data,
						"username": username,
						"filename": getattr(file_data, "filename", "?"),
						"size": getattr(file_data, "filesize", 0),
					}
				)

			for file_data in getattr(result, "locked_results", []):
				flat.append(
					{
						"file": file_data,
						"username": username,
						"filename": getattr(file_data, "filename", "?"),
						"size": getattr(file_data, "filesize", 0),
					}
				)

		return flat

	async def _collect_results(self, search_request: Any, timeout: float = 75.0) -> None:
		end = time.time() + timeout
		last_count = 0
		while time.time() < end:
			current_count = len(getattr(search_request, "results", []))
			if current_count > last_count:
				debug(f"[soulseek] Got {current_count} result(s)...")
				last_count = current_count
			await asyncio.sleep(0.5)

	def search(
		self,
		query: str,
		limit: int = 50,
		filters: Optional[Dict[str, Any]] = None,
		**kwargs: Any,
	) -> List[SearchResult]:
		filters = filters or {}

		try:
			flat_results = asyncio.run(self.perform_search(query, timeout=9.0, limit=limit))
			if not flat_results:
				return []

			music_results: List[dict] = []
			for item in flat_results:
				filename = item["filename"]
				ext = ("." + filename.rsplit(".", 1)[-1].lower()) if "." in filename else ""
				if ext in self.MUSIC_EXTENSIONS:
					music_results.append(item)

			if not music_results:
				return []

			enriched_results: List[dict] = []
			for item in music_results:
				filename = item["filename"]
				ext = ("." + filename.rsplit(".", 1)[-1].lower()) if "." in filename else ""

				display_name = filename.replace("\\", "/").split("/")[-1]
				path_parts = filename.replace("\\", "/").split("/")
				artist = path_parts[-3] if len(path_parts) >= 3 else ""
				album = path_parts[-2] if len(path_parts) >= 3 else (path_parts[-2] if len(path_parts) == 2 else "")

				base_name = display_name.rsplit(".", 1)[0] if "." in display_name else display_name
				track_num = ""
				title = base_name
				filename_artist = ""

				match = re.match(r"^(\d{1,3})\s*[\.\-]?\s+(.+)$", base_name)
				if match:
					track_num = match.group(1)
					rest = match.group(2)
					if " - " in rest:
						filename_artist, title = rest.split(" - ", 1)
					else:
						title = rest

				if filename_artist:
					artist = filename_artist

				enriched_results.append(
					{
						**item,
						"artist": artist,
						"album": album,
						"title": title,
						"track_num": track_num,
						"ext": ext,
					}
				)

			if filters:
				artist_filter = (filters.get("artist", "") or "").lower()
				album_filter = (filters.get("album", "") or "").lower()
				track_filter = (filters.get("track", "") or "").lower()

				if artist_filter or album_filter or track_filter:
					filtered: List[dict] = []
					for item in enriched_results:
						if artist_filter and artist_filter not in item["artist"].lower():
							continue
						if album_filter and album_filter not in item["album"].lower():
							continue
						if track_filter and track_filter not in item["title"].lower():
							continue
						filtered.append(item)
					enriched_results = filtered

			enriched_results.sort(key=lambda item: (item["ext"].lower() != ".flac", -item["size"]))

			results: List[SearchResult] = []
			for item in enriched_results:
				artist_display = item["artist"] if item["artist"] else "(no artist)"
				album_display = item["album"] if item["album"] else "(no album)"
				size_mb = int(item["size"] / 1024 / 1024)

				columns = [
					("Track", item["track_num"] or "?"),
					("Title", item["title"][:40]),
					("Artist", artist_display[:32]),
					("Album", album_display[:32]),
					("Size", f"{size_mb} MB"),
				]

				results.append(
					SearchResult(
						table="soulseek",
						title=item["title"],
						path=item["filename"],
						detail=f"{artist_display} - {album_display}",
						annotations=[f"{size_mb} MB", item["ext"].lstrip(".").upper()],
						media_kind="audio",
						size_bytes=item["size"],
						columns=columns,
						full_metadata={
							"username": item["username"],
							"filename": item["filename"],
							"artist": item["artist"],
							"album": item["album"],
							"track_num": item["track_num"],
							"ext": item["ext"],
						},
					)
				)

			return results

		except Exception as exc:
			log(f"[soulseek] Search error: {exc}", file=sys.stderr)
			return []

	def validate(self) -> bool:
		try:
			from aioslsk.client import SoulSeekClient  # noqa: F401

			return True
		except ImportError:
			return False


async def download_soulseek_file(
	username: str,
	filename: str,
	output_dir: Path = Path("./downloads"),
	timeout: int = 1200,
) -> Optional[Path]:
	"""Download a file from a Soulseek peer."""

	try:
		from aioslsk.client import SoulSeekClient
		from aioslsk.settings import CredentialsSettings, Settings
		from aioslsk.transfer.model import Transfer, TransferDirection
		from aioslsk.transfer.state import TransferState

		output_dir = Path(output_dir)
		output_dir.mkdir(parents=True, exist_ok=True)

		local_filename = filename.replace("\\", "/").split("/")[-1]
		output_user_dir = output_dir / username
		output_user_dir.mkdir(parents=True, exist_ok=True)
		output_path = (output_user_dir / local_filename)

		if output_path.exists():
			base = output_path.stem
			ext = output_path.suffix
			counter = 1
			while output_path.exists():
				output_path = output_user_dir / f"{base}_{counter}{ext}"
				counter += 1

		output_path = output_path.resolve()

		settings = Settings(credentials=CredentialsSettings(username=Soulseek.USERNAME, password=Soulseek.PASSWORD))
		client = SoulSeekClient(settings)

		with _suppress_aioslsk_noise():
			try:
				await client.start()
				await client.login()
				debug(f"[soulseek] Logged in as {Soulseek.USERNAME}")

				debug(f"[soulseek] Requesting download from {username}: {filename}")

				transfer = await client.transfers.add(Transfer(username, filename, TransferDirection.DOWNLOAD))
				transfer.local_path = str(output_path)
				await client.transfers.queue(transfer)

				start_time = time.time()
				last_log_time = 0.0
				while not transfer.is_finalized():
					if time.time() - start_time > timeout:
						log(f"[soulseek] Download timeout after {timeout}s", file=sys.stderr)
						return None

					if time.time() - last_log_time >= 5.0 and transfer.bytes_transfered > 0:
						progress = (transfer.bytes_transfered / transfer.filesize * 100) if transfer.filesize else 0
						debug(
							f"[soulseek] Progress: {progress:.1f}% "
							f"({transfer.bytes_transfered}/{transfer.filesize})"
						)
						last_log_time = time.time()

					await asyncio.sleep(1)

				if transfer.state.VALUE == TransferState.COMPLETE and transfer.local_path:
					downloaded_path = Path(transfer.local_path)
					if downloaded_path.exists():
						debug(f"[soulseek] Download complete: {downloaded_path}")
						return downloaded_path

					log(f"[soulseek] Transfer completed but file missing: {downloaded_path}", file=sys.stderr)
					return None

				log(
					f"[soulseek] Download failed: state={transfer.state.VALUE} "
					f"bytes={transfer.bytes_transfered}/{transfer.filesize}",
					file=sys.stderr,
				)
				return None

			finally:
				try:
					await client.stop()
				except Exception:
					pass

	except ImportError:
		log("[soulseek] aioslsk not installed. Install with: pip install aioslsk", file=sys.stderr)
		return None
	except Exception as exc:
		log(f"[soulseek] Download failed: {type(exc).__name__}: {exc}", file=sys.stderr)
		return None