main.py

from dotenv import load_dotenv
import os
# Carga inicial y forzada de variables de entorno
load_dotenv(dotenv_path=os.path.join(os.path.dirname(os.path.abspath(__file__)), '.env'), override=True)

import logging
import asyncio
from telegram import Update
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes
from openai import AzureOpenAI
from langchain_openai import AzureOpenAIEmbeddings
from knowledge.data_ingestion import KnowledgeBase
from langchain_chroma import Chroma
from chromadb.config import Settings
import sys
import json
from utils.logger import setup_logger
from knowledge.config_manager import ConfigManager
import re
from typing import List, Dict
from langchain_core.documents import Document

logger = setup_logger()

config_manager = ConfigManager()
bot_config = config_manager.get_bot_config()


def disable_chromadb_telemetry():
    try:
        import posthog

        posthog.disabled = True
        posthog.capture = lambda *args, **kwargs: None
    except Exception:
        pass

# Configuración de variables de Azure OpenAI Chat
AZURE_API_KEY = os.getenv("AZURE_OPENAI_API_KEY") or bot_config.get("azure", {}).get("chat", {}).get("api_key")
AZURE_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT") or bot_config.get("azure", {}).get("chat", {}).get("endpoint")
AZURE_API_VERSION = bot_config.get("azure", {}).get("chat", {}).get("api_version")
AZURE_DEPLOYMENT = bot_config.get("azure", {}).get("chat", {}).get("deployment")

if not AZURE_API_KEY:
    raise ValueError("AZURE_OPENAI_API_KEY not found in configuration or environment variables")
if not AZURE_ENDPOINT:
    raise ValueError("AZURE_OPENAI_ENDPOINT not found in configuration or environment variables")
if not AZURE_API_VERSION:
    raise ValueError("Azure API version not configured")
if not AZURE_DEPLOYMENT:
    raise ValueError("Azure deployment not configured")

client = AzureOpenAI(
    api_version=AZURE_API_VERSION,
    azure_endpoint=AZURE_ENDPOINT,
    api_key=AZURE_API_KEY
)


# Configuración explícita de Embeddings (Identica a la usada en build_knowledge_base.py)
embedding_function = AzureOpenAIEmbeddings(
    azure_deployment=bot_config.get("azure", {}).get("embeddings", {}).get("azure_deployment"),
    azure_endpoint=bot_config.get("azure", {}).get("embeddings", {}).get("endpoint") or os.getenv("AZURE_OPENAI_EMBEDDING_ENDPOINT"),
    api_key=bot_config.get("azure", {}).get("embeddings", {}).get("api_key") or os.getenv("AZURE_OPENAI_EMBEDDING_API_KEY"),
    api_version=bot_config.get("azure", {}).get("embeddings", {}).get("api_version")
)

knowledge_base = KnowledgeBase(
    chunk_size=500,
    chunk_overlap=50
)

disable_chromadb_telemetry()

# Inicialización segura de Chroma con los parámetros de persistencia correctos
try:
    knowledge_base.db = Chroma(
        persist_directory="./knowledge_base",
        embedding_function=embedding_function,
        client_settings=Settings(
            anonymized_telemetry=False,
            allow_reset=True,
            is_persistent=True
        )
    )
    if knowledge_base.db is None:
        raise RuntimeError("Chroma DB initialization returned None")
    count = knowledge_base.db._collection.count()
    logger.info(f"Database loaded successfully. Total documents: {count}")
except Exception as e:
    logger.error(f"Critical error loading Chroma DB: {e}")
    raise

knowledge_base.cache.set_query_function(knowledge_base._raw_query_knowledge)

active_conversations = {}

def load_agent_config():
    """Loads the agent configuration from the configuration file"""
    try:
        bot_config = config_manager.get_bot_config()
        instructions = format_agent_instructions(bot_config['instructions'])
        return {
            'name': bot_config['name'],
            'model': bot_config['model'],
            'instructions': instructions,
            'welcome_message': bot_config['welcome_message'],
            'question_template': bot_config['instructions']['question_template']
        }
    except Exception as e:
        logging.error(f"Error loading agent configuration: {e}")
        sys.exit(1)

def format_agent_instructions(config):
    """Formats the agent instructions"""
    sections = []
    sections.append(f"INSTRUCTIONS FOR THE ROLE: {config['role']}")
    sections.append("\nFORMATTING AND TONE:")
    for key, value in config.get('formatting', {}).items():
        sections.append(f"• {key.replace('_', ' ').title()}: {value}")
    sections.append("\nPROHIBITED:")
    sections.extend([f"• {item}" for item in config.get('prohibited', [])])
    if config.get('knowledge_limitations'):
        sections.append("\nKNOWLEDGE LIMITATIONS:")
        sections.append(config['knowledge_limitations'])
    if config.get('response_guidelines'):
        sections.append("\nRESPONSE GUIDELINES:")
        sections.extend([f"• {item}" for item in config['response_guidelines']])
    return "\n".join(sections)

agent_config = load_agent_config()

bot_agent = {
    'name': agent_config['name'],
    'instructions': agent_config['instructions'],
    'model': agent_config['model'],
    'question_template': agent_config['question_template']
}

async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    await update.message.reply_text(agent_config['welcome_message'])

async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Handles incoming messages."""
    try:
        if not update.message or not update.message.text:
            return

        message_text = update.message.text
        logger.info(f"Message from {update.effective_user.username}: {message_text}")

        is_private = update.message.chat.type == 'private'
        is_bot_mentioned = bool(update.message.entities and
            any(entity.type == 'mention' and
                context.bot.username in message_text[entity.offset:entity.offset + entity.length]
                for entity in update.message.entities))
        is_reply_to_bot = bool(update.message.reply_to_message and
            update.message.reply_to_message.from_user.id == context.bot.id)

        # Validación estricta para evitar procesar mensajes de grupos donde no se le invoca
        if not (is_private or is_bot_mentioned or is_reply_to_bot):
            logger.info("Message skipped: Not private, not mentioned, and not a reply to bot.")
            return

        chat_id = update.message.chat_id

        await context.bot.send_chat_action(chat_id=chat_id, action="typing")

        if chat_id not in active_conversations:
            active_conversations[chat_id] = []

        active_conversations[chat_id].append({
            "role": "user",
            "content": message_text
        })

        # Búsqueda vectorial local (rápida)
        relevant_info = knowledge_base.query_knowledge(message_text)
        context_text = "\n".join([doc.page_content for doc in relevant_info])

        system_prompt = get_optimized_prompt(message_text, context_text, bot_agent['instructions'])

        conversation = [{"role": "system", "content": system_prompt}]

        if len(active_conversations[chat_id]) > 1:
            conversation.extend(active_conversations[chat_id][-4:])

        conversation.append({
            "role": "user",
            "content": bot_agent['question_template'].format(question=message_text)
        })

        typing_task = asyncio.create_task(keep_typing(context.bot, chat_id))

        try:
            # Use the model defined in bot_config.json if provided, otherwise fall back to the environment deployment.
            model_to_use = bot_agent.get('model') or AZURE_DEPLOYMENT
            logger.info(f"Using model/deployment: {model_to_use}")

            # CRITICAL FIX: Run the blocking synchronous call in a separate thread (asyncio.to_thread)
            # This prevents the Telegram event loop from becoming unresponsive while Azure processes the request
            response = await asyncio.to_thread(
                client.chat.completions.create,
                model=model_to_use,
                messages=conversation,
                stream=False
            )

            bot_response = response.choices[0].message.content
            active_conversations[chat_id].append({
                "role": "assistant",
                "content": bot_response
            })

        finally:
            typing_task.cancel()
            try:
                await typing_task
            except asyncio.CancelledError:
                pass

        logger.info(f"Answer to {update.effective_user.username}: {bot_response}")

        await update.message.reply_text(
            bot_response,
            reply_to_message_id=update.message.message_id,
            parse_mode='Markdown'
        )

    except Exception as e:
        error_msg = f"Error processing message: {str(e)}"
        logger.error(error_msg)
        await update.message.reply_text("Sorry, there was an error processing your message.")

async def keep_typing(bot, chat_id):
    try:
        while True:
            await bot.send_chat_action(chat_id=chat_id, action="typing")
            await asyncio.sleep(3)
    except asyncio.CancelledError:
        pass

def process_context(knowledge_base, message_text):
    try:
        relevant_info = knowledge_base.query_knowledge(message_text)

        if not relevant_info:
            return ""

        context_parts = [doc.page_content for doc in relevant_info]
        return "\n\nRelevant Context:\n" + "\n---\n".join(context_parts)

    except Exception as e:
        logger.error(f"Error processing context: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        return ""

async def shutdown():
    """Graceful shutdown function for the bot"""
    if 'app' in globals() and app.is_running():
        await app.shutdown()
    print("Bot stopped gracefully")

def main():
    """Initialize and run the bot"""
    try:
        global app
        token = os.getenv('TELEGRAM_BOT_TOKEN')
        if not token:
            logging.error("TELEGRAM_BOT_TOKEN not found in environment variables")
            sys.exit(1)
            
        app = Application.builder().token(token).build()

        app.add_handler(CommandHandler("start", start))
        app.add_handler(MessageHandler(
            filters.TEXT & ~filters.COMMAND,
            handle_message
        ))

        print("Starting bot...")
        app.run_polling(allowed_updates=Update.ALL_TYPES)

    except Exception as e:
        logging.error(f"Critical error: {e}")
        sys.exit(1)

def _calculate_match_score(message_words, doc):
    """Calculate a relevance score for a given document"""
    keywords = set(doc.metadata.get('keywords', []))
    topics = set(doc.metadata.get('topics', []))
    return len(message_words & (keywords | topics)) * 2 or 1

def _detect_query_intent(message_text, config_file='config/keywords.json'):
    """Detect the intent of the query based on keywords"""
    with open(config_file, 'r') as f:
        keywords_config = json.load(f)

    query_lower = message_text.lower().replace('?', '')
    intents = {}
    for category, data in keywords_config.items():
        keywords = data.get('keywords', [])
        if any(kw.lower() in query_lower for kw in keywords):
            intents[category] = True

    return intents

def get_optimized_prompt(message_text, context_text, bot_instructions):
    """Optimizes the prompt for the Azure model"""
    max_context_length = 2000
    if len(context_text) > max_context_length:
        context_text = context_text[:max_context_length] + "..."

    return (
        f"{bot_instructions}\n\n"
        f"RELEVANT CONTEXT:\n{context_text}\n\n"
        f"USER QUESTION: {message_text}"
    )

def _build_prompt(self, user_message: str, context_docs: List[Document]) -> List[Dict[str, str]]:
    """Build the prompt for the AI model with context from knowledge base"""
    instructions = self.bot_config.get("instructions", {})

    source_info = ""
    if context_docs:
        source_names = set()
        for doc in context_docs:
            if 'filename' in doc.metadata and doc.metadata['filename']:
                source_names.add(doc.metadata['filename'])
        if source_names:
            source_info = f"\nInformation obtained from the following sources: {', '.join(source_names)}"

    system_message = f"{instructions['role']}\n\n"

    if 'formatting' in instructions:
        system_message += "FORMAT:\n"
        for key, value in instructions['formatting'].items():
            system_message += f"- {key}: {value}\n"

    if 'prohibited' in instructions:
        system_message += "\nPROHIBITED ACTIONS:\n"
        for item in instructions['prohibited']:
            system_message += f"- {item}\n"

    if 'knowledge_limitations' in instructions:
        system_message += f"\n{instructions['knowledge_limitations']}\n"

    context_text = ""
    if context_docs:
        context_text = "\nRELEVANT CONTEXT:\n"
        for i, doc in enumerate(context_docs):
            source_name = doc.metadata.get('filename', 'unknown source')
            context_text += f"\n--- Fragment #{i+1} of {source_name} ---\n"
            context_text += doc.page_content.strip() + "\n"

    logger.info(f"Building prompt with {len(context_docs)} context documents")

    messages = [{"role": "system", "content": system_message + context_text}]
    user_query = f"{user_message}{source_info}"
    messages.append({"role": "user", "content": user_query})
    return messages

if __name__ == "__main__":
    main()