import os
import google.generativeai as genai
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.core import Settings
from llama_index.core.llms import ChatMessage, MessageRole
import os

EMBEDDING_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
RETRIEVER_TOP_K = 10
RETRIEVER_SIMILARITY_CUTOFF = 0.7
RAG_FILES_DIR = "processed_data"
PROCESSED_DATA_FILE = "processed_data/processed_chunks.csv"

UPLOAD_FOLDER = "UPLOADED_DOCUMENTS"
PROCESSED_DATA_FILE = "processed_data/processed_chunks.csv"
INDEX_STATE_FILE = "processed_data/index_store.json"
RAG_FILES_DIR = "processed_data"

GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY', "AIzaSyDemsCp7JIdRNDRyP6DkYdMox1DLZwPcPE")
LLM_MODEL = "gemini-2.0-flash"

CHUNK_SIZE = 1024
CHUNK_OVERLAP = 256
MAX_CHUNK_SIZE = 2048
MIN_CHUNK_SIZE = 750
SIMILARITY_THRESHOLD = 0.7

RETRIEVER_TOP_K = 15
RETRIEVER_SIMILARITY_CUTOFF = 0.7

def setup_llm_settings():
    # Set embedding model first
    Settings.embed_model = HuggingFaceEmbedding(model_name=EMBEDDING_MODEL)
    
    # Only set LLM if API key is available
    if GOOGLE_API_KEY:
        try:
            llm = GoogleGenAI(model=LLM_MODEL, api_key=GOOGLE_API_KEY)
            Settings.llm = llm
        except Exception as e:
            print(f"Warning: Could not initialize Google GenAI LLM: {e}")
            Settings.llm = None
    else:
        print("Warning: GOOGLE_API_KEY not found. Setting LLM to None.")
        Settings.llm = None

# Intent Classification Prompt for preprocessing
INTENT_CLASSIFIER_PROMPT = """Проанализируй запрос пользователя и определи, к какой из следующих категорий он относится:

1. SUMMARY - Запрос на краткое изложение, саммари пункта, раздела или документа
   Примеры: "кратко расскажи о разделе 5", "саммари главы про безопасность", "что говорится в пункте 3.2"

2. SEARCH - Поиск конкретной информации, документа или пункта
   Примеры: "где найти информацию о...", "какой документ регулирует...", "в каком пункте говорится про..."

3. VALIDATION - Проверка корректности данных, соответствия требованиям
   Примеры: "правильно ли...", "соответствует ли требованиям...", "проверь корректность данных..."

4. ACTION_PLAN - Составление плана действий на основе документации
   Примеры: "составь план...", "какие шаги нужно предпринять...", "алгоритм действий для..."

Запрос пользователя: {user_query}

Ответь ТОЛЬКО одним словом: SUMMARY, SEARCH, VALIDATION или ACTION_PLAN"""

# Task-specific prompts
SUMMARY_PROMPT = """Ты специализированный помощник по анализу нормативной документации (НД) для создания кратких изложений и саммари.

Твоя задача: Предоставить краткое, структурированное изложение запрашиваемой информации из нормативных документов.

Правила создания саммари:
1. Обязательное указание источника: Всегда указывай источник информации в формате:
   - Для конкретного раздела/подраздела: Согласно разделу [X] и подразделу [X.X]: [Саммари]
   - Для общих пунктов: Согласно [Название документа] - [Номер и наименование пункта]: [Саммари]

2. Структурированность: Организуй информацию логично:
   - Основные положения
   - Ключевые требования
   - Важные исключения или особенности

3. Краткость и полнота: Включи все существенные моменты, но избегай избыточных деталей

4. Сохранение терминологии НД: Используй точную терминологию из документов

5. Если информация отсутствует: Информация для создания саммари по вашему запросу не найдена в нормативной документации.

Контекст: {context_str}

Запрос: {query_str}

Саммари:"""

SEARCH_PROMPT = """Ты специализированный помощник по поиску информации в нормативной документации (НД).

Твоя задача: Найти и предоставить точную информацию, документы или пункты согласно запросу пользователя.

Правила поиска информации:
1. Точное указание источника: Обязательно укажи точное местоположение информации:
   - Для конкретного раздела/подраздела: Согласно разделу [X] и подразделу [X.X]: [Найденная информация]
   - Для общих пунктов: Согласно [Название документа] - [Номер и наименование пункта]: [Найденная информация]

2. Релевантность: Предоставляй именно ту информацию, которая отвечает на запрос

3. Полнота ответа: Если найдено несколько релевантных источников, укажи все

4. Точность цитирования: Используй формулировки непосредственно из НД

5. Если информация не найдена: Информация по вашему запросу не была найдена в нормативной документации.

Контекст: {context_str}

Поисковый запрос: {query_str}

Результат поиска:"""

VALIDATION_PROMPT = """Ты специализированный помощник по проверке соответствия требованиям нормативной документации (НД).

Твоя задача: Проанализировать предоставленные данные или информацию на соответствие требованиям НД и дать заключение о корректности.

Правила проведения проверки:
1. Источник требований: Четко укажи, согласно каким пунктам НД проводится проверка:
   - Согласно разделу [X] и подразделу [X.X]: [Требование]
   - Согласно [Название документа] - [Номер и наименование пункта]: [Требование]

2. Структура заключения:
   - СООТВЕТСТВУЕТ/НЕ СООТВЕТСТВУЕТ требованиям
   - Обоснование заключения со ссылками на конкретные пункты НД
   - При несоответствии - указание на конкретные нарушения
   - Рекомендации по устранению несоответствий (если применимо)

3. Объективность: Основывайся исключительно на требованиях НД, без субъективных оценок

4. Если требования не найдены: Требования для проверки данной информации не найдены в нормативной документации.

Контекст: {context_str}

Запрос на проверку: {query_str}

Заключение о соответствии:"""

ACTION_PLAN_PROMPT = """Ты специализированный помощник по составлению планов действий на основе нормативной документации (НД).

Твоя задача: Создать пошаговый план действий, основанный на требованиях и процедурах, описанных в НД.

Правила составления плана действий:
1. Источник каждого шага: Для каждого шага указывай источник в НД:
   - Согласно разделу [X] и подразделу [X.X]: Шаг [N]
   - Согласно [Название документа] - [Номер и наименование пункта]: Шаг [N]

2. Структура плана:
   - Нумерованные шаги в логической последовательности
   - Четкие действия для каждого шага
   - Указание ответственных лиц/органов (если указано в НД)
   - Временные рамки (если указаны в НД)
   - Необходимые документы/материалы

3. Практичность: План должен быть выполнимым и соответствовать реальным процедурам

4. Полнота: Включи все обязательные этапы согласно НД

5. Если информация недостаточна: Недостаточно информации в нормативной документации для составления полного плана действий по вашему запросу.

Контекст: {context_str}

Запрос на план: {query_str}

План действий:"""

# Dictionary mapping intents to prompts
TASK_PROMPTS = {
    "SUMMARY": SUMMARY_PROMPT,
    "SEARCH": SEARCH_PROMPT,
    "VALIDATION": VALIDATION_PROMPT,
    "ACTION_PLAN": ACTION_PLAN_PROMPT
}

LLM_MODEL_PREPROCESS = "gemini-1.5-flash"

def classify_intent(user_query, llm=None):
    """Classify user intent to determine appropriate prompt"""
    if not llm:
        llm = GoogleGenAI(model=LLM_MODEL_PREPROCESS, temperature=0.1)
    
    try:
        intent_prompt = INTENT_CLASSIFIER_PROMPT.format(user_query=user_query)
        messages = [ChatMessage(role=MessageRole.USER, content=intent_prompt)]
        response = llm.chat(messages)
        intent = response.message.content.strip()
        
        # Validate intent and return default if invalid
        if intent in TASK_PROMPTS:
            return intent
        else:
            print(f"Unknown intent: {intent}, defaulting to SEARCH")
            return "SEARCH"
    except Exception as e:
        print(f"Intent classification failed: {e}, defaulting to SEARCH")
        return "SEARCH"

def get_task_specific_prompt(intent, context_str, query_str):
    """Get the appropriate prompt template based on intent"""
    prompt_template = TASK_PROMPTS.get(intent, SEARCH_PROMPT)
    return prompt_template.format(context_str=context_str, query_str=query_str)

def preprocess_query_with_context(user_query, chat_history=None, llm=None):
    if not chat_history:
        return user_query
        
    if not llm:
        llm = GoogleGenAI(model=LLM_MODEL_PREPROCESS, temperature=0.1)
    
    # Format chat history into a string for the prompt
    history_context = "\n".join([
        f"User: {item['user']}\nAssistant: {item['assistant']}"
        for item in chat_history[-3:]  # Consider only the last 3 exchanges for conciseness
    ])
    
    preprocessing_prompt = f"""Analyze the user's current question in the context of their chat history and improve it for better document retrieval.

Chat History:
{history_context}

Current Question: {user_query}

Tasks:
1. If the question refers to previous context, make it self-contained.
2. Add relevant keywords that would help find documents.
3. Maintain the legal/regulatory focus.
4. Keep it concise but specific.

Return ONLY the improved question:"""
    
    try:
        messages = [ChatMessage(role=MessageRole.USER, content=preprocessing_prompt)]
        response = llm.chat(messages)
        improved_query = response.message.content.strip()
        
        # Fallback to the original query if the preprocessing fails or provides an overly long response
        if len(improved_query) > len(user_query) * 3 or not improved_query:
            return user_query
            
        return improved_query
    except Exception as e:
        print(f"Query preprocessing failed: {e}")
        return user_query

def create_chat_context_prompt(base_response, chat_history=None):
    if not chat_history:
        return base_response
    
    base_aware_response = base_response
    
    if len(chat_history) > 0:
        last_exchange = chat_history[-1]
        if any(keyword in last_exchange['user'].lower() for keyword in ['закон', 'кодекс', 'статья']):
            # Add a conversational prefix
            base_aware_response = f"Продолжая тему нормативных документов: {base_response}"
    
    return base_aware_response

# Example usage function
def process_user_query(user_query, context_str, chat_history=None):
    # Step 1: Classify intent
    intent = classify_intent(user_query)
    print(f"Detected intent: {intent}")
    
    # Step 2: Preprocess query if needed
    processed_query = preprocess_query_with_context(user_query, chat_history)
    
    # Step 3: Get appropriate prompt
    task_prompt = get_task_specific_prompt(intent, context_str, processed_query)
    
    return task_prompt, intent