Medios-Macina/cmdlets/add_url.py

from __future__ import annotations

from typing import Any, Dict, Sequence
import json
import sys
from pathlib import Path

from . import register
import models
import pipeline as ctx
from helper import hydrus as hydrus_wrapper
from ._shared import Cmdlet, CmdletArg, normalize_hash
from helper.logger import log
from config import get_local_storage_path
from helper.local_library import LocalLibraryDB

CMDLET = Cmdlet(
    name="add-url",
    summary="Associate a URL with a file (Hydrus or Local).",
    usage="add-url [-hash <sha256>] <url>",
    args=[
        CmdletArg("-hash", description="Override the Hydrus file hash (SHA256) to target instead of the selected result."),
        CmdletArg("url", required=True, description="The URL to associate with the file."),
    ],
    details=[
        "- Adds the URL to the file's known URL list.",
    ],
)


@register(["add-url", "ass-url", "associate-url", "add_url"])  # aliases
def add(result: Any, args: Sequence[str], config: Dict[str, Any]) -> int:
    # Help
    try:
        if any(str(a).lower() in {"-?", "/?", "--help", "-h", "help", "--cmdlet"} for a in args):
            log(json.dumps(CMDLET, ensure_ascii=False, indent=2))
            return 0
    except Exception:
        pass
        
    from ._shared import parse_cmdlet_args
    parsed = parse_cmdlet_args(args, CMDLET)
    override_hash = parsed.get("hash")
    url_arg = parsed.get("url")
    
    if not url_arg:
        log("Requires a URL argument")
        return 1
    
    url_arg = str(url_arg).strip()
    if not url_arg:
        log("Requires a non-empty URL")
        return 1
        
    # Split by comma to handle multiple URLs
    urls_to_add = [u.strip() for u in url_arg.split(',') if u.strip()]
    
    # Handle @N selection which creates a list - extract the first item
    if isinstance(result, list) and len(result) > 0:
        result = result[0]
    
    # Helper to get field from both dict and object
    def get_field(obj: Any, field: str, default: Any = None) -> Any:
        if isinstance(obj, dict):
            return obj.get(field, default)
        else:
            return getattr(obj, field, default)

    success = False
    
    # 1. Try Local Library
    file_path = get_field(result, "file_path") or get_field(result, "path")
    if file_path and not override_hash:
        try:
            path_obj = Path(file_path)
            if path_obj.exists():
                storage_path = get_local_storage_path(config)
                if storage_path:
                    with LocalLibraryDB(storage_path) as db:
                        metadata = db.get_metadata(path_obj) or {}
                        known_urls = metadata.get("known_urls") or []
                        
                        local_changed = False
                        for url in urls_to_add:
                            if url not in known_urls:
                                known_urls.append(url)
                                local_changed = True
                                ctx.emit(f"Associated URL with local file {path_obj.name}: {url}")
                            else:
                                ctx.emit(f"URL already exists for local file {path_obj.name}: {url}")
                        
                        if local_changed:
                            metadata["known_urls"] = known_urls
                            # Ensure we have a hash if possible, but don't fail if not
                            if not metadata.get("hash"):
                                try:
                                    from helper.utils import sha256_file
                                    metadata["hash"] = sha256_file(path_obj)
                                except Exception:
                                    pass
                                    
                            db.save_metadata(path_obj, metadata)
                        
                        success = True
        except Exception as e:
            log(f"Error updating local library: {e}", file=sys.stderr)

    # 2. Try Hydrus
    hash_hex = normalize_hash(override_hash) if override_hash else normalize_hash(get_field(result, "hash_hex", None))
    
    if hash_hex:
        try:
            client = hydrus_wrapper.get_client(config)
            if client:
                for url in urls_to_add:
                    client.associate_url(hash_hex, url)
                    preview = hash_hex[:12] + ('…' if len(hash_hex) > 12 else '')
                    ctx.emit(f"Associated URL with Hydrus file {preview}: {url}")
                success = True
        except Exception as exc:
            # Only log error if we didn't succeed locally either
            if not success:
                log(f"Hydrus add-url failed: {exc}", file=sys.stderr)
                return 1

    if success:
        return 0
        
    if not hash_hex and not file_path:
        log("Selected result does not include a file path or Hydrus hash", file=sys.stderr)
        return 1
        
    return 1