src/database/weavservice.py

from functools import reduce
import weaviate as wvt
import datetime, os
from threading import Lock

from time import perf_counter, sleep
from weaviate.classes.config import Configure, Property, DataType
from weaviate.collections.classes.grpc import MetadataQuery
from weaviate.collections.collection import Collection
from weaviate.classes.init import AdditionalConfig, Timeout
from weaviate.classes.query import Filter
from weaviate.config import AdditionalConfig

from src.utils.logging import get_logger
from config import WeaviateConfiguration as wvtconf, AVAILABLE_LANGUAGES, HASH_FILE_PATH

logger = get_logger("weaviate_service")

_get_collection_name = lambda lang: f'{wvtconf.WEAVIATE_COLLECTION_BASENAME}_{lang}'
_collection_names = [_get_collection_name(lang) for lang in AVAILABLE_LANGUAGES]


class WeaviateService:
    """
    Provides an interface for interacting with the Weaviate vector database.
    Handles initialization, data import, and hybrid queries.
    """
    
    _instance = None 
    _init_lock = Lock()

    def __new__(cls):
        if cls._instance is None:
            with cls._init_lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self) -> None:
        """
        Initialize the Weaviate service.
        """
        if hasattr(self, '_initialized'):
            return 

        self._connection_type = 'local' if wvtconf.is_local() else 'cloud'
        self._client = None 
        self._client_lock = Lock()
        
        # Some parameters to ensure that the connection will not be closed 
        # during long pauses in conversations
        self._last_query_time = perf_counter()
        self._idle_timeout = 25 * 60
        self._initialized = True

        # Initialize the client for the first time 
        logger.info("Initializing Weaviate service...")
        try:
            self._init_client()
            logger.info("Weaviate service initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize Weaviate service: {e}")
            raise e


    def _init_client(self) -> wvt.WeaviateClient:
        """
        Initializes the weaviate client with additional configuration.
        Performs a warm-up querying to speed-up the subsequent calls.

        Returns:
            configured Weaviate client instance on successfull connection.
        
        Raises:
            WeaviateConnectionError of the last failed connection if connection fails after 3 retires.
        """
        # Returns the client if it hasn't been idling for too long 
        if self._client is not None:
            time_since_query = perf_counter() - self._last_query_time
            if time_since_query < self._idle_timeout:
                return self._client
            
            # The connection might be closed, clients has to be reconnected 
            logger.warning(f"Client has been idling for too long. Reconnecting to prevent server-side closure...")
            try:
                self._client.close()
            except Exception as _:
                pass 

            self._client = None 
        
        # Client initialization 
        with self._client_lock:
            if self._client:
                return self._client

            retries = 0
            last_exception: Exception = None
            while retries < 3:
                try:
                    if wvtconf.is_local():
                        self._client = wvt.connect_to_local()
                        break

                    self._client = wvt.connect_to_weaviate_cloud(
                        cluster_url=wvtconf.CLUSTER_URL,
                        auth_credentials=wvtconf.WEAVIATE_API_KEY,
                        additional_config=AdditionalConfig(
                            timeout=Timeout(
                                init=wvtconf.INIT_TIMEOUT, 
                                query=wvtconf.QUERY_TIMEOUT, 
                                insert=wvtconf.INSERT_TIMEOUT,
                            ),
                            skip_init_checks=False,
                        ),
                        headers={
                            "X-HuggingFace-Api-Key": wvtconf.HUGGING_FACE_API_KEY,
                        },
                    ) 

                    # Warm-up query 
                    logger.info("Running warm-up query to initialize server...")
                    try:
                        collection = _get_collection_name(AVAILABLE_LANGUAGES[0])
                        self._client.collections.exists(collection)
                        logger.info("Warm-up finished - server is ready!")
                    except Exception as warmup_err:
                        logger.warning(f"Warm-up query failed (non-critical): {warmup_err}")
                    
                    break
                except Exception as e:
                    last_exception = e
                    logger.warning(f"Failed to establish connection on try {retries}: {e}")
                    retries += 1 
                    sleep(1)
            
            if retries == 3:
                logger.error(f"Failed to establish connection after 3 retries!")
                raise last_exception

            logger.info(f"Successully connected to the {self._connection_type} weaviate database")
            self._last_query_time = perf_counter()
            return self._client


    def _select_collection(self, lang: str) -> tuple[Collection, str]:
        """
        Select a language-specific collection as the active working collection.

        Args:
            lang (str): Acceptable language code.

        Raises:
            weaviate.exceptions.WeaviateConnectionError: If the specified language collection does not exist.
        """
        if lang not in AVAILABLE_LANGUAGES:
            logger.error(f"No collection for language '{lang}' was found in the database")
            return None, ''

        collection_name = _get_collection_name(lang)
        logger.info(f"Using collection {collection_name}")

        client = self._init_client()
        return client.collections.use(collection_name), collection_name


    def batch_import(self, data_rows: list, lang: str) -> list:
        """
        Perform a batch import of multiple objects into the current collection.

        Args:
            data_rows (list): List of dictionaries representing the data rows to import.
            lang (str, optional): Language collection to use. If not provided, uses the current one.

        Returns:
            list[dict]: List of failed imports with error details, if any.

        Raises:
            If no active collection is available or a connection error was catched.
        """        
        collection, collection_name = self._select_collection(lang)
        if collection is None:
            logger.error("No working collection selected!")
            return []

        import_errors = []
        logger.info(f"Batch importing {len(data_rows)} rows into {collection_name}")
        
        try:
            with self._client_lock:
                with collection.batch.fixed_size(batch_size=100, concurrent_requests=2) as batch:
                    for idx, data_row in enumerate(data_rows):
                        try:
                            batch.add_object(properties=data_row)
                        except Exception as e:
                            import_errors.append({'index': idx, 'chunk_id': data_row['chunk_id'], 'error': str(e)})
                            
                        if idx % 20 == 0 and idx > 0:
                            if batch.number_errors > 0:
                                logger.info(f"Failed imports at index {idx}: {batch.number_errors}")
            
            self._last_query_time = perf_counter()
            logger.info(f"Batch import finished. Total errors: {len(import_errors)}")
            
        except Exception as e:
            if 'connection' in str(e).lower():
                logger.error(f"Connection error during batch import: {e}")
                self._client = None
            raise e  

        return import_errors
    
    @staticmethod
    def _create_property_filter(prop, values) -> Filter:
        match prop:
            case 'programs':
                return Filter.by_property('programs').contains_any(values)
            case _:
                return None


    def query(self, query: str, lang: str, property_filters: dict[str], limit: int = 5) -> dict:
        """
        Execute a hybrid semantic and keyword query against the active collection with automatic reconnection on idle timeout.

        Args:
            query (str): The query string.
            lang (str, optional): Language collection to use. If not provided, uses the current one.
            property_filters (dict[str, any]): Key-value pairs for metadata filtering. Keys correspond
            to document properties (e.g., 'program', 'topic'), and values are the required matches.
            Multiple filters are combined using logical AND.
            limit (int, optional): Maximum number of results to return. Defaults to 5.


        Returns:
            tuple: A tuple containing the query response and elapsed time.

        Raises:
            weaviate.exceptions.WeaviateConnectionError: If no active collection is available.
        """
        retry_count = 0 
        max_retries = 2
            
        filters = [self._create_property_filter(prop, values) 
                   for prop, values in property_filters.items()] if property_filters else None
        filters = reduce(lambda f1, f2: f1 & f2, filters)

        while retry_count < max_retries:
            try:
                collection, collection_name = self._select_collection(lang)
                if collection is None:
                    logger.error("No working collection selected upon starting of the querying!")
                    return [], 0
                
                logger.info(f"Querying collection {collection_name}")
                query_start_time = perf_counter()

                with self._client_lock:
                    resp = collection.query.hybrid(
                        query=query,
                        filters=filters,
                        limit=limit,
                        return_metadata=MetadataQuery.full()
                    )
                elapsed = perf_counter() - query_start_time
                self._last_query_time = perf_counter()
                logger.info(f"Querying retrieved {len(resp.objects)} objects in {elapsed:3.2f} seconds")

                return (resp, elapsed)
            except Exception as e:
                if any(err_type in str(e).lower() for err_type in ['reset', 'closed', 'grpc', 'unavailable']):
                    retry_count += 1 
                    logger.warning(f"Connection error detected: {e}. Retrying...")

                    if retry_count == max_retries:
                        raise e
                else: # Probably not a server issue
                    raise e


    def _create_collections(self):
        """
        Create and initialize language-specific collections.
        
        Creates collections for all available languages with vector configuration.
        """
        try:
            client = self._init_client()
            logger.info('Attempting collections creation...')
            
            vector_config = (
                Configure.Vectors.text2vec_transformers() if wvtconf.is_local() 
                else Configure.Vectors.text2vec_huggingface(
                    name='hsg_rag_embeddings',
                    source_properties=['body'],
                    model="sentence-transformers/all-MiniLM-L6-v2",
                )
            )
            
            successful_creations = 0
            
            with self._client_lock:
                for collection_name in _collection_names:
                    try:
                        client.collections.create(
                            name=collection_name,
                            properties=[
                                Property(name='body', data_type=DataType.TEXT),
                                Property(name='chunk_id', data_type=DataType.TEXT),
                                Property(name='document_id', data_type=DataType.TEXT),
                                Property(name='programs', data_type=DataType.TEXT_ARRAY),
                                Property(name='source', data_type=DataType.TEXT),
                                Property(name='date', data_type=DataType.DATE)
                            ],
                            vector_config=vector_config
                        )
                        logger.info(f"Created collection {collection_name}")
                        successful_creations += 1
                    except Exception as e:
                        logger.error(f"Failed to create collection '{collection_name}': {e}")
            
            self._last_query_time = perf_counter()
            
            if successful_creations == len(_collection_names):
                logger.info('All collections successfully instantiated')
            else:
                logger.warning(f"Only {successful_creations}/{len(_collection_names)} collections created")
                
        except Exception as e:
            logger.error(f"Collections creation failed: {e}")
            self._client = None
            raise e


    def _delete_collections(self):
        """
        Delete all existing collections from the database.
        
        Also removes the hash file if it exists.        
        """
        try:
            client = self._init_client()
            logger.info("Initiating deletion of stored collections...")
            
            deleted_count = 0
            with self._client_lock:
                for collection_name in _collection_names:
                    try:
                        if client.collections.exists(collection_name):
                            client.collections.delete(collection_name)
                            logger.info(f"Deleted collection {collection_name}")
                            deleted_count += 1
                        else:
                            logger.warning(f"Collection {collection_name} does not exist")
                    except Exception as e:
                        logger.error(f"Failed to delete collection {collection_name}: {e}")
            
            self._last_query_time = perf_counter()
            logger.info(f"Deleted {deleted_count}/{len(_collection_names)} collections")
            
            # Clean up hash file
            if os.path.exists(HASH_FILE_PATH):
                try:
                    logger.info(f"Removing hash file: {HASH_FILE_PATH}")
                    os.remove(HASH_FILE_PATH)
                    logger.info("Hash file deleted successfully")
                except Exception as e:
                    logger.error(f"Failed to delete hash file: {e}")
                    
        except Exception as e:
            logger.error(f"Collections deletion failed: {e}")
            self._client = None
            raise e
    

    def _reset_collections(self):
        self._delete_collections()
        self._create_collections()


    def _collect_chunk_ids(self) -> dict:
        client = self._init_client()
        try:
            ids = []
            with self._client_lock:
                for c in client.collections.list_all(simple=False):
                    coll = client.collections.get(c)
                    for obj in coll.iterator():
                        ids.append(obj.properties['chunk_id'])
            return ids
        except Exception as e:
            logger.error(f"Failed to collect chunk ids: {e}")
            raise e


    def _extract_data(self) -> dict:
        client = self._init_client()
        try:
            schema = []
            objects = {}
            with self._client_lock:
                for c in client.collections.list_all(simple=False):
                    coll = client.collections.get(c)
                    cfg = coll.config.get().to_dict()
                    schema.append(cfg)

                    objects[c] = []
                    for obj in coll.iterator(include_vector=True):
                        objects[c].append({
                            "uuid": obj.uuid,
                            "properties": obj.properties,
                            "vector": obj.vector,
                        })

            return {
                'schema':  schema,
                'objects': objects,
            }
        except Exception as e:
            logger.error(f"Failed to extract data from database: {e}")
            raise e


    def _create_backup(self) -> str:
        """
        Create a backup of the current database state and stores it under selected backup provider.

        Returns: backup id of the created backup.
        """
        try:
            if not wvtconf.BACKUP_METHOD:
                raise ValueError('Backup method is not selected!')
            if wvtconf.BACKUP_METHOD not in wvtconf.AVAILABLE_BACKUP_METHODS:
                raise ValueError(f"Selected backup method 'wvtconf.BACKUP_METHOD' is not supported!")
            if not wvtconf.BACKUP_PATH:
                raise ValueError("Backup directory is not set!")
            os.makedirs(wvtconf.BACKUP_PATH, exist_ok=True)

            backup_id = f"backup_{datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')}"
            logger.info(f"Initiating backup creation for {self._connection_type} database...")
            
            match wvtconf.BACKUP_METHOD:
                case 'manual':
                    import json 
                    
                    backup_path = os.path.join(wvtconf.BACKUP_PATH, backup_id)
                    os.makedirs(backup_path)
                    
                    db_data = self._extract_data()
                    data_backup = {
                        'creation_date': datetime.datetime.now().isoformat(),
                    }
                    
                    schema_backup_path = os.path.join(backup_path, 'schema.json')
                    with open(schema_backup_path, 'w', encoding='utf-8') as f:
                        json.dump(db_data['schema'], f, indent=2, default=str)
                    
                    objects_backup_path = os.path.join(backup_path, 'objects.json')
                    with open(objects_backup_path, 'w', encoding='utf-8') as f:
                        json.dump(db_data['objects'], f, indent=2, default=str)

                    data_backup_path = os.path.join(backup_path, 'data.json')
                    with open(data_backup_path, 'w', encoding='utf-8') as f:
                        json.dump(data_backup, f, indent=2, default=str)

                case 's3':
                    client = self._init_client()
                    with self._client_lock:
                        client.backup.create(
                            backup_id=backup_id,
                            backend="s3",
                            include_collections=_collection_names,
                            wait_for_completion=True,
                        ) 
                case _:
                    raise NotImplementedError()

                
            self._last_query_time = perf_counter()
            logger.info(f"Backup '{backup_id}' created successfully")
            
            return backup_id
        except Exception as e:
            logger.error(f"Backup creation failed: {e}")
            raise e


    def _restore_backup(self, backup_id: str):
        """
        Restore the database state from a backup.
        
        Restores specified collections from backup.
        
        Args:
            backup_id: ID of the backup to restore from
            
        Raises:
            Exception if backup restoration fails
        """
        self._delete_collections()

        try:
            if not wvtconf.BACKUP_METHOD:
                raise ValueError('Backup method is not selected!')
            if wvtconf.BACKUP_METHOD not in wvtconf.AVAILABLE_BACKUP_METHODS:
                raise ValueError(f"Selected backup method 'wvtconf.BACKUP_METHOD' is not supported!")
            if not wvtconf.BACKUP_PATH:
                raise ValueError("Backup directory is not set!")
            os.makedirs(wvtconf.BACKUP_PATH, exist_ok=True)
            
            backup_path = os.path.join(wvtconf.BACKUP_PATH, backup_id)
            if not os.path.exists(backup_path):
                raise RuntimeError(f"Directory for backup 'backup_id' does not exist in the backup directory!")  
            schema_backup_path = os.path.join(backup_path, 'schema.json')
            if not os.path.exists(schema_backup_path):
                raise RuntimeError(f"Schema backup is missing in the backup directory!")
            objects_backup_path = os.path.join(backup_path, 'objects.json')
            if not os.path.exists(objects_backup_path):
                raise RuntimeError(f"Objects backup is missing in the backup directory!")
            
            client = self._init_client()
            logger.info(f"Initiating restoration from backup '{backup_id}' for {self._connection_type} database...")
            
            with self._client_lock:
                match wvtconf.BACKUP_METHOD:
                    case 'manual':
                        import json

                        with open(schema_backup_path) as f:
                            schemas = json.load(f)
                            for cfg in schemas: 
                                client.collections.create_from_dict(cfg)

                        with open(objects_backup_path) as f:
                            data = json.load(f)
                            for name, objs in data.items():
                                logger.info(f"Restoring collection '{name}' with {len(objs)} objects...")
                                coll = client.collections.get(name)
                                
                                with coll.batch.dynamic() as batch:
                                    for o in objs:
                                        o['properties']['date'] = o['properties']['date'] \
                                                .replace(" ", "T").replace("+00:00", "Z")
                                        batch.add_object(
                                            uuid=o["uuid"],
                                            properties=o["properties"],
                                            vector=o["vector"]
                                        )
                                logger.info(f"Collection '{name}' restored successfully")
                    case 's3':                        
                        client.backup.restore(
                            backup_id=backup_id,
                            backend="s3",
                            wait_for_completion=True,
                            roles_restore="all",
                            users_restore="all",
                        )
                    case _:
                        raise NotImplementedError()

            self._last_query_time = perf_counter()
            logger.info(f"Backup '{backup_id}' restored successfully")
            
        except Exception as e:
            error_msg = str(e).lower()
            if 'connection' in error_msg:
                logger.error(f"Connection error during backup restore: {e}. Will reconnect on next operation.")
                self._client = None
            logger.error(f"Backup restoration failed: {e}")
            raise e


    def _checkhealth(self) -> bool:
        """
        Check the connectivity and health status of the Weaviate database.
        
        Verifies:
        - Connection to the database
        - Database metadata and version
        - Existence of all expected collections
        - Module availability
        
        Returns:
            True if all health checks pass, False otherwise
        """
        try:
            client = self._init_client()
            
            # Check basic connectivity
            is_connected = False
            with self._client_lock:
                is_connected = client.is_connected()
            
            connection_status = "✓ OK" if is_connected else "✗ ERROR"
            logger.info(f"Connection to {self._connection_type} database: {connection_status}")
            
            if not is_connected:
                logger.error("Database connection check failed")
                return False
            
            # Get and log metadata
            try:
                with self._client_lock:
                    metainfo = client.get_meta()
                
                # Format module information
                modules = metainfo.get('modules', {})
                modules_list = list(modules.keys()) if isinstance(modules, dict) else modules
                modules_str = ', '.join(str(m) for m in modules_list) if modules_list else 'None'
                
                # Truncate long module strings for logging
                if len(modules_str) > 50:
                    modules_str = modules_str[:47] + '...'
                
                # Log connection details
                if wvtconf.is_local():
                    logger.info(
                        f"Database metadata: "
                        f"HOSTNAME={metainfo.get('hostname', 'unknown')}, "
                        f"VERSION={metainfo.get('version', 'unknown')}, "
                        f"MODULES={modules_str}"
                    )
                else:
                    logger.info(
                        f"Database metadata: "
                        f"VERSION={metainfo.get('version', 'unknown')}, "
                        f"MODULES={modules_str}"
                    )
                
            except Exception as e:
                logger.warning(f"Could not retrieve database metadata: {e}")
            
            # Check collection existence
            all_collections_exist = True
            
            with self._client_lock:
                for collection_name in _collection_names:
                    try:
                        exists = client.collections.exists(collection_name)
                        status = "✓ OK" if exists else "✗ MISSING"
                        logger.info(f"Collection '{collection_name}': {status}")
                        
                        if not exists:
                            all_collections_exist = False
                            
                    except Exception as e:
                        logger.error(f"Error checking collection '{collection_name}': {e}")
                        all_collections_exist = False
            
            # Update last health check time
            self._last_query_time = perf_counter()
            
            # Log overall health status
            if is_connected and all_collections_exist:
                logger.info("✓ Database health check PASSED - All systems operational")
                return True
            else:
                logger.warning("✗ Database health check FAILED - Some issues detected")
                return False
                
        except Exception as e:
            error_msg = str(e).lower()
            if 'connection' in error_msg:
                logger.error(f"Connection error during health check: {e}. Will reconnect on next operation.")
                self._client = None
            logger.error(f"Health check failed: {e}")
            return False


def parse_arguments():
    """
    Parse command-line arguments for managing Weaviate collections.

    Returns:
        argparse.Namespace: Parsed command-line arguments.
    """
    import argparse
    
    parser = argparse.ArgumentParser(
        description='Weaviate database management utility'
    )
    group = parser.add_mutually_exclusive_group()
    
    group.add_argument(
        '-dc', "--delete_collections", 
        action='store_true', 
        help='Delete all collections from the database'
    )
    group.add_argument(
        '-cc', "--create_collections", 
        action='store_true', 
        help='Initialize collections for different language contents'
    )
    group.add_argument(
        '-rc', "--redo_collections", 
        action='store_true', 
        help='Delete and recreate all collections'
    )
    group.add_argument(
        '-ch', "--checkhealth", 
        action='store_true', 
        help='Check database connection and collection existence'
    )
    group.add_argument(
        '-cb', "--create_backup", 
        action='store_true', 
        help='Create a backup of the current database state'
    )
    group.add_argument(
        '-rb', "--restore_backup", 
        type=str, 
        metavar='BACKUP_ID',
        help='Restore database from a backup (provide backup_id)'
    )
    
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_arguments()
    service = WeaviateService()

    if args.create_backup:
       service._create_backup()

    if args.restore_backup:
        service._restore_backup(args.restore_backup)

    if any([args.delete_collections, args.redo_collections]):
        service._delete_collections()

    if any([args.create_collections, args.redo_collections]):
        service._create_collections()
    
    if any([args.checkhealth, args.create_collections, args.redo_collections]):
        service._checkhealth()