modules/data_management/index_manager.py

import os
import logging
import json
import shutil
from pathlib import Path
import pandas as pd
from datetime import datetime, timedelta
import hashlib
import uuid
import faiss

from modules.data_management.index_utils import validate_index_directory
from modules.data_management.index_utils import check_indexing_availability, initialize_embedding_model
from modules.data_management.hash_utils import generate_data_hash
from modules.data_management.index_utils import check_index_integrity

from modules.config.paths import INDICES_DIR

from modules.config.ai_settings import (
    CHUNK_SIZE,
    CHUNK_OVERLAP,
    EXCLUDED_EMBED_METADATA_KEYS,
    EXCLUDED_LLM_METADATA_KEYS
)

logger = logging.getLogger(__name__)

# Перевірка доступності модулів для індексування
INDEXING_AVAILABLE = check_indexing_availability()
INDEXING_MODULES = {
    "VectorStoreIndex": None,
    "StorageContext": None,
    "SimpleDocumentStore": None,
    "TokenTextSplitter": None,
    "BM25Retriever": None,
    "FaissVectorStore": None,
    "Settings": None
}



def _generate_data_hash(self, df):
    """
    Генерація хешу для DataFrame для ідентифікації унікальних даних.
    
    Args:
        df (pandas.DataFrame): DataFrame для хешування
        
    Returns:
        str: Хеш даних
    """
    # Використовуємо основні колонки для хешування
    key_columns = ['Issue key', 'Summary', 'Status', 'Issue Type', 'Created', 'Updated']
    
    return generate_data_hash(df, key_columns)

class IndexManager:
    """
    Менеджер для створення та управління індексами даних (FAISS, BM25).
    """
    def __init__(self, base_indices_dir="temp/indices"):
        """
        Ініціалізація менеджера індексів.
        
        Args:
            base_indices_dir (str): Базова директорія для зберігання індексів
        """
        self.base_indices_dir = Path(base_indices_dir) if base_indices_dir else INDICES_DIR
        self.base_indices_dir.mkdir(exist_ok=True, parents=True)
        
        # Перевірка доступності модулів для індексування
        self.indexing_available = INDEXING_AVAILABLE
        if not self.indexing_available:
            logger.warning("Функціональність індексування недоступна. Встановіть необхідні пакети.")
    
    def create_indices_for_session(self, session_id, merged_df, indices_dir=None):
        """
        Створення індексів для даних сесії.
        
        Args:
            session_id (str): Ідентифікатор сесії
            merged_df (pandas.DataFrame): DataFrame з об'єднаними даними
            indices_dir (str, optional): Директорія для збереження індексів. 
                                        Якщо None, використовується директорія сесії.
                                        
        Returns:
            dict: Інформація про створені індекси
        """
        if not self.indexing_available:
            return {"error": "Функціональність індексування недоступна. Встановіть необхідні пакети."}
        
        try:
            # Визначаємо директорію для індексів
            indices_path = Path(indices_dir) if indices_dir else self.base_indices_dir / session_id
            indices_path.mkdir(exist_ok=True, parents=True)
            
            # Генеруємо хеш для даних
            data_hash = self._generate_data_hash(merged_df)
            
            # Перевіряємо, чи існують індекси для цих даних
            existing_indices = self._find_indices_by_hash(data_hash)
            
            if existing_indices:
                return self._reuse_existing_indices(existing_indices, indices_path, session_id, data_hash, merged_df)
            
            # Створюємо нові індекси
            return self._create_new_indices(indices_path, session_id, data_hash, merged_df)
            
        except Exception as e:
            logger.error(f"Помилка при створенні індексів: {e}")
            return {"error": f"Помилка при створенні індексів: {str(e)}"}
    
    def _reuse_existing_indices(self, existing_indices, indices_path, session_id, data_hash, merged_df):
        """
        Повторне використання існуючих індексів.
        
        Args:
            existing_indices (str): Шлях до існуючих індексів
            indices_path (Path): Шлях для нових індексів
            session_id (str): Ідентифікатор сесії
            data_hash (str): Хеш даних
            merged_df (pandas.DataFrame): DataFrame з даними
            
        Returns:
            dict: Інформація про скопійовані індекси
        """
        logger.info(f"Знайдено існуючі індекси для даних з хешем {data_hash}")
        
        try:
            # Спочатку очищаємо цільову директорію
            if indices_path.exists():
                for item in indices_path.iterdir():
                    if item.is_file():
                        item.unlink()
                    elif item.is_dir():
                        shutil.rmtree(item)
            
            # Копіюємо індекси
            for item in Path(existing_indices).iterdir():
                if item.is_file():
                    shutil.copy2(item, indices_path)
                elif item.is_dir():
                    shutil.copytree(item, indices_path / item.name)
            
            logger.info(f"Індекси успішно скопійовано в {indices_path}")
            
            # Оновлюємо метадані
            metadata = {
                "session_id": session_id,
                "created_at": datetime.now().isoformat(),
                "data_hash": data_hash,
                "rows_count": len(merged_df),
                "columns_count": len(merged_df.columns),
                "copied_from": str(existing_indices)
            }
            
            with open(indices_path / "metadata.json", "w", encoding="utf-8") as f:
                json.dump(metadata, f, ensure_ascii=False, indent=2)
            
            return {
                "success": True,
                "indices_dir": str(indices_path),
                "data_hash": data_hash,
                "reused_existing": True,
                "source": str(existing_indices)
            }
                
        except Exception as copy_err:
            logger.error(f"Помилка при копіюванні індексів: {copy_err}")
            # Продовжуємо створення нових індексів
            return self._create_new_indices(indices_path, session_id, data_hash, merged_df)
    
    def _create_new_indices(self, indices_path, session_id, data_hash, merged_df):
        """
        Створення нових індексів.
        Зберігає індекси у форматі, сумісному з jira_hybrid_chat.py.
        """
        if not INDEXING_AVAILABLE:
            return {"error": "Функціональність індексування недоступна"}
                    
        try:
            logger.info(f"Створення нових індексів для сесії {session_id}")
            
            # Імпортуємо необхідні модулі напряму
            from llama_index.core import VectorStoreIndex, StorageContext, Settings
            from llama_index.core.storage.docstore import SimpleDocumentStore
            from llama_index.core.node_parser import TokenTextSplitter
            from llama_index.retrievers.bm25 import BM25Retriever
            from llama_index.vector_stores.faiss import FaissVectorStore
            import faiss
            
            # Ініціалізуємо модель ембедингів
            from modules.data_management.index_utils import initialize_embedding_model
            embed_model = initialize_embedding_model()
            
            # Отримуємо розмірність ембедингів динамічно
            import numpy as np
            test_embedding = embed_model.get_text_embedding("Тестовий текст")
            embed_dim = len(test_embedding)
            logger.info(f"Розмірність ембедингів: {embed_dim}")
            
            # Конвертуємо DataFrame в документи
            documents = self._convert_dataframe_to_documents(merged_df)
            
            # Створюємо розділювач тексту
            text_splitter = TokenTextSplitter(
                chunk_size=CHUNK_SIZE,
                chunk_overlap=CHUNK_OVERLAP
            )
            
            # Встановлюємо формат збереження на JSON через глобальні налаштування
            # Це важливо для сумісності з jira_hybrid_chat.py
            Settings.persist_json_format = True
            
            # Створюємо FAISS індекс
            faiss_index = faiss.IndexFlatL2(embed_dim)
            
            # Створюємо контекст зберігання
            docstore = SimpleDocumentStore()
            vector_store = FaissVectorStore(faiss_index=faiss_index)
            
            storage_context = StorageContext.from_defaults(
                docstore=docstore,
                vector_store=vector_store
            )
            
            # Встановлюємо модель ембедингів у налаштуваннях
            Settings.embed_model = embed_model
            
            # Створюємо індекс
            index = VectorStoreIndex.from_documents(
                documents,
                storage_context=storage_context,
                transformations=[text_splitter]
            )
            
            # Зберігаємо індекс у форматі JSON (через глобальні налаштування)
            # НЕ передаємо json_format як аргумент
            index.storage_context.persist(persist_dir=str(indices_path))
            
            # Створюємо BM25 retriever
            bm25_retriever = BM25Retriever.from_defaults(
                docstore=docstore,
                similarity_top_k=10
            )
            
            # Зберігаємо параметри BM25
            bm25_dir = indices_path / "bm25"
            bm25_dir.mkdir(exist_ok=True)
            
            with open(bm25_dir / "params.json", "w", encoding="utf-8") as f:
                json.dump({"similarity_top_k": 10}, f)
            
            # Зберігаємо метадані
            metadata = {
                "session_id": session_id,
                "created_at": datetime.now().isoformat(),
                "data_hash": data_hash,
                "rows_count": len(merged_df),
                "columns_count": len(merged_df.columns),
                "embedding_model": embed_model.__class__.__name__,
                "embedding_dim": embed_dim,
                "format": "json"  # Вказуємо використаний формат збереження
            }
            
            with open(indices_path / "metadata.json", "w", encoding="utf-8") as f:
                json.dump(metadata, f, ensure_ascii=False, indent=2)

            with open(indices_path / "indices.valid", "w", encoding="utf-8") as f:
                f.write(f"Indices created at {datetime.now().isoformat()}")

            logger.info(f"Створено файл-маркер indices.valid")
            
            return {
                "success": True,
                "indices_dir": str(indices_path),
                "data_hash": data_hash
            }
            
        except Exception as e:
            logger.error(f"Помилка при створенні нових індексів: {e}")
            return {"error": f"Помилка при створенні нових індексів: {str(e)}"}

    def _save_bm25_data(self, indices_path, bm25_retriever):
        """
        Збереження даних для BM25 retriever.
        
        Args:
            indices_path (Path): Шлях до директорії індексів
            bm25_retriever (BM25Retriever): Об'єкт BM25Retriever
            
        Returns:
            bool: True, якщо дані успішно збережені, False у випадку помилки
        """
        try:
            # Створюємо директорію для BM25
            bm25_dir = indices_path / "bm25"
            bm25_dir.mkdir(exist_ok=True)
            
            # Зберігаємо параметри BM25
            bm25_params = {
                "similarity_top_k": bm25_retriever.similarity_top_k,
                "alpha": getattr(bm25_retriever, "alpha", 0.75),
                "beta": getattr(bm25_retriever, "beta", 0.75),
                "index_creation_time": datetime.now().isoformat()
            }
            
            with open(bm25_dir / "params.json", "w", encoding="utf-8") as f:
                json.dump(bm25_params, f, ensure_ascii=False, indent=2)
            
            logger.info(f"Дані BM25 збережено в {bm25_dir}")
            return True
            
        except Exception as e:
            logger.error(f"Помилка при збереженні даних BM25: {e}")
            return False
    
    def _convert_dataframe_to_documents(self, df):
        """
        Конвертує DataFrame в документи для індексування.
        
        Args:
            df (pandas.DataFrame): DataFrame для конвертації
            
        Returns:
            list: Список документів
        """
        try:
            # Імпортуємо Document напряму
            from llama_index.core import Document
            
            documents = []
            
            # Перебираємо рядки DataFrame
            for idx, row in df.iterrows():
                # Створюємо текст документа
                text = f"Issue Key: {row.get('Issue key', '')}\n"
                text += f"Summary: {row.get('Summary', '')}\n"
                text += f"Status: {row.get('Status', '')}\n"
                text += f"Issue Type: {row.get('Issue Type', '')}\n"
                
                # Додаємо опис, якщо він є
                if 'Description' in row and pd.notna(row['Description']):
                    text += f"Description: {row['Description']}\n"
                
                # Додаємо коментарі, якщо вони є
                if 'Comments' in row and pd.notna(row['Comments']):
                    text += f"Comments: {row['Comments']}\n"
                
                # Створюємо метадані
                metadata = {
                    "issue_key": row.get('Issue key', ''),
                    "summary": row.get('Summary', ''),
                    "status": row.get('Status', ''),
                    "issue_type": row.get('Issue Type', ''),
                    "created": str(row.get('Created', '')),
                    "updated": str(row.get('Updated', ''))
                }
                
                # Створюємо документ
                doc = Document(
                    text=text,
                    metadata=metadata
                )
                
                documents.append(doc)
            
            logger.info(f"Створено {len(documents)} документів з DataFrame")
            return documents
            
        except Exception as e:
            logger.error(f"Помилка при конвертації DataFrame в документи: {e}")
            raise

    def _generate_data_hash(self, df):
        """
        Генерація хешу для DataFrame для ідентифікації унікальних даних.
        
        Args:
            df (pandas.DataFrame): DataFrame для хешування
            
        Returns:
            str: Хеш даних
        """
        try:
            # Використовуємо основні колонки для хешування
            key_columns = ['Issue key', 'Summary', 'Status', 'Issue Type', 'Created', 'Updated']
            
            # Фільтруємо тільки наявні колонки
            available_columns = [col for col in key_columns if col in df.columns]
            
            if not available_columns:
                # Якщо немає жодної ключової колонки, використовуємо всі дані
                data_str = df.to_json()
            else:
                # Інакше використовуємо тільки ключові колонки
                data_str = df[available_columns].to_json()
            
            # Створюємо хеш
            hash_object = hashlib.sha256(data_str.encode())
            data_hash = hash_object.hexdigest()
            
            return data_hash
            
        except Exception as e:
            logger.error(f"Помилка при генерації хешу даних: {e}")
            # У випадку помилки повертаємо випадковий хеш
            return str(uuid.uuid4())
    
    def _find_indices_by_hash(self, data_hash):
        """
        Пошук існуючих індексів за хешем даних.
        
        Args:
            data_hash (str): Хеш даних
            
        Returns:
            str: Шлях до директорії з індексами або None, якщо не знайдено
        """
        try:
            # Перебираємо всі піддиректорії в базовій директорії індексів
            for index_dir in self.base_indices_dir.iterdir():
                if not index_dir.is_dir():
                    continue
                
                # Перевіряємо метадані
                metadata_file = index_dir / "metadata.json"
                if not metadata_file.exists():
                    continue
                
                try:
                    with open(metadata_file, "r", encoding="utf-8") as f:
                        metadata = json.load(f)
                    
                    # Перевіряємо хеш
                    if metadata.get("data_hash") == data_hash:
                        # Перевіряємо наявність необхідних файлів
                        if validate_index_directory(index_dir):
                            logger.info(f"Знайдено існуючі індекси з відповідним хешем: {index_dir}")
                            return str(index_dir)
                        else:
                            logger.warning(f"Знайдено індекси з відповідним хешем, але вони неповні: {index_dir}")
                except Exception as e:
                    logger.error(f"Помилка при перевірці метаданих {metadata_file}: {e}")
            
            logger.info(f"Не знайдено існуючих індексів з хешем {data_hash}")
            return None
            
        except Exception as e:
            logger.error(f"Помилка при пошуку індексів за хешем: {e}")
            return None

    def cleanup_old_indices(self, max_age_days=7, max_indices=20):
        """
        Очищення застарілих індексів.
        
        Args:
            max_age_days (int): Максимальний вік індексів у днях
            max_indices (int): Максимальна кількість індексів для зберігання
            
        Returns:
            int: Кількість видалених директорій індексів
        """
        try:
            # Перевіряємо, чи існує базова директорія
            if not self.base_indices_dir.exists():
                return 0
            
            # Отримуємо список директорій індексів
            index_dirs = []
            
            for index_dir in self.base_indices_dir.iterdir():
                if not index_dir.is_dir():
                    continue
                
                # Перевіряємо метадані для отримання часу створення
                metadata_file = index_dir / "metadata.json"
                created_at = None
                
                if metadata_file.exists():
                    try:
                        with open(metadata_file, "r", encoding="utf-8") as f:
                            metadata = json.load(f)
                        created_at = metadata.get("created_at")
                    except Exception:
                        pass
                
                # Якщо немає метаданих, використовуємо час створення директорії
                if not created_at:
                    created_at = datetime.fromtimestamp(index_dir.stat().st_mtime).isoformat()
                
                # Додаємо інформацію про директорію
                index_dirs.append({
                    "path": str(index_dir),
                    "created_at": created_at
                })
            
            # Якщо немає директорій для обробки, повертаємо 0
            if not index_dirs:
                return 0
            
            # Сортуємо директорії за часом створення (від найновіших до найстаріших)
            index_dirs.sort(key=lambda x: x["created_at"], reverse=True)
            
            # Визначаємо директорії для видалення
            dirs_to_delete = []
            
            # 1. Залишаємо max_indices найновіших директорій
            if len(index_dirs) > max_indices:
                dirs_to_delete.extend(index_dirs[max_indices:])
            
            # 2. Перевіряємо, чи є серед залишених застарілі директорії
            cutoff_date = (datetime.now() - timedelta(days=max_age_days)).isoformat()
            
            for index_info in index_dirs[:max_indices]:
                if index_info["created_at"] < cutoff_date:
                    dirs_to_delete.append(index_info)
            
            # Видаляємо директорії
            deleted_count = 0
            
            for dir_info in dirs_to_delete:
                try:
                    dir_path = Path(dir_info["path"])
                    if dir_path.exists():
                        shutil.rmtree(dir_path)
                        logger.info(f"Видалено застарілу директорію індексів: {dir_path}")
                        deleted_count += 1
                except Exception as e:
                    logger.error(f"Помилка при видаленні директорії {dir_info['path']}: {e}")
            
            return deleted_count
            
        except Exception as e:
            logger.error(f"Помилка при очищенні застарілих індексів: {e}")
            return 0
    
    def load_indices(self, indices_dir):
        """
        Завантаження індексів з директорії.
        
        Args:
            indices_dir (str): Шлях до директорії з індексами
            
        Returns:
            tuple: (VectorStoreIndex, BM25Retriever) або (None, None) у випадку помилки
        """
        if not self.indexing_available:
            logger.warning("Функціональність індексування недоступна. Встановіть необхідні пакети.")
            return None, None
        
        try:
            # Перевіряємо цілісність індексів
            is_valid, message = check_index_integrity(indices_dir)
            if not is_valid:
                logger.error(f"Індекси не пройшли перевірку цілісності: {message}")
                return None, None

            indices_path = Path(indices_dir)
            
            if not indices_path.exists():
                logger.error(f"Директорія індексів не існує: {indices_dir}")
                return None, None
            
            # Перевіряємо наявність необхідних файлів
            if not (indices_path / "docstore.json").exists():
                logger.error(f"Директорія індексів не містить необхідних файлів: {indices_dir}")
                return None, None
            
            # Імпортуємо необхідні модулі
            StorageContext = INDEXING_MODULES.get("StorageContext")
            VectorStoreIndex = INDEXING_MODULES.get("VectorStoreIndex")
            BM25Retriever = INDEXING_MODULES.get("BM25Retriever")
            
            # Завантажуємо контекст зберігання
            storage_context = StorageContext.from_defaults(persist_dir=str(indices_path))
            
            # Завантажуємо індекс
            index = VectorStoreIndex.from_storage_context(storage_context)
            
            # Створюємо BM25 retriever
            bm25_retriever = BM25Retriever.from_defaults(
                docstore=storage_context.docstore,
                similarity_top_k=10
            )
            
            # Завантажуємо параметри BM25, якщо вони є
            bm25_params_file = indices_path / "bm25" / "params.json"
            if bm25_params_file.exists():
                try:
                    with open(bm25_params_file, "r", encoding="utf-8") as f:
                        bm25_params = json.load(f)
                    
                    # Встановлюємо параметри
                    if "similarity_top_k" in bm25_params:
                        bm25_retriever.similarity_top_k = bm25_params["similarity_top_k"]
                except Exception as e:
                    logger.warning(f"Помилка при завантаженні параметрів BM25: {e}")
            
            logger.info(f"Індекси успішно завантажено з {indices_dir}")
            
            return index, bm25_retriever
            
        except Exception as e:
            logger.error(f"Помилка при завантаженні індексів: {e}")
            return None, None