path: root/cogs/utils/file_search.py

# This file is part of Senator Sperma
# Licensed under the GPL-3.0 License
# Copyright (C) 2026 Salmonik @ salmonik@tuta.io
# https://codeberg.org/salmonik/senator-sperma

import json
import logging
import re
from pathlib import Path
import aiofiles
from rapidfuzz import fuzz
import aiohttp
import os
from dotenv import load_dotenv

load_dotenv()

logger = logging.getLogger(__name__)

class SearchUtils:
    """Utilities for searching and scoring audio files."""
    
    @staticmethod
    def normalize_name(p: Path, keep_ext=False):
        """Preprocess filenames: lower, replace separators, strip extensions optionally."""
        name = str(p.name if keep_ext else p.stem)
        return name.lower().replace('_', ' ').replace('-', ' ')
    
    @staticmethod
    def length_penalty(query, candidate):
        """Penalize tiny filenames when query is long; returns multiplier [0.0,1.0]."""
        ql = len(query)
        cl = len(candidate)
        if ql == 0: 
            return 1.0
        ratio = cl / ql
        # Prefer candidate lengths near query length; gentle falloff otherwise
        return math.exp(-0.5 * ((math.log(max(ratio, 1e-6)))**2))
    
    @staticmethod
    def combined_score(query, candidate):
        """Calculate a combined fuzzy matching score using multiple strategies.
        
        Args:
            query: The search query
            candidate: The candidate string to score against
            
        Returns:
            float: Score between 0-100, where higher means better match
        """
        if not query or not candidate:
            return 0.0
            
        query = query.lower().strip()
        candidate = candidate.lower()
        
        # Check for direct matches first
        direct_match_score = SearchUtils._check_direct_matches(query, candidate)
        if direct_match_score is not None:
            return direct_match_score
        
        # Handle short queries specially
        short_query_score = SearchUtils._handle_short_queries(query, candidate)
        if short_query_score is not None:
            return short_query_score
        
        # For longer queries, use complex scoring
        return SearchUtils._calculate_complex_score(query, candidate)
    
    @staticmethod
    def _check_direct_matches(query: str, candidate: str) -> float:
        """Check for direct substring matches."""
        if query in candidate:
            base_score = 100
            if re.search(rf'\b{re.escape(query)}\b', candidate):
                base_score = 110
            return min(100, base_score)
        return 0
    
    @staticmethod
    def _handle_short_queries(query: str, candidate: str) -> float:
        """Handle scoring for short queries (2-3 chars)."""
        if len(query) <= 3:
            if re.search(rf'\b{re.escape(query)}', candidate):
                return 90
            if query in candidate:
                return 80
        return 0
    
    @staticmethod
    def _calculate_complex_score(query: str, candidate: str) -> float:
        """Calculate complex score for longer queries."""
        query_words = {word for word in re.split(r'\W+', query) if len(word) > 1}
        if not query_words:
            return 0.0
        
        candidate_words = set(re.split(r'\W+', candidate))
        word_ratio = SearchUtils._calculate_word_ratio(query_words, candidate_words)
        
        if word_ratio == 0:
            return 0.0
        
        # Calculate various similarity metrics
        similarity_scores = SearchUtils._get_similarity_scores(query, candidate)
        word_scores = SearchUtils._get_word_scores(query, candidate, query_words, candidate_words)
        
        # Calculate final weighted score
        final_score = SearchUtils._calculate_weighted_score(similarity_scores, word_scores, word_ratio)
        
        # Apply penalties and boosts
        final_score = SearchUtils._apply_adjustments(final_score, query, candidate, word_ratio)
        
        return min(100, max(0, int(final_score)))
    
    @staticmethod
    def _calculate_word_ratio(query_words: set, candidate_words: set) -> float:
        """Calculate the ratio of matching words."""
        matching_words = sum(1 for word in query_words if word in candidate_words)
        return matching_words / len(query_words)
    
    @staticmethod
    def _get_similarity_scores(query: str, candidate: str) -> dict:
        """Get various fuzzy similarity scores."""
        return {
            'ratio': fuzz.ratio(query, candidate),
            'partial': fuzz.partial_ratio(query, candidate),
            'token_set': fuzz.token_set_ratio(query, candidate)
        }
    
    @staticmethod
    def _get_word_scores(query: str, candidate: str, query_words: set, candidate_words: set) -> dict:
        """Calculate word-based scores."""
        query_tokens = set(query.split())
        candidate_tokens = set(candidate.split())
        
        word_intersection = len(query_tokens & candidate_tokens)
        word_union = len(query_tokens | candidate_tokens)
        jaccard_sim = word_intersection / word_union if word_union > 0 else 0
        
        prefix_bonus = SearchUtils._calculate_prefix_bonus(query_words, candidate_words)
        
        return {
            'jaccard': jaccard_sim,
            'prefix': min(prefix_bonus, 100)
        }
    
    @staticmethod
    def _calculate_prefix_bonus(query_words: set, candidate_words: set) -> float:
        """Calculate prefix matching bonus."""
        prefix_bonus = 0
        for q_word in query_words:
            for c_word in candidate_words:
                if c_word.startswith(q_word):
                    prefix_bonus += (len(q_word) / max(len(c_word), 1)) * 20
        return prefix_bonus
    
    @staticmethod
    def _calculate_weighted_score(similarity_scores: dict, word_scores: dict, word_ratio: float) -> float:
        """Calculate the final weighted score."""
        return (
            0.40 * similarity_scores['token_set'] +
            0.25 * similarity_scores['ratio'] +
            0.15 * similarity_scores['partial'] +
            0.15 * (word_scores['jaccard'] * 100) +
            0.05 * word_scores['prefix']
        )
    
    @staticmethod
    def _apply_adjustments(score: float, query: str, candidate: str, word_ratio: float) -> float:
        """Apply length penalty and word ratio boost."""
        import math
        
        # Apply length penalty
        score *= min(1.0, 0.9 + (0.1 * SearchUtils.length_penalty(query, candidate)))
        
        # Apply word ratio boost
        score = score * (0.7 + (0.3 * word_ratio))
        
        return score

class UniversalSearcher:
    """Handles searching for files in configured paths with support for multiple file types."""
    
    FILE_TYPE_EXTENSIONS = {
        'Audio': {'.mp4', '.mp3', '.wav', '.opus', '.m4a', '.flac', '.aac', '.ogg'},
        'Video': {'.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.m4v', '.3gp', '.gif'},
        'Image': {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp', '.svg', '.ico'}
    }
    
    @staticmethod
    def _is_valid_file(file_path: Path, filetype: str = None) -> bool:
        """Check if file is valid for the specified filetype."""
        if not file_path.is_file():
            return False
            
        if filetype is None:
            # Search all file types
            all_extensions = set()
            for extensions in UniversalSearcher.FILE_TYPE_EXTENSIONS.values():
                all_extensions.update(extensions)
            return file_path.suffix.lower() in all_extensions
        
        filetype = filetype.capitalize()
        if filetype not in UniversalSearcher.FILE_TYPE_EXTENSIONS:
            return False
            
        valid_extensions = UniversalSearcher.FILE_TYPE_EXTENSIONS[filetype]
        return file_path.suffix.lower() in valid_extensions
    
    @staticmethod
    def _calculate_match_score(query: str, file_name: str, autocomplete: bool) -> tuple[int, str]:
        """Calculate match score for a file."""
        if not query and autocomplete:
            return 0, file_name
        
        if query and len(query) >= 2:
            if query in file_name:
                logger.debug(f"Found direct match: {file_name} (contains '{query}')")
            
            score = SearchUtils.combined_score(query, file_name)
            logger.debug(f"File: {file_name} - Score: {score}")
            
            if score >= 30:
                return score, file_name
        
        return -1, file_name
    
    @staticmethod
    async def _search_single_path(path: Path, query: str, max_results: int, autocomplete: bool, filetype: str = None) -> list[tuple[int, str]]:
        """Search for files in a single path."""
        matches = []
        
        logger.debug(f"Searching in directory: {path} (filetype: {filetype or 'All'})")
        
        for file_path in path.rglob('*'):
            if not UniversalSearcher._is_valid_file(file_path, filetype):
                continue
            
            file_name = file_path.name.lower()
            score, _ = UniversalSearcher._calculate_match_score(query, file_name, autocomplete)
            
            if score >= 0:
                matches.append((score, str(file_path.absolute())))
            
            if len(matches) >= max_results * 5:
                break
        
        return matches
    
    @staticmethod
    def _get_search_paths(data: dict, filetype: str = None) -> list:
        """Get the search paths from settings."""
        return data.get('file_paths', [])
    
    @classmethod
    async def _search_paths(cls, search_paths: list, query: str, max_results: int, autocomplete: bool, filetype: str = None) -> tuple[list, int]:
        """Search through the provided paths and return matches and file count."""
        matches = []
        file_count = 0
        
        for path_str in search_paths:
            path = Path(path_str)
            if not path.exists():
                logger.warning(f"Path does not exist: {path}")
                continue
            
            path_matches = await cls._search_single_path(path, query, max_results, autocomplete, filetype)
            matches.extend(path_matches)
            file_count += len(path_matches)
            
            if len(matches) >= max_results * 5:
                break
        
        return matches, file_count
    
    @classmethod
    async def search_local(cls, query: str = None, max_results: int = 25, autocomplete: bool = False, filetype: str = None):
        """Search for files with fuzzy matching and optional autocomplete support.
        
        Args:
            query: Search query string (can be None for autocomplete)
            max_results: Maximum number of results to return
            autocomplete: If True, returns results even with empty query
            filetype: Type of files to search for ('Audio', 'Video', 'Image', or None for all)
            
        Returns:
            List of matching file paths with scores
        """
        if not query and not autocomplete:
            return []
            
        query = (query or '').lower().strip()
        logger.info(f"Starting local search for query: '{query}' (filetype: {filetype or 'All'})")
        
        try:
            async with aiofiles.open('settings.json') as f:
                content = await f.read()
                data = json.loads(content)
                
                search_paths = cls._get_search_paths(data, filetype)
                logger.info(f"Searching in paths: {search_paths}")
                
                matches, file_count = await cls._search_paths(search_paths, query, max_results, autocomplete, filetype)
                
                logger.info(f"Scanned {file_count} files total")
                
                matches.sort(reverse=True, key=lambda x: x[0])
                result = [path for score, path in matches[:max_results]]
                
                logger.info(f"Search complete. Found {len(result)} matches out of {file_count} files.")
                if not result and query:
                    logger.warning(f"No matches found for query: '{query}'. Try a different search term.")
                return result
                
        except Exception as e:
            logger.error(f"Error during local search: {str(e)}", exc_info=True)
            return []
    
    @classmethod
    async def search_local_autocomplete(cls, query: str = None, max_results: int = 10, filetype: str = None):
        """Wrapper for autocomplete searches that returns more results and allows empty query.
        
        Args:
            query: Search query string (can be None for autocomplete)
            max_results: Maximum number of results to return
            filetype: Type of files to search for ('Audio', 'Video', 'Image', or None for all)
        """
        # For autocomplete, we want to show more options and allow empty queries
        if not query:
            return []
        
        return await cls.search_local(query, max_results=max_results, autocomplete=True, filetype=filetype)
    
    @staticmethod
    async def _load_baseurl():
        """Load baseurl from settings.json"""
        try:
            async with aiofiles.open('settings.json') as f:
                content = await f.read()
                data = json.loads(content)
                return data.get('baseurl', 'http://localhost:8080').rstrip('/')
        except Exception as e:
            logger.error(f"Error loading baseurl: {e}")
            return 'http://localhost:8080'
    
    @staticmethod
    async def search_copyparty(query: str, max_results: int = 25):
        """Search for files using copyparty API (jPOST method).
        
        Args:
            query: Search query string
            max_results: Maximum number of results to return
            
        Returns:
            List of matching file URLs
        """


        if not query:
            return []
            
        query = query.strip()
        logger.info(f"Starting copyparty search for query: '{query}'")
        query = f'name like "*{query}*"'
        
        try:
            baseurl = await UniversalSearcher._load_baseurl()
            password = os.getenv("COPYPARTY_PASSWORD")
            
            async with aiohttp.ClientSession() as session:
                headers = {"Content-Type": "application/json"}
                if password:
                    headers["PW"] = password
                async with session.post(f"{baseurl}/?srch", headers=headers, data=json.dumps({"n": max_results, "q": query})) as response:
                    if response.status == 200:
                        data = await response.json()
                        logger.info(f"Copyparty response: {data}")
                        files = UniversalSearcher._extract_files_from_response(data)
                        results = UniversalSearcher._convert_files_to_urls(files, baseurl, max_results)
                        
                        logger.info(f"Copyparty search complete. Found {len(results)} matches.")
                        return results
                    else:
                        logger.error(f"Copyparty API returned status {response.status}")
                        return []
                        
        except aiohttp.ClientError as e:
            logger.error(f"HTTP error during copyparty search: {e}")
            return []
        except Exception as e:
            logger.error(f"Error during copyparty search: {str(e)}", exc_info=True)
            return []
    
    @staticmethod
    def _extract_files_from_response(data):
        """Extract files list from copyparty API response."""
        if isinstance(data, dict):
            if 'hits' in data:
                return data['hits']
            elif 'files' in data:
                return data['files']
            elif 'results' in data:
                return data['results']
            else:
                return data if isinstance(data, list) else []
        elif isinstance(data, list):
            return data
        else:
            return []
    
    @staticmethod
    def _convert_files_to_urls(files, baseurl, max_results):
        """Convert file info objects to URLs."""
        results = []
        
        for file_info in files[:max_results]:
            if isinstance(file_info, dict):
                file_url = UniversalSearcher._extract_url_from_dict(file_info, baseurl)
                if file_url:
                    results.append(file_url)
            elif isinstance(file_info, str):
                # If it's just a string path/URL
                if file_info.startswith('http'):
                    results.append(file_info)
                else:
                    results.append(f"{baseurl}/{file_info.lstrip('/')}")
        
        return results
    
    @staticmethod
    def _extract_url_from_dict(file_info, baseurl):
        """Extract URL from file info dictionary."""
        if 'url' in file_info:
            return file_info['url']
        elif 'rp' in file_info:
            return f"{baseurl}/{file_info['rp']}"
        elif 'path' in file_info:
            return f"{baseurl}/{file_info['path'].lstrip('/')}"
        elif 'name' in file_info:
            return f"{baseurl}/{file_info['name'].lstrip('/')}"
        else:
            return None
    
    @staticmethod
    async def search_copyparty_autocomplete(query: str, max_results: int = 10):
        """Wrapper for autocomplete searches using copyparty.
        
        Args:
            query: Search query string
            max_results: Maximum number of results to return
            
        Returns:
            List of matching file URLs
        """
        if not query:
            return []
        
        return await UniversalSearcher.search_copyparty(query, max_results)