from flask import Flask, render_template, request, jsonify, session
import google.generativeai as genai
import os
from dotenv import load_dotenv
import http.client
import json
from werkzeug.utils import secure_filename
import tempfile

app = Flask(__name__)
app.secret_key = os.urandom(24)  # For session management
load_dotenv()

# Configure Google AI
genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))

# Safety settings
safety_settings = [
    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
]

# System prompt
SYSTEM_PROMPT = """
# Prompt System pour Mariam, IA conçu par youssouf 
[Your existing system prompt content here]
"""

# Initialize Gemini model
model = genai.GenerativeModel('gemini-2.0-flash-exp', 
                            tools='code_execution',
                            safety_settings=safety_settings,
                            system_instruction=SYSTEM_PROMPT)

def perform_web_search(query):
    conn = http.client.HTTPSConnection("google.serper.dev")
    payload = json.dumps({"q": query})
    headers = {
        'X-API-KEY': '9b90a274d9e704ff5b21c0367f9ae1161779b573',
        'Content-Type': 'application/json'
    }
    try:
        conn.request("POST", "/search", payload, headers)
        res = conn.getresponse()
        data = json.loads(res.read().decode("utf-8"))
        return data
    except Exception as e:
        print(f"Web search error: {e}")
        return None
    finally:
        conn.close()

def format_search_results(data):
    if not data:
        return "Aucun résultat trouvé"
    
    result = []
    
    if 'knowledgeGraph' in data:
        kg = data['knowledgeGraph']
        result.append({
            'type': 'knowledge',
            'title': kg.get('title', ''),
            'description': kg.get('description', ''),
            'category': kg.get('type', '')
        })
    
    if 'organic' in data:
        for item in data['organic'][:3]:
            result.append({
                'type': 'organic',
                'title': item['title'],
                'snippet': item['snippet'],
                'link': item['link']
            })
    
    if 'peopleAlsoAsk' in data:
        for item in data['peopleAlsoAsk'][:2]:
            result.append({
                'type': 'question',
                'question': item['question'],
                'answer': item['snippet']
            })
    
    return result

UPLOAD_FOLDER = 'temp'
ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png', 'pdf', 'txt', 'mp3', 'mp4'}

os.makedirs(UPLOAD_FOLDER, exist_ok=True)

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/')
def home():
    if 'chat_history' not in session:
        session['chat_history'] = []
    return render_template('index.html', messages=session['chat_history'])

@app.route('/send_message', methods=['POST'])
def send_message():
    try:
        data = request.json
        message = data.get('message')
        web_search_enabled = data.get('web_search', False)
        
        if not message:
            return jsonify({'error': 'No message provided'}), 400

        # Initialize chat if not in session
        if 'chat' not in session:
            session['chat'] = model.start_chat(history=[])

        # Perform web search if enabled
        web_results = None
        if web_search_enabled:
            web_results = perform_web_search(message)
            if web_results:
                formatted_results = format_search_results(web_results)
                message = f"""Question: {message}\n\nRésultats de recherche web:\n{formatted_results}\n\nPourrais-tu analyser ces informations et me donner une réponse complète?"""

        # Send message to Gemini
        response = session['chat'].send_message(message)
        
        # Update chat history
        session['chat_history'].append({
            'role': 'user',
            'content': message
        })
        session['chat_history'].append({
            'role': 'assistant',
            'content': response.text
        })
        
        return jsonify({
            'response': response.text,
            'web_results': web_results
        })

    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': 'No selected file'}), 400
    
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(UPLOAD_FOLDER, filename)
        file.save(filepath)
        
        try:
            gemini_file = genai.upload_file(filepath)
            return jsonify({'success': True, 'filename': filename})
        except Exception as e:
            return jsonify({'error': str(e)}), 500
    
    return jsonify({'error': 'Invalid file type'}), 400

@app.route('/clear_chat', methods=['POST'])
def clear_chat():
    session.pop('chat_history', None)
    session.pop('chat', None)
    return jsonify({'success': True})

if __name__ == '__main__':
    app.run(debug=True)