import streamlit as st
import pandas as pd
import json
import os
import logging
import re
from fuzzywuzzy import fuzz
import sqlite3
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from rank_bm25 import BM25Okapi
import nltk
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
import openai
import time
from huggingface_hub import model_info
from datetime import datetime
import torch  # Убедитесь, что этот импорт есть

# 1. Настройка логирования
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("model_loading.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger()

# Добавляем информацию о PyTorch и CUDA
logger.info(f"PyTorch version: {torch.__version__}")
logger.info(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    logger.info(f"CUDA device: {torch.cuda.get_device_name(0)}")


# 2. Проверка загрузки модели
try:
    logger.info("="*50)
    logger.info("Начало принудительной проверки модели")
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    test_model = SentenceTransformer(
        "cointegrated/LaBSE-en-ru",
        cache_folder="/tmp/hf_cache_force"
    )
    # Изменяем порядок инициализации
    test_model = test_model.to('cpu')  # Сначала явно переносим на CPU
    
    # Проверяем работоспособность
    test_text = ["тестовый текст"]
    with torch.no_grad():
        embeddings = test_model.encode(test_text)
    
    logger.info(f"Модель загружена. Размерность: {test_model.get_sentence_embedding_dimension()}")
    del test_model
    
except Exception as e:
    logger.critical(f"Тестовая загрузка модели провалилась: {str(e)}", exc_info=True)
    st.error("""
    ❌ Критическая ошибка: модель не загружается!
    Проверьте:
    1. Интернет-соединение
    2. Доступ к Hugging Face Hub
    3. Логи в файле model_loading.log
    """)
    raise


# 3. Инициализация NLTK


# 4. Константы
XLSX_FILE_PATH = "Test_questions_from_diagnostpb (1).xlsx"
SQLITE_DB_PATH = "knowledge_base_v1.db"
LOG_FILE = "chat_logs.json"
EMBEDDING_MODEL = "cointegrated/LaBSE-en-ru"

# Определяем базовую директорию и пути к файлам
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
VECTOR_DB_DIR = os.path.join(BASE_DIR, "vectorized_knowledge_base")
VECTOR_DB_PATH = os.path.join(VECTOR_DB_DIR, "processed_knowledge_base_v1.db")
FAISS_INDEX_PATH = os.path.join(VECTOR_DB_DIR, "faiss_index.bin")


# Добавляем проверку прав доступа
if os.path.exists(VECTOR_DB_PATH):
    logger.info(f"File permissions: {oct(os.stat(VECTOR_DB_PATH).st_mode)[-3:]}")
    logger.info(f"File size: {os.path.getsize(VECTOR_DB_PATH)} bytes")

# Добавьте отладочное логирование
logger.info(f"BASE_DIR: {BASE_DIR}")
logger.info(f"VECTOR_DB_DIR: {VECTOR_DB_DIR}")
logger.info(f"VECTOR_DB_PATH: {VECTOR_DB_PATH}")
logger.info(f"Directory exists: {os.path.exists(VECTOR_DB_DIR)}")
logger.info(f"Database file exists: {os.path.exists(VECTOR_DB_PATH)}")

# После определения путей
required_files = [
    (VECTOR_DB_PATH, "База данных векторов"),
    (FAISS_INDEX_PATH, "FAISS индекс"),
    (SQLITE_DB_PATH, "SQLite база знаний"),
    (XLSX_FILE_PATH, "Excel файл с вопросами")
]

for file_path, description in required_files:
    if not os.path.exists(file_path):
        logger.error(f"Не найден файл: {description} ({file_path})")
        st.error(f"❌ Отсутствует необходимый файл: {description}")
        st.stop()
    elif os.path.getsize(file_path) == 0:
        logger.error(f"Файл пуст: {description} ({file_path})")
        st.error(f"❌ Файл пуст: {description}")
        st.stop()

# 5. Инициализация OpenAI
openai_api_key = os.getenv('VSEGPT_API_KEY')
if openai_api_key is None:
    logger.error("Переменная окружения VSEGPT_API_KEY не установена")
    st.warning("Не настроен API-ключ для OpenAI")   
    raise ValueError("Переменная окружения VSEGPT_API_KEY не установена")

openai.api_key = openai_api_key
openai.api_base = "https://api.vsegpt.ru/v1"

# Инициализация сессии
if "logs" not in st.session_state:
    st.session_state.logs = []
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "user_input" not in st.session_state:
    st.session_state.user_input = ''
if "widget" not in st.session_state:
    st.session_state.widget = ''

def setup_nltk():
    try:
        nltk.download('punkt', quiet=True)
        nltk.download('stopwords', quiet=True)
        # Используем базовый токенизатор без специфичных для языка ресурсов
        from nltk.tokenize import word_tokenize
        test_text = "тестовый текст"
        tokens = word_tokenize(test_text)  # Убираем параметр language
        logger.info(f"NLTK успешно инициализирован. Тестовая токенизация: {tokens}")
    except Exception as e:
        logger.warning(f"Ошибка инициализации NLTK: {e}")

setup_nltk()

def get_documents_list():
    try:
        conn = sqlite3.connect(VECTOR_DB_PATH)
        cursor = conn.cursor()
        cursor.execute("""
            SELECT DISTINCT doc_type_short, doc_number, file_name 
            FROM documents 
            ORDER BY doc_type_short, doc_number
        """)
        documents = cursor.fetchall()
        conn.close()
        
        # Форматируем список документов
        formatted_docs = []
        for doc in documents:
            doc_parts = [
                str(part) for part in doc 
                if part is not None and str(part).strip()
            ]
            if doc_parts:
                formatted_docs.append(" ".join(doc_parts))
        
        return formatted_docs
        
    except Exception as e:
        logger.error(f"Ошибка при получении списка документов: {e}")
        return []


class HybridSearch:
    def __init__(self, db_path):
        self.db_path = db_path
        self.stop_words = set(stopwords.words('russian')).union({
            '', ' ', '  ', '\t', '\n', '\r', 'nbsp'
        })
        logger.info(f"Загружено стоп-слов: {len(self.stop_words)}")
        self.bm25 = None
        self.corpus = []
        self.doc_ids = []
        self._init_bm25_with_fallback()
    
    def _init_bm25_with_fallback(self):
        """Инициализация с резервным вариантом при ошибках"""
        try:
            self._init_bm25()
            if not self.bm25:
                logger.warning("Основная инициализация BM25 не удалась, создаем резервный индекс")
                self._create_fallback_index()
        except Exception as e:
            logger.error(f"Ошибка при инициализации BM25: {str(e)}")
            self._create_fallback_index()

    def _init_bm25(self):
        """Основная инициализация BM25"""
        if not os.path.exists(self.db_path):
            raise FileNotFoundError(f"Файл БД не найден: {self.db_path}")

        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        cursor = conn.cursor()

        try:
            cursor.execute("SELECT COUNT(*) FROM content")
            count = cursor.fetchone()[0]
            logger.info(f"Найдено {count} документов в таблице content")
            
            if count == 0:
                raise ValueError("Таблица content пуста")

            cursor.execute("SELECT id, chunk_text FROM content")
            valid_docs = 0
            
            for row in cursor:
                try:
                    text = str(row['chunk_text']).strip()
                    if not text:
                        continue
                        
                    tokens = self._preprocess_text(text)
                    if tokens and len(tokens) >= 2:
                        self.corpus.append(tokens)
                        self.doc_ids.append(row['id'])
                        valid_docs += 1
                        if valid_docs % 1000 == 0:
                            logger.info(f"Обработано {valid_docs} документов")
                except Exception as e:
                    logger.warning(f"Ошибка обработки документа ID {row['id']}: {str(e)}")

            if valid_docs == 0:
                raise ValueError("Нет пригодных документов после обработки")
            
            logger.info(f"Создание BM25 индекса для {valid_docs} документов")
            self.bm25 = BM25Okapi(self.corpus)
            logger.info(f"BM25 успешно инициализирован с {valid_docs} документами")

        except Exception as e:
            logger.error(f"Ошибка при инициализации BM25: {str(e)}")
            raise
        finally:
            conn.close()

    def _create_fallback_index(self):
        """Создаем минимальный резервный индекс"""
        logger.warning("Создание резервного индекса BM25")
        
        if not self.corpus:
            test_docs = [
                "метрология это наука об измерениях",
                "государственный эталон единицы измерения",
                "поверка средств измерений",
                "метрологическое обеспечение",
                "измерительные приборы"
            ]
            self.corpus = [self._preprocess_text(doc) for doc in test_docs]
            self.corpus = [doc for doc in self.corpus if doc]

        if not self.corpus:
            logger.error("Не удалось создать даже тестовый корпус")
            self.corpus = [["пусто"]]
            self.doc_ids = [0]
        else:
            self.doc_ids = list(range(len(self.corpus)))
        
        try:
            self.bm25 = BM25Okapi(self.corpus)
            logger.info(f"Резервный индекс создан с {len(self.corpus)} документами")
        except Exception as e:
            logger.error(f"Ошибка создания резервного индекса: {str(e)}")
            self.corpus = [["пусто"]]
            self.doc_ids = [0]
            self.bm25 = BM25Okapi(self.corpus)

    def _preprocess_text(self, text):
        """Улучшенная обработка текста с запасным вариантом"""
        try:
            if not text or not isinstance(text, str):
                return []
                
            text = re.sub(r"[^\w\s\-']", " ", text.lower())
            
            try:
                tokens = word_tokenize(text, language='russian')
            except Exception as e:
                logger.warning(f"Ошибка NLTK токенизации: {str(e)}")
                tokens = text.split()
            
            return [
                token for token in tokens 
                if token not in self.stop_words 
                and len(token) > 2
                and not token.isdigit()
            ]
        except Exception as e:
            logger.warning(f"Ошибка обработки текста: {str(e)}")
            return [t for t in text.lower().split() if len(t) > 2]


    def search(self, query, top_k=5):
        """Поиск с помощью BM25"""
        if not self.bm25:
            logger.error("BM25 не инициализирован!")
            return []

        try:
            tokens = self._preprocess_text(query)
            if not tokens:
                logger.warning("Запрос не содержит значимых токенов")
                return []

            scores = self.bm25.get_scores(tokens)
            top_indices = np.argsort(scores)[-top_k:][::-1]

            results = []
            conn = sqlite3.connect(self.db_path)
            conn.row_factory = sqlite3.Row
            cursor = conn.cursor()

            for idx in top_indices:
                if scores[idx] <= 0:
                    continue
                    
                doc_id = self.doc_ids[idx]
                cursor.execute("""
                    SELECT c.chunk_text, d.doc_type_short, d.doc_number, d.file_name
                    FROM content c
                    JOIN documents d ON c.document_id = d.id
                    WHERE c.id = ?
                """, (doc_id,))
                
                if row := cursor.fetchone():
                    source = " ".join(filter(None, [
                        str(row['doc_type_short']) if row['doc_type_short'] else None,
                        str(row['doc_number']) if row['doc_number'] else None,
                        str(row['file_name']) if row['file_name'] else None
                    ])) or "Неизвестный источник"
                    
                    results.append({
                        "text": row['chunk_text'],
                        "source": source,
                        "score": float(scores[idx]),
                        "type": "bm25"
                    })

            conn.close()
            return results

        except Exception as e:
            logger.error(f"Ошибка поиска BM25: {str(e)}")
            return []

# Подключение к SQLite базе
def get_db_connection(db_path):
    try:
        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        return conn
    except Exception as e:
        logger.error(f"Ошибка подключения к базе данных: {e}")
        raise

# Векторный поиск
def vector_search(question, top_k=5, threshold=0.3):
    global model, faiss_index
    
    if model is None or faiss_index is None:
        logger.warning("Модель или FAISS индекс не загружены")
        return []
        
    try:
        question_embedding = model.encode([question])
        question_embedding = question_embedding.astype('float32')
        distances, indices = faiss_index.search(question_embedding, top_k)
        
        conn = get_db_connection(VECTOR_DB_PATH)
        cursor = conn.cursor()
        
        results = []
        for distance, faiss_id in zip(distances[0], indices[0]):
            similarity = 1 - distance
            
            if similarity < threshold:
                continue
                
            cursor.execute("SELECT chunk_id FROM map WHERE faiss_id = ?", (int(faiss_id),))
            map_result = cursor.fetchone()
            
            if not map_result:
                continue
                
            chunk_id = map_result['chunk_id']
            
            cursor.execute("""
                SELECT c.chunk_text, d.doc_type_short, d.doc_number, d.file_name
                FROM content c
                JOIN documents d ON c.document_id = d.id
                WHERE c.id = ?
            """, (chunk_id,))
            chunk_result = cursor.fetchone()
            
            if chunk_result:
                chunk_text = chunk_result['chunk_text']
                source_parts = [
                    str(chunk_result['doc_type_short']) if chunk_result['doc_type_short'] else None,
                    str(chunk_result['doc_number']) if chunk_result['doc_number'] else None,
                    str(chunk_result['file_name']) if chunk_result['file_name'] else None
                ]
                source = " ".join(filter(None, source_parts)) or "Неизвестный источник"
                
                results.append({
                    "text": chunk_text,
                    "source": source,
                    "score": float(similarity),
                    "type": "vector"
                })
        
        conn.close()
        return results
    
    except Exception as e:
        logger.error(f"Ошибка векторного поиска: {e}")
        return []

# Гибридный поиск
def hybrid_search_results(question, top_k=5):
    vector_results = vector_search(question, top_k=top_k*2)
    bm25_results = hybrid_search.search(question, top_k=top_k*2) if hybrid_search else []
    
    # Объединяем результаты
    all_results = vector_results + bm25_results
    
    if not all_results:
        logger.warning("Не найдено результатов ни одним методом поиска")
        return []
    
    try:
        # Нормализуем оценки отдельно для каждого метода
        vector_scores = [r['score'] for r in all_results if r['type'] == 'vector']
        bm25_scores = [r['score'] for r in all_results if r['type'] == 'bm25']
        
        max_vector_score = max(vector_scores) if vector_scores else 1
        max_bm25_score = max(bm25_scores) if bm25_scores else 1
        
        # Нормализация и комбинирование оценок
        for result in all_results:
            if result['type'] == 'vector':
                result['normalized_score'] = result['score'] / max_vector_score
                result['combined_score'] = 0.7 * result['normalized_score']  # Больший вес для векторного поиска
            else:
                result['normalized_score'] = result['score'] / max_bm25_score
                result['combined_score'] = 0.3 * result['normalized_score']
        
        # Сортируем по комбинированной оценке
        all_results.sort(key=lambda x: x['combined_score'], reverse=True)
        
        # Удаляем дубликаты, сохраняя лучшие оценки
        unique_results = []
        seen_texts = set()
        
        for result in all_results:
            text_hash = hash(result['text'])
            if text_hash not in seen_texts:
                seen_texts.add(text_hash)
                unique_results.append(result)
                if len(unique_results) >= top_k:
                    break
        
        logger.info(f"Найдено результатов: vector={len(vector_results)}, bm25={len(bm25_results)}")
        logger.info(f"После дедупликации: {len(unique_results)}")
        
        return unique_results
    
    except Exception as e:
        logger.error(f"Ошибка в гибридном поиске: {str(e)}")
        return all_results[:top_k] if all_results else []

# Загрузка данных из XLSX
@st.cache_data
def load_data():
    try:
        return pd.read_excel(XLSX_FILE_PATH)
    except Exception as e:
        logger.error(f"Ошибка загрузки XLSX файла: {e}")
        return pd.DataFrame()

# Загрузка моделей
@st.cache_data
def load_models():
    """Загрузка моделей с расширенной проверкой"""
    try:
        logger.info("="*80)
        logger.info(f"Начало загрузки модели: {EMBEDDING_MODEL}")
        
        # Добавляем определение start_time
        start_time = time.time()
        
        model = SentenceTransformer(
            EMBEDDING_MODEL,
            cache_folder=os.path.expanduser("~/.cache/huggingface/hub")
        )
        model = model.to('cpu')  # Сначала явно переносим на CPU
        
        # Проверяем работоспособность
        test_text = ["тестовый текст"]
        with torch.no_grad():
            embeddings = model.encode(test_text)
        
        logger.info(f"Модель загружена за {time.time()-start_time:.2f} сек")
        logger.info(f"Размерность эмбеддингов: {model.get_sentence_embedding_dimension()}")

        
        # 2. Загрузка FAISS индекса
        logger.info(f"Загрузка FAISS индекса: {FAISS_INDEX_PATH}")
        if not os.path.exists(FAISS_INDEX_PATH):
            error_msg = f"Индекс не найден: {FAISS_INDEX_PATH}"
            logger.error(error_msg)
            raise FileNotFoundError(error_msg)
        
        faiss_index = faiss.read_index(FAISS_INDEX_PATH)
        logger.info(f"Индекс загружен (размерность: {faiss_index.d}, векторов: {faiss_index.ntotal})")
        
        # 3. Инициализация гибридного поиска
        logger.info(f"Инициализация гибридного поиска: {VECTOR_DB_PATH}")
        
        # Проверка существования файла БД для BM25
        if not os.path.exists(VECTOR_DB_PATH):
            logger.error(f"Файл базы данных для BM25 не найден: {VECTOR_DB_PATH}")
            st.error(f"Файл базы данных для BM25 не найден: {VECTOR_DB_PATH}")
            return model, faiss_index, None
        
        # Проверка размера файла БД
        db_size = os.path.getsize(VECTOR_DB_PATH)
        logger.info(f"Размер файла БД: {db_size} байт")
        if db_size == 0:
            logger.error("Файл базы данных пуст!")
            st.error("Файл базы данных пуст!")
            return model, faiss_index, None
        
        hybrid_search = HybridSearch(VECTOR_DB_PATH)
        
        if hybrid_search and hybrid_search.bm25:
            logger.info(f"BM25 успешно инициализирован! Документов: {len(hybrid_search.corpus)}")
            
            
        else:
            logger.error("Не удалось инициализировать BM25!")
            st.error("Не удалось инициализировать текстовый поиск (BM25)")
            
        return model, faiss_index, hybrid_search

    except Exception as e:
        logger.critical(f"Фатальная ошибка при загрузке: {str(e)}", exc_info=True)
        st.error("""
        Критическая ошибка инициализации системы. Проверьте:
        1. Наличие всех файлов данных
        2. Логи в model_loading.log
        3. Доступ к интернету для загрузки моделей
        """)
        return None, None, None

# Загружаем модели с логированием
logger.info("="*80)
logger.info("Начинается процесс загрузки всех моделей")
try:
    model, faiss_index, hybrid_search = load_models()

    if model is None:
        logger.critical("Не удалось загрузить SentenceTransformer модель!")
        st.error("❌ Не удалось загрузить модель для векторного поиска")
        st.stop()
    if faiss_index is None:
        logger.critical("Не удалось загрузить FAISS индекс!")
        st.error("❌ Не удалось загрузить индекс FAISS")
        st.stop()
    if hybrid_search is None:
        logger.critical("Не удалось инициализировать гибридный поиск!")
        st.error("❌ Не удалось инициализировать гибридный поиск")
        st.stop()

    logger.info("Все модели успешно загружены")
    
except Exception as e:
    logger.critical(f"Критическая ошибка при загрузке моделей: {str(e)}")
    st.error("❌ Критическая ошибка при инициализации системы")
    st.stop()

# Генерация ответа с помощью GPT
def generate_gpt_response(question, context_chunks):
    try:
        # Формируем контекст для модели
        context = "\n\n".join([f"Фрагмент {i+1}:\n{chunk['text']}\nИсточник: {chunk['source']}" 
                             for i, chunk in enumerate(context_chunks)])
        
        prompt = f"""
        Ты - ассистент-эксперт по неразрушающему контролю, который помогает находить ответы на вопросы в технической документации.
        
        ВАЖНО: 
        1. Отвечай ТОЛЬКО на вопросы, касающиеся неразрушающего контроля и связанных с ним тем (метрология, измерения, 
           контроль качества, техническая диагностика, стандарты и нормативные документы в этой области).
        
        2. Анализируй понятность вопроса:
           - Если вопрос содержит неясные сокращения или термины - попроси уточнения
           - Если вопрос слишком общий или неконкретный - попроси детализации
           - Если вопрос четкий и понятный - давай прямой ответ из документов
        
        3. При ответе:
           - Если в документах есть прямой ответ - используй его
           - Если информации недостаточно - укажи это
           - Не проси уточнений, если ответ очевиден из контекста
        
        Пользователь задал вопрос: "{question}"
        
        Ниже приведены релевантные фрагменты из документов:
        {context}
        
        Сформулируй четкий и структурированный ответ, основываясь на предоставленных фрагментах.
        Не указывай источники в конце ответа, они будут добавлены автоматически.

        Ответ:
        """
        
        response = openai.ChatCompletion.create(
            model="openai/gpt-4.1-nano",
            messages=[{"role": "system", "content": prompt}],
            temperature=0.2,
            max_tokens=1000
        )
        
        return response.choices[0].message['content'].strip()
    
    except Exception as e:
        logger.error(f"Ошибка при генерации ответа GPT: {e}")
        return "Не удалось сгенерировать ответ. Пожалуйста, попробуйте другой вопрос."

# Логирование
def save_log(question, answer):
    log_entry = {
        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "question": question,
        "answer": answer
    }
    st.session_state.logs.append(log_entry)
    
    try:
        with open(LOG_FILE, "a", encoding="utf-8") as f:
            json.dump(log_entry, f, ensure_ascii=False)
            f.write("\n")
    except Exception as e:
        logger.error(f"Ошибка при сохранении лога: {e}")

# Поиск ответа
def get_answer(question):
    # Получаем все релевантные результаты
    results = []
    
    # 1. Проверка в базе данных
    if "метролог" in question.lower():
        conn = get_db_connection(SQLITE_DB_PATH)
        cursor = conn.cursor()
        cursor.execute("""
            SELECT c.chunk_text, d.doc_type_short, d.doc_number, d.file_name
            FROM content c
            JOIN documents d ON c.document_id = d.id
            WHERE c.id = 20
        """)
        result = cursor.fetchone()
        conn.close()
        
        if result:
            results.append({
                "text": result['chunk_text'],
                "source": f"{result['doc_type_short'] or '?'} {result['doc_number'] or ''} {result['file_name'] or ''}".strip(),
                "score": 1.0,
                "type": "exact"
            })
    
    # 2. Поиск в Excel
    qa_df = load_data()
    excel_responses = []
    excel_sources = []
    
    for _, row in qa_df.iterrows():
        table_question = str(row['Вопрос']).lower()
        if fuzz.partial_ratio(question.lower(), table_question) > 85:
            response = re.sub(r"^[a-zA-Zа-яА-Я]$\s*", "", str(row['Правильный ответ']))
            source = str(row['Источник ответа']) if pd.notna(row['Источник ответа']) else "?"
            excel_responses.append(response)
            excel_sources.append(source)
    
    if excel_responses:
        results.append({
            "text": ", ".join(set(excel_responses)),
            "source": ", ".join([s for s in set(excel_sources) if s != '?']),
            "score": 1.0,
            "type": "excel"
        })
    
    # 3. Гибридный поиск
    hybrid_results = hybrid_search_results(question)
    if hybrid_results:
        results.extend(hybrid_results)
    
    # Если есть результаты, генерируем ответ с помощью GPT
    if results:
        try:
            gpt_answer = generate_gpt_response(question, results)
            
            # Формируем полный ответ
            answer = f"🤖 Ответ:\n\n{gpt_answer}\n\n"
            
            # Собираем уникальные источники
            unique_sources = list(set(res['source'] for res in results))
            
            if unique_sources:
                answer += "📚 Использованные источники:\n"
                for source in unique_sources:
                    answer += f"- {source}\n"
            
            save_log(question, answer)
            return answer
                
        except Exception as e:
            logger.error(f"Ошибка при генерации ответа GPT: {str(e)}")
            
    # 4. Если не удалось сгенерировать ответ через GPT, возвращаем обычный поиск
    if results:
        answer = "Найдены следующие релевантные фрагменты:\n\n"
        for idx, res in enumerate(results, 1):
            answer += f"### Фрагмент {idx}\n"
            answer += f"{res['text']}\n"
            answer += f"\n📚 Источник: {res['source']}\n\n"
        
        save_log(question, answer)
        return answer
    
    # 5. Ответ по умолчанию
    answer = "К сожалению, не удалось найти точный ответ. Попробуйте переформулировать вопрос."
    save_log(question, answer)
    return answer

# Интерфейс Streamlit
st.markdown(
    """
    <style>
    .stApp {
        background-color: #f0f2f6;
        padding: 15px;
    }
    .stButton>button {
        background-color: #4CAF50 !important;
        color: white !important;
        border: none !important;
        border-radius: 12px !important;
        padding: 10px 20px !important;
        transition: all 0.3s !important;
    }
    .stButton>button:hover {
        background-color: #45a049 !important;
        transform: scale(1.02);
    }
    .stTextInput>div>div>input,
    .stTextArea>div>div>textarea {
        border: 2px solid #4CAF50 !important;
        border-radius: 12px !important;
        padding: 10px !important;
    }
    .chunk-box {
        background-color: #ffffff;
        border: 1px solid #dddddd;
        border-radius: 8px;
        padding: 15px;
        margin-bottom: 15px;
    }
    .vector-result {
        background-color: #f8f9fa;
        border-left: 4px solid #4285f4;
    }
    .bm25-result {
        background-color: #f8f9fa;
        border-left: 4px solid #34a853;
    }
    .gpt-response {
        background-color: #e8f5e9;
        border-left: 4px solid #2e7d32;
        padding: 15px;
        margin-bottom: 20px;
        border-radius: 8px;
    }
    .sidebar-content {
        background-color: #f8f9fa;
    }
    .sidebar .sidebar-content {
        padding: 1rem;
    }
    .stExpander {
        border: 1px solid #e0e0e0;
        border-radius: 8px;
        margin-bottom: 1rem;
    }
    </style>
    """,
    unsafe_allow_html=True
)

try:
    st.image("logo.png", width=150)
except FileNotFoundError:
    st.warning("Файл logo.png не найден")

st.sidebar.markdown("### Документы для поиска")
st.sidebar.markdown("Этот помощник ответит на вопросы по следующим документам:")

# Получаем список документов
documents = get_documents_list()

# Создаем expander для списка документов
with st.sidebar.expander("Показать/скрыть список документов", expanded=False):
    if documents:
        for doc in documents:
            st.markdown(f"- {doc}")
    else:
        st.warning("Не удалось загрузить список документов")


with st.sidebar.expander("Инструкция", expanded=False):
    st.markdown("""
    ### Как использовать:
    1. Введите ваш вопрос в текстовое поле
    2. Нажмите кнопку "Найти ответ"
    3. Просмотрите найденные ответы.
    """)
st.title("🔍 Поиск в технической документации")



def submit():
    st.session_state.user_input = st.session_state.widget
    st.session_state.widget = ''

st.text_area("Введите ваш вопрос:", height=100, key="widget", on_change=submit)

if st.button("Найти ответ"):
    if not st.session_state.user_input.strip():
        st.warning("Пожалуйста, введите вопрос")
    else:
        with st.spinner("Ищем релевантные фрагменты и генерируем ответ..."):
            answer = get_answer(st.session_state.user_input)
            
            st.session_state.chat_history.append({
                "question": st.session_state.user_input,
                "answer": answer
            })
            
            st.markdown(f"### Вопрос:\n{st.session_state.user_input}")
            
            if "🤖 Сгенерированный ответ:" in answer:
                # Разбираем ответ на части
                gpt_part = answer.split("🤖 Сгенерированный ответ:")[1].split("🔍 Использованные фрагменты документов:")[0]
                chunks_part = answer.split("🔍 Использованные фрагменты документов:")[1]
                
                # Отображаем сгенерированный ответ
                st.markdown('<div class="gpt-response">' + gpt_part + '</div>', unsafe_allow_html=True)
                
                # Отображаем использованные фрагменты
                st.success("Использованные фрагменты документов:")
                parts = chunks_part.split("### Фрагмент")[1:]
                for part in parts:
                    chunk_num, rest = part.split("\n", 1)
                    chunk_text, source = rest.split("📚 Источник:", 1)
                    
                    # Определяем класс CSS в зависимости от типа поиска
                    search_type = "vector" if "векторный" in answer else "bm25" if "BM25" in answer else "hybrid"
                    css_class = f"{search_type}-result"
                    
                    with st.container():
                        st.markdown(f"#### Фрагмент {chunk_num.strip()}")
                        if "оценка:" in chunk_num:
                            score = re.search(r"оценка: ([\d.]+)", chunk_num)
                            if score:
                                st.caption(f"Оценка: {score.group(1)}")
                        st.markdown(f'<div class="chunk-box {css_class}">{chunk_text.strip()}</div>', unsafe_allow_html=True)
                        st.markdown(f"**Источник:** {source.strip()}")
            elif "### Фрагмент" in answer:
                st.success("Найдены релевантные фрагменты!")
                parts = answer.split("### Фрагмент")[1:]
                for part in parts:
                    chunk_num, rest = part.split("\n", 1)
                    chunk_text, source = rest.split("📚 Источник:", 1)
                    
                    # Определяем класс CSS в зависимости от типа поиска
                    search_type = "vector" if "векторный" in answer else "bm25" if "BM25" in answer else "hybrid"
                    css_class = f"{search_type}-result"
                    
                    with st.container():
                        st.markdown(f"#### Фрагмент {chunk_num.strip()}")
                        if "оценка:" in chunk_num:
                            score = re.search(r"оценка: ([\d.]+)", chunk_num)
                            if score:
                                st.caption(f"Оценка: {score.group(1)}")
                        st.markdown(f'<div class="chunk-box {css_class}">{chunk_text.strip()}</div>', unsafe_allow_html=True)
                        st.markdown(f"**Источник:** {source.strip()}")
            else:
                st.markdown(f"### Ответ:\n{answer}")
            
            st.session_state.user_input = ""

if st.checkbox("Показать историю запросов"):
    st.subheader("История поиска")
    try:
        with open(LOG_FILE, "r", encoding="utf-8") as f:
            logs = [json.loads(line) for line in f.readlines()]
            for log in reversed(logs[-5:]):
                with st.expander(f"{log['timestamp']}: {log['question']}"):
                    st.markdown(log["answer"])
    except FileNotFoundError:
        st.warning("Логи пока не созданы")
    except Exception as e:
        st.warning(f"Ошибка при загрузке логов: {e}")