StashHandler/stash_handler.py

#!/usr/bin/env python3
"""StashHandler — bridge between qBittorrent and Stash."""

import argparse
import logging
import os
import signal
import shutil
import time
from pathlib import Path

import db
from hasher import hash_file
from qbit_client import QBitClient

log = logging.getLogger("stash_handler")

DEFAULT_MEDIA_EXTENSIONS = ".mp4,.mkv,.avi,.wmv,.mov,.webm,.flv,.m4v"


def load_config():
    """Build config dict from environment variables."""
    return {
        "qbit_url": os.environ.get("QBIT_URL", "http://qbittorrent:8080"),
        "qbit_username": os.environ.get("QBIT_USERNAME", "admin"),
        "qbit_password": os.environ.get("QBIT_PASSWORD", "adminadmin"),
        "staging_dir": os.environ.get("STAGING_DIR", "/staging"),
        "import_dir": os.environ.get("IMPORT_DIR", "/import"),
        "poll_interval": int(os.environ.get("POLL_INTERVAL", "60")),
        "media_extensions": set(
            os.environ.get("MEDIA_EXTENSIONS", DEFAULT_MEDIA_EXTENSIONS).split(",")
        ),
        "hash_algorithm": os.environ.get("HASH_ALGORITHM", "sha256"),
        "move_mode": os.environ.get("MOVE_MODE", "copy"),
        "qbit_category": os.environ.get("QBIT_CATEGORY", "Stash"),
        "db_path": os.environ.get("DB_PATH", "/data/seen.db"),
        "log_level": os.environ.get("LOG_LEVEL", "INFO"),
    }


def setup_logging(level="INFO"):
    logging.basicConfig(
        level=getattr(logging, level.upper(), logging.INFO),
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )


def is_media(path, extensions):
    return path.suffix.lower() in extensions


def walk_media(root, extensions):
    """Yield all media file paths under root."""
    root = Path(root)
    for p in sorted(root.rglob("*")):
        if p.is_file() and is_media(p, extensions):
            yield p


def transfer_file(src, dst, mode):
    """Move, hardlink, or copy src to dst."""
    dst.parent.mkdir(parents=True, exist_ok=True)
    if mode == "hardlink":
        try:
            os.link(src, dst)
        except OSError:
            log.warning("Hardlink failed, falling back to copy: %s", src)
            shutil.copy2(src, dst)
    elif mode == "copy":
        shutil.copy2(src, dst)
    else:  # move
        shutil.move(str(src), str(dst))


# --- Daemon mode ---

def daemon(cfg):
    qbit = QBitClient(cfg["qbit_url"], cfg["qbit_username"], cfg["qbit_password"])
    conn = db.init_db(cfg["db_path"])
    extensions = cfg["media_extensions"]
    staging = Path(cfg["staging_dir"])
    import_dir = Path(cfg["import_dir"])
    algorithm = cfg["hash_algorithm"]
    move_mode = cfg["move_mode"]
    poll_interval = cfg["poll_interval"]

    running = True

    def stop(*_):
        nonlocal running
        log.info("Shutting down...")
        running = False

    signal.signal(signal.SIGINT, stop)
    signal.signal(signal.SIGTERM, stop)

    log.info("Daemon started — polling every %ds", poll_interval)

    while running:
        try:
            torrents = qbit.get_completed_torrents(cfg["qbit_category"])
        except Exception:
            log.exception("Failed to fetch torrents")
            _sleep(poll_interval, lambda: running)
            continue

        for t in torrents:
            if not running:
                break
            info_hash = t["hash"]
            name = t["name"]

            if db.is_torrent_processed(conn, info_hash):
                continue

            log.info("Processing torrent: %s (%s)", name, info_hash)
            torrent_path = staging / t.get("content_path", name).lstrip("/")
            if not torrent_path.exists():
                # Fall back: content_path from qBit may be absolute or relative
                torrent_path = staging / name

            new_count = 0
            dup_count = 0
            for fpath in walk_media(torrent_path if torrent_path.is_dir() else torrent_path.parent, extensions):
                file_hash = hash_file(fpath, algorithm)
                if db.is_known(conn, file_hash):
                    db.increment_seen(conn, file_hash)
                    dup_count += 1
                    log.debug("Duplicate skipped: %s", fpath.name)
                else:
                    rel = fpath.relative_to(staging)
                    dst = import_dir / rel
                    transfer_file(fpath, dst, move_mode)
                    db.record_file(conn, file_hash, str(rel), "torrent", info_hash)
                    new_count += 1
                    log.info("Imported: %s", rel)

            db.record_torrent(conn, info_hash, name)
            log.info("Torrent done: %s — %d new, %d duplicates", name, new_count, dup_count)

        _sleep(poll_interval, lambda: running)

    conn.close()
    log.info("Shutdown complete")


def _sleep(seconds, check_fn, step=1):
    """Sleep in small steps so we can respond to shutdown quickly."""
    elapsed = 0
    while elapsed < seconds and check_fn():
        time.sleep(min(step, seconds - elapsed))
        elapsed += step


# --- Scan mode ---

def scan(cfg, path, do_import):
    conn = db.init_db(cfg["db_path"])
    extensions = cfg["media_extensions"]
    import_dir = Path(cfg["import_dir"])
    algorithm = cfg["hash_algorithm"]
    move_mode = cfg["move_mode"]
    scan_root = Path(path)

    total = 0
    new = 0
    duplicates = 0

    for fpath in walk_media(scan_root, extensions):
        total += 1
        file_hash = hash_file(fpath, algorithm)

        if db.is_known(conn, file_hash):
            db.increment_seen(conn, file_hash)
            duplicates += 1
            log.debug("Already known: %s", fpath.name)
        else:
            if do_import:
                rel = fpath.relative_to(scan_root)
                dst = import_dir / rel
                transfer_file(fpath, dst, move_mode)
                log.info("Imported: %s", rel)
            db.record_file(conn, file_hash, str(fpath), "scan" if do_import else "bulk")
            new += 1

    conn.close()

    print(f"\nScan complete: {path}")
    print(f"  Total files:    {total}")
    print(f"  New:            {new}")
    print(f"  Already known:  {duplicates}")
    if do_import:
        print(f"  Imported to:    {import_dir}")
    else:
        print("  (record-only mode — no files moved)")


# --- CLI ---

def main():
    parser = argparse.ArgumentParser(description="StashHandler")
    sub = parser.add_subparsers(dest="command")

    sub.add_parser("daemon", help="Poll qBittorrent and import new files")

    scan_p = sub.add_parser("scan", help="Bulk-scan files and record/import them")
    scan_p.add_argument("path", nargs="?", default="/staging", help="Directory to scan")
    scan_p.add_argument("--import", dest="do_import", action="store_true",
                        help="Move new files to import dir (default: record-only)")

    args = parser.parse_args()
    cfg = load_config()
    setup_logging(cfg["log_level"])

    if args.command == "scan":
        scan(cfg, args.path, args.do_import)
    else:
        daemon(cfg)


if __name__ == "__main__":
    main()