mfinder/db/files_sql.py

import threading
import time
import datetime
import re
from sqlalchemy import create_engine, or_, func, and_, text
from sqlalchemy import Column, TEXT, Numeric, DateTime
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy.pool import QueuePool
from sqlalchemy.exc import OperationalError, PendingRollbackError, NoResultFound
from mfinder import DB_URL, LOGGER
from mfinder.utils.helpers import unpack_new_file_id
import asyncio

BASE = declarative_base()

class Files(BASE):
    __tablename__ = "files"
    file_name = Column(TEXT, primary_key=True)
    file_id = Column(TEXT)
    file_ref = Column(TEXT)
    file_size = Column(Numeric)
    file_type = Column(TEXT)
    mime_type = Column(TEXT)
    caption = Column(TEXT)
    created_at = Column(DateTime, server_default=func.now())

    def __init__(
        self, file_name, file_id, file_ref, file_size, file_type, mime_type, caption, created_at=None
    ):
        self.file_name = file_name
        self.file_id = file_id
        self.file_ref = file_ref
        self.file_size = file_size
        self.file_type = file_type
        self.mime_type = mime_type
        self.caption = caption
        self.created_at = created_at or datetime.datetime.utcnow()

class SearchSuggestionsCache(BASE):
    __tablename__ = "search_suggestions_cache"
    
    original_query = Column(TEXT, primary_key=True)
    corrected_query = Column(TEXT, nullable=True)
    keywords = Column(TEXT, nullable=True)  # Comma separated
    suggestions = Column(TEXT, nullable=True)  # Comma separated
    created_at = Column(DateTime, server_default=func.now(), onupdate=func.now())

    def __init__(self, original_query, corrected_query, keywords, suggestions):
        self.original_query = original_query.strip().lower()
        self.corrected_query = corrected_query
        self.keywords = keywords
        self.suggestions = suggestions

def start() -> scoped_session:
    connect_args = {}
    if DB_URL and DB_URL.startswith(("postgres://", "postgresql://")):
        connect_args["sslmode"] = "require"
    engine = create_engine(
        DB_URL,
        connect_args=connect_args,
        poolclass=QueuePool,
        pool_size=10,
        max_overflow=20,
        pool_pre_ping=True,
        pool_recycle=1800
    )
    BASE.metadata.bind = engine
    BASE.metadata.create_all(engine)
    
    # DDL migrations: ensure created_at column on files table and pg_trgm indexes
    try:
        with engine.connect() as conn:
            conn.execute(text("ALTER TABLE files ADD COLUMN IF NOT EXISTS created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP"))
            conn.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm"))
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_files_file_name_trgm ON files USING gin (file_name gin_trgm_ops)"))
            conn.execute(text("CREATE INDEX IF NOT EXISTS idx_files_caption_trgm ON files USING gin (caption gin_trgm_ops)"))
            conn.commit()
    except Exception as e:
        LOGGER.warning(f"Failed to create GIN trigram indexes or extension: {e}")
        
    return scoped_session(sessionmaker(bind=engine, autoflush=False))

SESSION = start()
INSERTION_LOCK = threading.RLock()

import asyncio
import functools
from sqlalchemy.exc import OperationalError, PendingRollbackError

def db_retry(func):
    if asyncio.iscoroutinefunction(func):
        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            for attempt in range(3):
                try:
                    return await func(*args, **kwargs)
                except (OperationalError, PendingRollbackError) as e:
                    LOGGER.warning("Database connection error in %s (attempt %s): %s. Reconnecting...", func.__name__, attempt + 1, str(e))
                    try:
                        SESSION.rollback()
                        SESSION.remove()
                    except Exception:
                        pass
                    if attempt == 2:
                        raise
                    await asyncio.sleep(0.5)
        return wrapper
    else:
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            for attempt in range(3):
                try:
                    return func(*args, **kwargs)
                except (OperationalError, PendingRollbackError) as e:
                    LOGGER.warning("Database connection error in %s (attempt %s): %s. Reconnecting...", func.__name__, attempt + 1, str(e))
                    try:
                        SESSION.rollback()
                        SESSION.remove()
                    except Exception:
                        pass
                    if attempt == 2:
                        raise
                    import time
                    time.sleep(0.5)
        return wrapper

def reconnect_session(max_retries=5, delay=5):
    """Attempt to reconnect to the database a specified number of times with a delay."""
    for attempt in range(max_retries):
        try:
            global SESSION
            SESSION = start()
            return SESSION
        except OperationalError as e:
            LOGGER.warning(f"Database connection failed: {e}. Retrying in {delay} seconds...")
            time.sleep(delay)
    raise Exception("Failed to reconnect to the database after multiple attempts")

@db_retry
async def save_file(media):
    """Save a media file to the database."""
    file_id, file_ref = unpack_new_file_id(media.file_id)
    with INSERTION_LOCK:
        try:
            # Check if file with same file_id exists
            file = SESSION.query(Files).filter_by(file_id=file_id).one()
            LOGGER.warning("%s is already saved in the database", media.file_name)
        except NoResultFound:
            try:
                # Check if file with same file_name and file_size exists
                file = SESSION.query(Files).filter_by(file_name=media.file_name, file_size=media.file_size).one()
                LOGGER.warning("%s with size %s is already saved in the database", media.file_name, media.file_size)
            except NoResultFound:
                # Create a new file record
                file = Files(
                    file_name=media.caption if media.caption else media.file_name,
                    file_id=file_id,
                    file_ref=file_ref,
                    file_size=media.file_size,
                    file_type=media.file_type,
                    mime_type=media.mime_type,
                    caption=media.caption if media.caption else media.file_name,
                )
                LOGGER.info("%s is saved in the database", media.file_name)
                SESSION.add(file)
                SESSION.commit()
                return True
            except Exception as e:
                LOGGER.warning("Error occurred while saving file in the database: %s", str(e))
                SESSION.rollback()
                return False
        except Exception as e:
            LOGGER.warning("Error occurred while saving file in the database: %s", str(e))
            SESSION.rollback()
            return False
        finally:
            try:
                SESSION.close()
            except Exception as close_error:
                LOGGER.error(f"Error closing session: {close_error}")

async def get_filter_results(query, page=1, per_page=10):
    """Get filtered results from the database."""
    retries = 3
    while retries > 0:
        try:
            with INSERTION_LOCK:
                offset = (page - 1) * per_page
                search = query.split()
                conditions = []
                for word in search:
                    conditions.append(
                        or_(
                            Files.file_name.ilike(f"%{word}%"),
                            Files.caption.ilike(f"%{word}%"),
                        )
                    )
                combined_condition = and_(*conditions)
                files_query = (
                    SESSION.query(Files)
                    .filter(combined_condition)
                    .order_by(Files.file_name)
                )
                total_count = files_query.count()
                files = files_query.offset(offset).limit(per_page).all()
                return files, total_count
        except PendingRollbackError:
            SESSION.rollback()
            retries -= 1
            continue
        except OperationalError as e:
            LOGGER.warning(f"OperationalError: {e}. Retrying...")
            reconnect_session()
            retries -= 1
        except Exception as e:
            LOGGER.warning(f"Error occurred while retrieving filter results: {e}")
            return [], 0
        finally:
            try:
                SESSION.close()
            except Exception as close_error:
                LOGGER.error(f"Error closing session: {close_error}")
    return [], 0

async def get_precise_filter_results(query, page=1, per_page=10):
    """Get precise filtered results from the database."""
    retries = 3
    while retries > 0:
        try:
            with INSERTION_LOCK:
                offset = (page - 1) * per_page
                search = query.split()
                conditions = []
                for word in search:
                    conditions.append(
                        or_(
                            func.concat(" ", Files.file_name, " ").ilike(f"% {word} %"),
                            func.concat(" ", Files.caption, " ").ilike(f"% {word} %"),
                        )
                    )
                combined_condition = and_(*conditions)
                files_query = (
                    SESSION.query(Files)
                    .filter(combined_condition)
                    .order_by(Files.file_name)
                )
                total_count = files_query.count()
                files = files_query.offset(offset).limit(per_page).all()
                return files, total_count
        except PendingRollbackError:
            SESSION.rollback()
            retries -= 1
            continue
        except OperationalError as e:
            LOGGER.warning(f"OperationalError: {e}. Retrying...")
            reconnect_session()
            retries -= 1
        except Exception as e:
            LOGGER.warning(f"Error occurred while retrieving filter results: {e}")
            return [], 0
        finally:
            try:
                SESSION.close()
            except Exception as close_error:
                LOGGER.error(f"Error closing session: {close_error}")
    return [], 0

GENERIC_KEYWORDS_RE = re.compile(
    r'\b(1080p|720p|480p|4k|2k|hdrip|webdl|web-dl|bluray|blu-ray|hdtv|brrip|dvdrip|hevc|x264|x265|h264|h265|aac|dd5\.1|dd\+|ddp5\.1|ddp|dd|dual|multi|esub|esubs|sub|subs|tamil|telugu|hindi|malayalam|kannada|english|bengali|marathi|punjabi|movie|season|s\d+ep\d+|s\d+|ep\d+|combined)\b',
    re.IGNORECASE
)

def clean_movie_title(text: str) -> str:
    if not text:
        return ""
    # Strip common channel tags (e.g. starting with @)
    title = re.sub(r'@[a-zA-Z0-9_]+', '', text)
    # Strip leading non-alphanumeric characters
    title = re.sub(r'^[^a-zA-Z0-9]+', '', title).strip()
    
    # Try to extract the title including the year in parentheses, e.g. "Movie Name (2024)"
    year_match = re.search(r'^(.*?)\s*\((\d{4})\)', title)
    if year_match:
        prefix = year_match.group(1)
        year = year_match.group(2)
        # Clean prefix from common separators like [, but NOT mid-word hyphens/apostrophes
        # Split by [ or ' - ' (with spaces)
        prefix = re.split(r'\[|\s+-\s+', prefix)[0]
        prefix = re.sub(r'\s+', ' ', prefix).strip()
        # Remove trailing non-alphanumeric except punctuation
        prefix = re.sub(r'[^a-zA-Z0-9\-\'\s]+$', '', prefix).strip()
        return f"{prefix} ({year})"
    
    # Otherwise, split at [ or ' - ' (with spaces) or (
    title = re.split(r'\[|\s+-\s+|\(', title)[0]
    # Remove generic keywords
    title = GENERIC_KEYWORDS_RE.sub('', title)
    title = re.sub(r'\s+', ' ', title).strip()
    # Remove trailing non-alphanumeric except punctuation
    title = re.sub(r'[^a-zA-Z0-9\-\'\s]+$', '', title).strip()
    return title

async def get_trigram_filter_results(query, page=1, per_page=10, threshold=0.3):
    """Get fuzzy filtered results from the database using trigram similarity and re-ranked in Python."""
    import difflib

    retries = 3
    while retries > 0:
        try:
            with INSERTION_LOCK:
                offset = (page - 1) * per_page
                # Set dynamic trigram word similarity threshold for matching
                SESSION.execute(text(f"SET LOCAL pg_trgm.word_similarity_threshold = {threshold}"))
                
                # Fetch up to 100 candidate fuzzy matches to re-rank
                sql_query = text("""
                    SELECT file_name, file_id, file_ref, file_size, file_type, mime_type, caption
                    FROM files
                    WHERE file_name %> :query OR caption %> :query
                    LIMIT 100
                """)
                
                params = {"query": query}
                result = SESSION.execute(sql_query, params).fetchall()
                
                # Clean the query
                q_clean = clean_movie_title(query).lower()
                if not q_clean:
                    q_clean = query.lower()
                
                candidates = []
                for row in result:
                    title_name = clean_movie_title(row.file_name)
                    ratio_name = difflib.SequenceMatcher(None, q_clean, title_name.lower()).ratio()
                    
                    ratio_caption = 0.0
                    if row.caption:
                        title_caption = clean_movie_title(row.caption)
                        ratio_caption = difflib.SequenceMatcher(None, q_clean, title_caption.lower()).ratio()
                        
                    best_ratio = max(ratio_name, ratio_caption)
                    candidates.append((row, best_ratio))
                    
                # Re-rank candidates by ratio descending, then file_name
                candidates.sort(key=lambda x: (-x[1], x[0].file_name))
                
                total_count = len(candidates)
                sliced_candidates = candidates[offset : offset + per_page]
                
                files = []
                for row, ratio in sliced_candidates:
                    files.append(Files(
                        file_name=row.file_name,
                        file_id=row.file_id,
                        file_ref=row.file_ref,
                        file_size=row.file_size,
                        file_type=row.file_type,
                        mime_type=row.mime_type,
                        caption=row.caption
                    ))
                return files, total_count
        except PendingRollbackError:
            SESSION.rollback()
            retries -= 1
            continue
        except OperationalError as e:
            LOGGER.warning(f"OperationalError: {e}. Retrying...")
            reconnect_session()
            retries -= 1
        except Exception as e:
            LOGGER.warning(f"Error occurred while retrieving trigram results: {e}")
            return [], 0
        finally:
            try:
                SESSION.close()
            except Exception as close_error:
                LOGGER.error(f"Error closing session: {close_error}")
    return [], 0

async def get_file_details(file_id):
    """Get file details based on file_id."""
    retries = 3
    while retries > 0:
        try:
            with INSERTION_LOCK:
                file_details = SESSION.query(Files).filter_by(file_id=file_id).all()
                return file_details
        except PendingRollbackError:
            SESSION.rollback()
            retries -= 1
            continue
        except OperationalError as e:
            LOGGER.warning(f"OperationalError: {e}. Retrying...")
            reconnect_session()
            retries -= 1
        except Exception as e:
            LOGGER.warning(f"Error occurred while retrieving file details: {e}")
            return []
        finally:
            try:
                SESSION.close()
            except Exception as close_error:
                LOGGER.error(f"Error closing session: {close_error}")
    return []

async def delete_file(media):
    """Delete a file record from the database."""
    file_id, file_ref = unpack_new_file_id(media.file_id)
    retries = 3
    while retries > 0:
        try:
            with INSERTION_LOCK:
                file = SESSION.query(Files).filter_by(file_id=file_id).first()
                if file:
                    SESSION.delete(file)
                    SESSION.commit()
                    return True
                return "Not Found"
                LOGGER.warning("File to delete not found: %s", str(file_id))
        except PendingRollbackError:
            SESSION.rollback()
            retries -= 1
            continue
        except OperationalError as e:
            LOGGER.warning(f"OperationalError: {e}. Retrying...")
            reconnect_session()
            retries -= 1
        except Exception as e:
            LOGGER.warning(f"Error occurred while deleting file: {e}")
            SESSION.rollback()
            return False
        finally:
            try:
                SESSION.close()
            except Exception as close_error:
                LOGGER.error(f"Error closing session: {close_error}")
    return False

@db_retry
async def get_keyword_db_suggestions(query: str, limit: int = 30) -> list:
    """
    Extracts significant keywords (actor names, franchise words) from query
    and finds matching unique movie titles directly from uploaded files in DB.
    """
    import re
    STOP_WORDS = {
        'new', 'movie', 'movies', 'tamil', 'telugu', 'hindi', 'malayalam', 'kannada',
        'english', 'full', 'download', 'hd', 'mp4', 'mkv', '1080p', '720p', '480p',
        'best', 'latest', 'film', 'films', 'cinema', '2023', '2024', '2025', '2026',
        'part', 'all', 'series', 'collection', 'link', 'file', 'files', 'video'
    }

    GENERIC_RE = re.compile(
        r'\b(1080p|720p|480p|4k|2k|hdrip|webdl|web-dl|bluray|blu-ray|hdtv|brrip|dvdrip|hevc|x264|x265|h264|h265|aac|dd5\.1|dd\+|ddp5\.1|ddp|dd|dual|multi|esub|esubs|sub|subs|tamil|telugu|hindi|malayalam|kannada|english|bengali|marathi|punjabi|movie|season|s\d+ep\d+|s\d+|ep\d+|combined)\b',
        re.IGNORECASE
    )

    def _extract_clean_title(text: str) -> str:
        if not text:
            return ""
        t = text.replace("@FIREBOLTOFFICIAL", "")
        t = re.split(r'\(\d{4}\)', t)[0]
        t = t.split('[')[0]
        t = GENERIC_RE.sub('', t)
        t = re.sub(r'^[^a-zA-Z0-9]+|[^a-zA-Z0-9]+$', '', t)
        t = re.sub(r'\s+', ' ', t).strip()
        return t

    words = [w.strip().lower() for w in re.findall(r'\b\w+\b', query) if len(w) >= 3]
    keywords = [w for w in words if w not in STOP_WORDS]

    if not keywords:
        return []

    retries = 3
    while retries > 0:
        try:
            with INSERTION_LOCK:
                conditions = [Files.file_name.ilike(f"%{kw}%") for kw in keywords]
                files = (
                    SESSION.query(Files.file_name)
                    .filter(or_(*conditions))
                    .order_by(Files.file_name.desc())
                    .limit(150)
                    .all()
                )

                extracted_titles = []
                seen = set()
                for (f_name,) in files:
                    title = _extract_clean_title(f_name)
                    if title and len(title) >= 2:
                        t_key = title.lower()
                        if t_key not in seen:
                            seen.add(t_key)
                            extracted_titles.append(title)
                            if len(extracted_titles) >= limit:
                                break

                return extracted_titles
        except Exception as e:
            LOGGER.warning(f"Error extracting DB keyword suggestions: {e}")
            retries -= 1
            await asyncio.sleep(0.2)
        finally:
            try:
                SESSION.close()
            except Exception:
                pass
    return []

async def count_files():
    """Count the total number of files in the database."""
    retries = 3
    while retries > 0:
        try:
            with INSERTION_LOCK:
                total_count = SESSION.query(Files).count()
                return total_count
        except PendingRollbackError:
            SESSION.rollback()
            retries -= 1
            continue
        except OperationalError as e:
            LOGGER.warning(f"OperationalError: {e}. Retrying...")
            reconnect_session()
            retries -= 1
        except Exception as e:
            LOGGER.warning(f"Error occurred while counting files: {e}")
            return 0
        finally:
            try:
                SESSION.close()
            except Exception as close_error:
                LOGGER.error(f"Error closing session: {close_error}")
    return 0

async def keep_alive():
    """Keep the database connection alive."""
    while True:
        try:
            with SESSION() as session:
                session.execute("SELECT 1")
            await asyncio.sleep(180)
        except Exception as e:
            LOGGER.warning(f"Keep-alive error: {e}")
            await asyncio.sleep(10)

async def main():
    asyncio.create_task(keep_alive())

if __name__ == "__main__":
    asyncio.run(main())