import os
from telegram import Update
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from src.rag_core import get_retriever, get_llm

TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
ALLOWED_USERS = [int(u) for u in os.getenv("ALLOWED_TELEGRAM_USERS", "").split(",") if u]

retriever, _, _ = get_retriever()
llm = get_llm()
prompt = ChatPromptTemplate.from_messages([
    ("system", "Nutze folgenden Kontext für die Antwort:\n\n{context}"),
    ("human", "{input}"),
])
chain = create_retrieval_chain(retriever, create_stuff_documents_chain(llm, prompt))

def check_auth(update: Update) -> bool:
    return update.effective_user.id in ALLOWED_USERS

async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if not check_auth(update):
        await update.message.reply_text("Zugriff verweigert.")
        return
    await update.message.reply_text("Hallo! Ich bin dein DocuMind-G4 RAG-Bot. Frag mich etwas.")

async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
    if not check_auth(update):
        return
    
    question = update.message.text
    # "Denke nach" Indikator
    await context.bot.send_chat_action(chat_id=update.effective_chat.id, action='typing')
    
    try:
        response = chain.invoke({"input": question})
        answer = response["answer"]
        sources = list({doc.metadata.get("source", "Unbekannt") for doc in response.get("context", [])})
        
        reply = f"{answer}\n\n**Quellen:**\n" + "\n".join(f"- {s}" for s in sources)
        await update.message.reply_text(reply, parse_mode='Markdown')
    except Exception as e:
        await update.message.reply_text(f"Es gab einen Fehler: {str(e)}")

def main():
    app = Application.builder().token(TOKEN).build()
    app.add_handler(CommandHandler("start", start))
    app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, handle_message))
    print("Telegram Bot gestartet...")
    app.run_polling()

if __name__ == "__main__":
    main()