Medios-Macina/cmdlets/add_relationship.py

"""Add file relationships in Hydrus based on relationship tags in sidecar."""

from __future__ import annotations

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

from helper.logger import log

from . import register
import models
import pipeline as ctx
from helper import hydrus as hydrus_wrapper
from ._shared import Cmdlet, CmdletArg, parse_cmdlet_args, normalize_result_input, should_show_help, get_field
from helper.folder_store import read_sidecar, find_sidecar


CMDLET = Cmdlet(
    name="add-relationship",
    summary="Associate file relationships (king/alt/related) in Hydrus based on relationship tags in sidecar.",
    usage="@1-3 | add-relationship -king @4    OR    add-relationship -path <file>    OR    @1,@2,@3 | add-relationship",
    arg=[
        CmdletArg("path", type="string", description="Specify the local file path (if not piping a result)."),
        CmdletArg("-king", type="string", description="Explicitly set the king hash/file for relationships (e.g., -king @4 or -king hash)"),
        CmdletArg("-type", type="string", description="Relationship type for piped items (default: 'alt', options: 'king', 'alt', 'related')"),
    ],
    detail=[
        "- Mode 1: Pipe multiple items, first becomes king, rest become alts (default)",
        "- Mode 2: Use -king to explicitly set which item/hash is the king: @1-3 | add-relationship -king @4",
        "- Mode 3: Read relationships from sidecar (format: 'relationship: hash(king)<HASH>,hash(alt)<HASH>...')",
        "- Supports three relationship types: king (primary), alt (alternative), related (other versions)",
        "- When using -king, all piped items become the specified relationship type to the king",
    ],
)


def _normalise_hash_hex(value: Optional[str]) -> Optional[str]:
    """Normalize a hash hex string to lowercase 64-char format."""
    if not value or not isinstance(value, str):
        return None
    normalized = value.strip().lower()
    if len(normalized) == 64 and all(c in '0123456789abcdef' for c in normalized):
        return normalized
    return None


def _extract_relationships_from_tag(tag_value: str) -> Dict[str, list[str]]:
    """Parse relationship tag like 'relationship: hash(king)<HASH>,hash(alt)<HASH>'.
    
    Returns a dict like {"king": ["HASH1"], "alt": ["HASH2"], ...}
    """
    result: Dict[str, list[str]] = {}
    if not isinstance(tag_value, str):
        return result
    
    # Match patterns like hash(king)HASH or hash(type)HASH (no angle brackets)
    pattern = r'hash\((\w+)\)([a-fA-F0-9]{64})'
    matches = re.findall(pattern, tag_value)
    
    for rel_type, hash_value in matches:
        normalized = _normalise_hash_hex(hash_value)
        if normalized:
            if rel_type not in result:
                result[rel_type] = []
            result[rel_type].append(normalized)
    
    return result


def _resolve_king_reference(king_arg: str) -> Optional[str]:
    """Resolve a king reference like '@4' to its actual hash or path.
    
    Supports:
    - Direct hash: '0123456789abcdef...' (64 chars)
    - Selection reference: '@4' (resolves from pipeline context)
    
    Returns:
    - For Hydrus items: normalized hash
    - For local storage items: file path
    - None if not found
    """
    if not king_arg:
        return None
    
    # Check if it's already a valid hash
    normalized = _normalise_hash_hex(king_arg)
    if normalized:
        return normalized
    
    # Try to resolve as @N selection from pipeline context
    if king_arg.startswith('@'):
        try:
            # Get the result items from the pipeline context
            from pipeline import get_last_result_items
            items = get_last_result_items()
            if not items:
                log(f"Cannot resolve {king_arg}: no search results in context", file=sys.stderr)
                return None
            
            # Parse @N to get the index (1-based)
            index_str = king_arg[1:]  # Remove '@'
            index = int(index_str) - 1  # Convert to 0-based
            
            if 0 <= index < len(items):
                item = items[index]
                
                # Try to extract hash from the item (could be dict or object)
                item_hash = (
                    get_field(item, 'hash_hex')
                    or get_field(item, 'hash')
                    or get_field(item, 'file_hash')
                )
                
                if item_hash:
                    normalized = _normalise_hash_hex(item_hash)
                    if normalized:
                        return normalized
                
                # If no hash, try to get file path (for local storage)
                file_path = (
                    get_field(item, 'file_path')
                    or get_field(item, 'path')
                    or get_field(item, 'target')
                )
                
                if file_path:
                    return str(file_path)
                
                log(f"Item {king_arg} has no hash or path information", file=sys.stderr)
                return None
            else:
                log(f"Index {king_arg} out of range", file=sys.stderr)
                return None
        except (ValueError, IndexError) as e:
            log(f"Cannot resolve {king_arg}: {e}", file=sys.stderr)
            return None
    
    return None


def _refresh_relationship_view_if_current(target_hash: Optional[str], target_path: Optional[str], other: Optional[str], config: Dict[str, Any]) -> None:
    """If the current subject matches the target, refresh relationships via get-relationship."""
    try:
        from cmdlets import get_relationship as get_rel_cmd  # type: ignore
    except Exception:
        return

    try:
        subject = ctx.get_last_result_subject()
        if subject is None:
            return

        def norm(val: Any) -> str:
            return str(val).lower()

        target_hashes = [norm(v) for v in [target_hash, other] if v]
        target_paths = [norm(v) for v in [target_path, other] if v]

        subj_hashes: list[str] = []
        subj_paths: list[str] = []
        if isinstance(subject, dict):
            subj_hashes = [norm(v) for v in [subject.get("hydrus_hash"), subject.get("hash"), subject.get("hash_hex"), subject.get("file_hash")] if v]
            subj_paths = [norm(v) for v in [subject.get("file_path"), subject.get("path"), subject.get("target")] if v]
        else:
            subj_hashes = [norm(getattr(subject, f, None)) for f in ("hydrus_hash", "hash", "hash_hex", "file_hash") if getattr(subject, f, None)]
            subj_paths = [norm(getattr(subject, f, None)) for f in ("file_path", "path", "target") if getattr(subject, f, None)]

        is_match = False
        if target_hashes and any(h in subj_hashes for h in target_hashes):
            is_match = True
        if target_paths and any(p in subj_paths for p in target_paths):
            is_match = True
        if not is_match:
            return

        refresh_args: list[str] = []
        if target_hash:
            refresh_args.extend(["-hash", target_hash])
        get_rel_cmd._run(subject, refresh_args, config)
    except Exception:
        pass


@register(["add-relationship", "add-rel"])  # primary name and alias
def _run(result: Any, _args: Sequence[str], config: Dict[str, Any]) -> int:
    """Associate file relationships in Hydrus.
    
    Two modes of operation:
    1. Read from sidecar: Looks for relationship tags in the file's sidecar (format: "relationship: hash(king)<HASH>,hash(alt)<HASH>")
    2. Pipeline mode: When piping multiple results, the first becomes "king" and subsequent items become "alt"

    Returns 0 on success, non-zero on failure.
    """
    # Help
    if should_show_help(_args):
        log(json.dumps(CMDLET, ensure_ascii=False, indent=2))
        return 0
    
    # Parse arguments using CMDLET spec
    parsed = parse_cmdlet_args(_args, CMDLET)
    arg_path: Optional[Path] = None
    king_arg = parsed.get("king")  # New: explicit king argument
    rel_type = parsed.get("type", "alt")  # New: relationship type (default: alt)
    
    if parsed:
        # Get the first arg value (e.g., -path)
        first_arg_name = CMDLET.get("args", [{}])[0].get("name") if CMDLET.get("args") else None
        if first_arg_name and first_arg_name in parsed:
            arg_value = parsed[first_arg_name]
            try:
                arg_path = Path(str(arg_value)).expanduser()
            except Exception:
                arg_path = Path(str(arg_value))

    # Handle @N selection which creates a list
    # Use normalize_result_input to handle both single items and lists
    items_to_process = normalize_result_input(result)
    
    if not items_to_process and not arg_path:
        log("No items provided to add-relationship (no piped result and no -path)", file=sys.stderr)
        return 1
    
    # If no items from pipeline, just process the -path arg
    if not items_to_process and arg_path:
        items_to_process = [{"file_path": arg_path}]

    # Import local storage utilities
    from helper.folder_store import LocalLibrarySearchOptimizer
    from config import get_local_storage_path
    
    local_storage_path = get_local_storage_path(config) if config else None
    
    # Check if any items have Hydrus hashes (file_hash or hash_hex fields)
    has_hydrus_hashes = any(
        (isinstance(item, dict) and (item.get('hash_hex') or item.get('hash')))
        or (hasattr(item, 'hash_hex') or hasattr(item, 'hash'))
        for item in items_to_process
    )
    
    # Only try to initialize Hydrus if we actually have Hydrus hashes to work with
    hydrus_client = None
    if has_hydrus_hashes:
        try:
            hydrus_client = hydrus_wrapper.get_client(config)
        except Exception as exc:
            log(f"Hydrus unavailable, will use local storage: {exc}", file=sys.stderr)
    
    # Use local storage if it's available and either Hydrus is not available or items are local files
    use_local_storage = local_storage_path and (not has_hydrus_hashes or (arg_path and arg_path.exists()))

    # Resolve the king reference once (if provided)
    king_hash = None
    if king_arg:
        # Resolve the king reference (could be @4 or a direct hash)
        king_hash = _resolve_king_reference(king_arg)
        if not king_hash:
            log(f"Failed to resolve king argument: {king_arg}", file=sys.stderr)
            return 1
        log(f"Using king hash: {king_hash}", file=sys.stderr)

    # Process each item in the list
    for item_idx, item in enumerate(items_to_process):
        # Extract hash and path from current item
        file_hash = None
        file_path_from_result = None
        
        if isinstance(item, dict):
            file_hash = item.get("hash_hex") or item.get("hash")
            file_path_from_result = item.get("file_path") or item.get("path") or item.get("target")
        else:
            file_hash = getattr(item, "hash_hex", None) or getattr(item, "hash", None)
            file_path_from_result = getattr(item, "file_path", None) or getattr(item, "path", None)
        
        # PIPELINE MODE with Hydrus: Track relationships using hash
        if file_hash and hydrus_client:
            file_hash = _normalise_hash_hex(file_hash)
            if not file_hash:
                log("Invalid file hash format", file=sys.stderr)
                return 1
            
            # If explicit -king provided, use it
            if king_hash:
                try:
                    hydrus_client.set_relationship(file_hash, king_hash, rel_type)
                    log(
                        f"[add-relationship] Set {rel_type} relationship: {file_hash} <-> {king_hash}",
                        file=sys.stderr
                    )
                    _refresh_relationship_view_if_current(file_hash, file_path_from_result, king_hash, config)
                except Exception as exc:
                    log(f"Failed to set relationship: {exc}", file=sys.stderr)
                    return 1
            else:
                # Original behavior: no explicit king, first becomes king, rest become alts
                try:
                    existing_king = ctx.load_value("relationship_king")
                except Exception:
                    existing_king = None
                
                # If this is the first item, make it the king
                if not existing_king:
                    try:
                        ctx.store_value("relationship_king", file_hash)
                        log(f"Established king hash: {file_hash}", file=sys.stderr)
                        continue  # Move to next item
                    except Exception:
                        pass
                
                # If we already have a king and this is a different hash, link them
                if existing_king and existing_king != file_hash:
                    try:
                        hydrus_client.set_relationship(file_hash, existing_king, rel_type)
                        log(
                            f"[add-relationship] Set {rel_type} relationship: {file_hash} <-> {existing_king}",
                            file=sys.stderr
                        )
                        _refresh_relationship_view_if_current(file_hash, file_path_from_result, existing_king, config)
                    except Exception as exc:
                        log(f"Failed to set relationship: {exc}", file=sys.stderr)
                        return 1
        
        # LOCAL STORAGE MODE: Handle relationships for local files
        elif use_local_storage and file_path_from_result:
            try:
                file_path_obj = Path(str(file_path_from_result))
                
                if not file_path_obj.exists():
                    log(f"File not found: {file_path_obj}", file=sys.stderr)
                    return 1
                
                if king_hash:
                    # king_hash is a file path from _resolve_king_reference (or a Hydrus hash)
                    king_file_path = Path(str(king_hash)) if king_hash else None
                    if king_file_path and king_file_path.exists():
                        with LocalLibrarySearchOptimizer(local_storage_path) as db:
                            db.set_relationship(file_path_obj, king_file_path, rel_type)
                        log(f"Set {rel_type} relationship: {file_path_obj.name} -> {king_file_path.name}", file=sys.stderr)
                        _refresh_relationship_view_if_current(None, str(file_path_obj), str(king_file_path), config)
                    else:
                        log(f"King file not found or invalid: {king_hash}", file=sys.stderr)
                        return 1
                else:
                    # Original behavior: first becomes king, rest become alts
                    try:
                        king_path = ctx.load_value("relationship_king_path")
                    except Exception:
                        king_path = None
                    
                    if not king_path:
                        try:
                            ctx.store_value("relationship_king_path", str(file_path_obj))
                            log(f"Established king file: {file_path_obj.name}", file=sys.stderr)
                            continue  # Move to next item
                        except Exception:
                            pass
                    
                    if king_path and king_path != str(file_path_obj):
                        try:
                            with LocalLibrarySearchOptimizer(local_storage_path) as db:
                                db.set_relationship(file_path_obj, Path(king_path), rel_type)
                            log(f"Set {rel_type} relationship: {file_path_obj.name} -> {Path(king_path).name}", file=sys.stderr)
                            _refresh_relationship_view_if_current(None, str(file_path_obj), str(king_path), config)
                        except Exception as exc:
                            log(f"Failed to set relationship: {exc}", file=sys.stderr)
                            return 1
            except Exception as exc:
                log(f"Local storage error: {exc}", file=sys.stderr)
                return 1
    
    return 0

    # FILE MODE: Read relationships from sidecar (legacy mode - for -path arg only)
    log("Note: Use piping mode for easier relationships. Example: 1,2,3 | add-relationship", file=sys.stderr)
    
    # Resolve media path from -path arg or result target
    target = getattr(result, "target", None) or getattr(result, "path", None)
    media_path = arg_path if arg_path is not None else Path(str(target)) if isinstance(target, str) else None
    if media_path is None:
        log("Provide -path <file> or pipe a local file result", file=sys.stderr)
        return 1
    
    # Validate local file
    if str(media_path).lower().startswith(("http://", "https://")):
        log("This cmdlet requires a local file path, not a URL", file=sys.stderr)
        return 1
    if not media_path.exists() or not media_path.is_file():
        log(f"File not found: {media_path}", file=sys.stderr)
        return 1

    # Build Hydrus client
    try:
        client = hydrus_wrapper.get_client(config)
    except Exception as exc:
        log(f"Hydrus client unavailable: {exc}", file=sys.stderr)
        return 1
    
    if client is None:
        log("Hydrus client unavailable", file=sys.stderr)
        return 1

    # Read sidecar to find relationship tags
    sidecar_path = find_sidecar(media_path)
    if sidecar_path is None:
        log(f"No sidecar found for {media_path.name}", file=sys.stderr)
        return 1
    
    try:
        _, tags, _ = read_sidecar(sidecar_path)
    except Exception as exc:
        log(f"Failed to read sidecar: {exc}", file=sys.stderr)
        return 1

    # Find relationship tags (format: "relationship: hash(king)<HASH>,hash(alt)<HASH>,hash(related)<HASH>")
    relationship_tags = [t for t in tags if isinstance(t, str) and t.lower().startswith("relationship:")]
    
    if not relationship_tags:
        log(f"No relationship tags found in sidecar", file=sys.stderr)
        return 0  # Not an error, just nothing to do
    
    # Get the file hash from result (should have been set by add-file)
    file_hash = getattr(result, "hash_hex", None)
    if not file_hash:
        log("File hash not available (run add-file first)", file=sys.stderr)
        return 1
    
    file_hash = _normalise_hash_hex(file_hash)
    if not file_hash:
        log("Invalid file hash format", file=sys.stderr)
        return 1
    
    # Parse relationships from tags and apply them
    success_count = 0
    error_count = 0
    
    for rel_tag in relationship_tags:
        try:
            # Parse: "relationship: hash(king)<HASH>,hash(alt)<HASH>,hash(related)<HASH>"
            rel_str = rel_tag.split(":", 1)[1].strip()  # Get part after "relationship:"
            
            # Parse relationships
            rels = _extract_relationships_from_tag(f"relationship: {rel_str}")
            
            # Set the relationships in Hydrus
            for rel_type, related_hashes in rels.items():
                if not related_hashes:
                    continue
                
                for related_hash in related_hashes:
                    # Don't set relationship between hash and itself
                    if file_hash == related_hash:
                        continue
                    
                    try:
                        hydrus_client.set_relationship(file_hash, related_hash, rel_type)
                        log(
                            f"[add-relationship] Set {rel_type} relationship: "
                            f"{file_hash} <-> {related_hash}",
                            file=sys.stderr
                        )
                        success_count += 1
                    except Exception as exc:
                        log(f"Failed to set {rel_type} relationship: {exc}", file=sys.stderr)
                        error_count += 1
        
        except Exception as exc:
            log(f"Failed to parse relationship tag: {exc}", file=sys.stderr)
            error_count += 1
    
    if success_count > 0:
        log(f"Successfully set {success_count} relationship(s) for {media_path.name}", file=sys.stderr)
        ctx.emit(f"add-relationship: {media_path.name} ({success_count} relationships set)")
        return 0
    elif error_count == 0:
        log(f"No relationships to set", file=sys.stderr)
        return 0  # Success with nothing to do
    else:
        log(f"Failed with {error_count} error(s)", file=sys.stderr)
        return 1