import os
import google.generativeai as genai
from flask import Flask, request, jsonify
from flask_cors import CORS
from google.generativeai import protos

app = Flask(__name__)
CORS(app)

# --- Konfiguration & Gemini Initialisierung ---
GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
if not GOOGLE_API_KEY:
    raise ValueError("GOOGLE_API_KEY nicht gesetzt! Bitte als HF Space Secret hinterlegen.")

genai.configure(api_key=GOOGLE_API_KEY)

# Korrekte Tool-Initialisierung für Google Search
search_tool = protos.Tool(
    google_search_retrieval=protos.Tool.GoogleSearchRetrieval(
        disable_attribution=False,
        max_return_results=3
    )
)

model = genai.GenerativeModel(
    'gemini-1.5-pro',
    tools=[search_tool],
    system_instruction="Du bist Moejra, eine hilfsbereite KI-Lernbegleitung. Antworte im unterstützenden Stil und nutze Webrecherche bei Unsicherheiten. Generiere bei Bildanfragen explizit ein Bild."
)

chat = model.start_chat(history=[])

# --- Wissensdatenbank (RAG) Funktion ---
def retrieve_info_from_kb(query):
    relevant_text = []
    knowledge_base_dir = "knowledge_base"
    
    if not os.path.exists(knowledge_base_dir):
        return ""

    query_keywords = [word.lower() for word in query.split() if len(word) > 2]

    for filename in os.listdir(knowledge_base_dir):
        if filename.endswith(".txt"):
            filepath = os.path.join(knowledge_base_dir, filename)
            try:
                with open(filepath, "r", encoding="utf-8") as f:
                    content = f.read()
                    if any(keyword in content.lower() for keyword in query_keywords):
                        relevant_text.append(content)
            except Exception as e:
                print(f"Dateilesefehler: {e}")

    return "\n---\n".join(relevant_text) if relevant_text else ""

# --- API-Endpunkt für den Chat ---
@app.route('/chat', methods=['POST'])
def handle_chat():
    user_input = request.json.get('text')
    if not user_input:
        return jsonify({"error": "Kein Text übermittelt"}), 400

    try:
        # RAG-Phase
        kb_info = retrieve_info_from_kb(user_input)
        
        # Prompt-Konstruktion
        prompt_parts = [
            f"Nutzeranfrage: {user_input}",
            f"Kontext aus Wissensdatenbank:\n{kb_info}" if kb_info else "Kein relevanter Kontext gefunden",
            "Antworte im unterstützenden Moejra-Stil. Bei Unsicherheiten recherchiere im Web."
        ]

        # Gemini-Abfrage
        response = chat.send_message("\n".join(prompt_parts))
        
        # Antwortverarbeitung
        response_text = ""
        image_data = []
        
        for part in response.candidates[0].content.parts:
            if part.text:
                response_text += part.text
            if hasattr(part, 'image'):
                img_data = part.image.data
                if hasattr(img_data, 'mime_type') and hasattr(img_data, 'data'):
                    image_data.append({
                        "mime": img_data.mime_type,
                        "data": f"data:{img_data.mime_type};base64,{img_data.data}"
                    })

        return jsonify({
            "text": response_text.strip(),
            "images": image_data
        })

    except Exception as e:
        print(f"API-Fehler: {str(e)}")
        return jsonify({
            "text": "😞 Entschuldige, ein unerwarteter Fehler ist aufgetreten. Bitte versuche es später erneut.",
            "images": []
        }), 500

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)