from flask import Flask, request, jsonify, send_file
from pymongo import MongoClient
from pydantic import BaseModel
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
from flask_cors import CORS
from functools import wraps
import datetime
import uvicorn
import bcrypt
import jwt
import os
import re
import uuid
from gtts import gTTS
from bson import ObjectId

# app = FastAPI()
app = Flask(__name__)

CORS(app, resources={r"/*": {"origins": ["https://bhagvat-ras.static.hf.space", "http://localhost:7860"]}})

SECRET_KEY = os.getenv("VERA_API_KEY")

# Rate limiter (create it before using @limiter.limit)
limiter = Limiter(get_remote_address, app=app, default_limits=["200/hour"])


# MongoDB connection
client = MongoClient("mongodb+srv://veratheai24:veratheai24@cluster0.b15rk44.mongodb.net/?retryWrites=true&w=majority&appName=Cluster0")

# Database connections
app_info_db = client['app_info_db']
user_db = client['user_db']
users_collection = user_db['users']  # Collection for user data
chat_db = client['user_chat_db']
chat_collection = chat_db['chats']  # Single collection for all chat history
shastra_db = client['shastra']
donation_db = client['donation_db']
donation_collection = donation_db['donations']


# Token-required decorator
def token_required(f):
    @wraps(f)
    def decorated(*args, **kwargs):
        token = request.headers.get('Authorization')  # Expecting 'Bearer <JWT>'
        if not token:
            return {"message": "Token is missing!"}, 403

        try:
            # Remove 'Bearer' from the token string if present
            token = token.split(" ")[1] if " " in token else token
            decoded_token = jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
        except jwt.ExpiredSignatureError:
            return {"message": "Token has expired!"}, 403
        except jwt.InvalidTokenError:
            return {"message": "Invalid token!"}, 403

        # Store decoded token data in kwargs for the route function
        kwargs['user_info'] = decoded_token
        return f(*args, **kwargs)
    return decorated


# Function to fetch all chapters from 'shastra' database
@app.route("/get_chapters", methods=["GET"])
def get_chapters():
    try:
        # List all collections in the 'shastra' database
        collections = shastra_db.list_collection_names()
        data = {}

        # Fetch data from each collection and store it in 'data'
        for collection_name in collections:
            collection_data = list(shastra_db[collection_name].find({}, {"_id": 0}))  # Exclude '_id' from results
            data[collection_name] = collection_data

        # Debugging output for verification
        # print("Fetched collections:", collections)
        # print("Fetched data sample:", data)

        return jsonify({"success": True, "data": data}), 200

    except Exception as e:
        error_message = f"Error fetching chapters: {e}"
        print(error_message)
        return jsonify({"success": False, "message": error_message}), 500


@app.route("/get_chapter_details", methods=["GET"])
def get_chapter_details():
    try:
        # Get the query parameter for chapterTitle
        chapter_title = request.args.get("chapterTitle")
        print(f"Requested chapterTitle: {chapter_title}")

        if not chapter_title:
            return jsonify({"success": False, "message": "chapterTitle parameter is required."}), 400

        # Fetch the chapter details from the 'Bhagwat Charcha' collection
        collection = shastra_db["Bhagwat Mahapuran"]

        # Use aggregation to search for the chapter inside the nested 'chapters' array
        result = collection.aggregate([
            {"$unwind": "$chapters"},  # Unwind the chapters array
            {"$match": {"chapters.chapterTitle": chapter_title}},  # Match the specific chapter
            {"$project": {"_id": 0, "chapterDetails": "$chapters"}}  # Project only the matching chapter
        ])

        # Get the first result from the aggregation
        chapter_details = next(result, None)

        if not chapter_details:
            return jsonify({"success": False, "message": "Chapter not found."}), 404

        # Return the chapter details
        return jsonify({"success": True, "data": chapter_details["chapterDetails"]}), 200

    except Exception as e:
        error_message = f"Error fetching chapter details: {e}"
        print(error_message)
        return jsonify({"success": False, "message": error_message}), 500


@app.route("/get_chapter_titles", methods=["GET"])
def get_chapter_titles():
    try:
        # List all collections in the 'shastra' database
        collections = shastra_db.list_collection_names()
        data = {}

        # Iterate over each collection to gather chapter titles
        for collection_name in collections:
            collection_data = list(shastra_db[collection_name].find({}, {"_id": 0, "chapterTitle": 1, "skandaTitle": 1, "chapters": 1}))
            data[collection_name] = []

            for entry in collection_data:
                if "chapters" in entry and isinstance(entry["chapters"], list):
                    # Handle grouped chapters under a Skanda
                    skanda_entry = {
                        "skandaTitle": entry.get("skandaTitle", ""),
                        "chapters": [
                            {"chapterTitle": chapter["chapterTitle"]} for chapter in entry["chapters"]
                        ]
                    }
                    data[collection_name].append(skanda_entry)
                else:
                    # Handle standalone chapters
                    data[collection_name].append({
                        "chapterTitle": entry.get("chapterTitle", "")
                    })

        return jsonify({"success": True, "data": data}), 200

    except Exception as e:
        error_message = f"Error fetching chapter titles: {e}"
        print(error_message)
        return jsonify({"success": False, "message": error_message}), 500
       


# Function to fetch the audio URL for a given sloke
def get_audio_from_db(sloke_text):
    try:
        # Access the collection
        collection = shastra_db["Bhagwat Mahapuran"]  # Replace with your collection name
        print(f"Searching for sloke: {sloke_text}")  # Debugging log
        
        # Use aggregation to search for the chapter containing the sloke
        result = collection.aggregate([
            {"$unwind": "$chapters"},  # Unwind the chapters array
            {"$match": {"chapters.slokes": {"$regex": sloke_text}}},
            # {"$match": {"chapters.slokes": {"$regex": r"^(श्लोक\s*\d+:)?\s*" + re.escape(sloke_text)}}},  # Match the sloke text with optional prefix
            {"$project": {
                "_id": 0,
                "slokes": "$chapters.slokes",
                "audioLinks": "$chapters.audioLinks"
            }}  # Project slokes and audioLinks
        ])
        
        # Get the first result from the aggregation
        chapter_details = next(result, None)
        if not chapter_details:
            return None  # No matching sloke found

        # Split the slokes into individual slokes
        slokes = chapter_details["slokes"].split("\n")  # Assuming slokes are newline-separated
        audio_links = chapter_details.get("audioLinks", [])  # Get the audio links array
        # print(f"Audio Links: {audio_links}")  # Debugging log
        # Extract the slok number from the provided text
        slok_number_match = re.search(r"श्लोक\s*(\d+)", sloke_text)
        if not slok_number_match:
            print(f"Slok number not found in text: {sloke_text}")
            return None
        
        slok_number = int(slok_number_match.group(1))  # Extracted slok number
        print(f"Slok Number: {slok_number}")  # Debugging log
        # The slok number corresponds to the index in the audioLinks array (0-based index)
        if 1 <= slok_number <= len(audio_links):
            # Return the audio link for the slok number
            return audio_links[slok_number - 1]
        else:
            print(f"Audio link for slok number {slok_number} not found")
            return None

    except Exception as e:
        print(f"Error fetching audio URL: {e}")
        return None
        

def convert_drive_link(shareable_link):
    if "drive.google.com" in shareable_link and "/file/d/" in shareable_link:
        file_id = shareable_link.split('/d/')[1].split('/')[0]
        return f"https://drive.google.com/uc?export=download&id={file_id}"
    return shareable_link


@app.route('/get_audio_url', methods=['GET'])
def get_audio_url():
    sloke_text = request.args.get('slokeText')  # Use slokeText as the parameter
    if not sloke_text:
        return jsonify({'success': False, 'message': 'Sloke text is required'}), 400
    
    # Fetch the audio URL from MongoDB
    audio_url = get_audio_from_db(sloke_text)

    if audio_url:
        # audio_url = convert_drive_link(audio_url)
        return jsonify({'success': True, 'audioUrl': audio_url})
    else:
        return jsonify({'success': False, 'message': 'Audio not found'}), 404


    
# Function to register a user and store their chat history in the same collection
@app.route("/register_user", methods=["POST"])
def register_user():
    data = request.get_json()
    name = data.get("name")
    username = data.get("username")
    email = data.get("email")
    password = data.get("password")
    device_id = data.get("device_id")

    # Hash the password for security
    hashed_password = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt())

    # Check if the user already exists in MongoDB
    existing_user = users_collection.find_one({"email": email})
    if existing_user:
        return {"message": "User already registered!"}, 409  # Use 409 Conflict

    # Register the user in the 'users' collection
    user_data = {
        "name": name,
        "username": username,
        "email": email,
        "password": hashed_password,
        "device_id": device_id,
        "registered_on": datetime.datetime.now()
    }
    users_collection.insert_one(user_data)

    # Generate JWT token
    token = jwt.encode({
        "email": email,
        "exp": datetime.datetime.utcnow() + datetime.timedelta(days=14)  # Expire in 14 days
    }, SECRET_KEY, algorithm="HS256")

    return {"success": True, "message": f"User {username} registered!", "token": token}, 201


@app.route('/login', methods=['POST'])
def login():
    data = request.get_json()
    email = data.get('username')
    password = data.get('password')

    # Fetch the user from the MongoDB database by email
    user = users_collection.find_one({"email": email})
    
    if not user or not bcrypt.checkpw(password.encode('utf-8'), user['password']):
        return jsonify({"message": "Invalid credentials!"}), 401

    # Generate a JWT token
    token = jwt.encode({
        "email": email,
        "exp": datetime.datetime.utcnow() + datetime.timedelta(days=14)  # Token expires in 14 days
    }, SECRET_KEY, algorithm="HS256")

    return jsonify({"success": True, "message": "Login successful!", "token": token}), 200


# Function to store chat history for unregistered users
@app.route('/store_anonymous_chat', methods=["POST"])
def store_anonymous_chat():
    data = request.get_json()
    user_prompt = data.get('user_data')
    bot_prompt = data.get('bot_data')

    bot_data_list = bot_prompt.split('\n')
    category = bot_data_list[0] if bot_data_list else "General"
    bot_data = '\n'.join(bot_data_list[1:])
    
    anonymous_chat_entry = {
        "user_type": "anonymous",  # Flag to identify unregistered users
        "user_data": user_prompt,
        "bot_data": bot_data,
        "category": category,
        "timestamp": datetime.datetime.now()
    }

    # Save chat asynchronously using Celery
    chat_collection.insert_one(anonymous_chat_entry)
    
    # chat_collection.insert_one(anonymous_chat_entry)
    return {"message": "Anonymous chat saved successfully!"}, 201


# Function to store chat history for registered users
@app.route('/store_registered_chat/<username>', methods=["POST"])
def store_registered_chat(username: str):
    data = request.get_json()
    user_prompt = data.get('user_data')
    bot_prompt = data.get('bot_data')

    # Check if the user exists
    user_exists = users_collection.find_one({"email": username})
    if not user_exists:
        return {"message": "User not found!"}, 404

    bot_data_list = bot_prompt.split('\n')
    category = bot_data_list[0] if bot_data_list else "General"
    bot_data = '\n'.join(bot_data_list[1:])
    # Store chat in the unified chat collection with a user identifier
    chat_entry = {
        "user_type": "registered",
        "username": username,
        "user_data": user_prompt,
        "bot_data": bot_data,
        "category": category,
        "timestamp": datetime.datetime.now()
    }

    # Save chat asynchronously using Celery
    chat_collection.insert_one(chat_entry)
    
    # chat_collection.insert_one(chat_entry)
    return {"message": "Chat saved successfully for registered user!"}, 201

    
# Function to register a user from google and store their chat history in the same collection 
@app.route("/register_google_user", methods=["POST"])
def register_google_user():
    data = request.get_json()
    name = data.get("name")
    username = data.get("username")
    email = data.get("email")
    google_id = data.get("google_id")
    device_id = data.get("device_id")

    # Check if the user already exists in MongoDB
    existing_user = users_collection.find_one({"email": email})
    if existing_user:
        # Generate JWT token
        token = jwt.encode({
            "email": email,
            "exp": datetime.datetime.utcnow() + datetime.timedelta(days=14)  # Expire in 14 days
        }, SECRET_KEY, algorithm="HS256")
        
        return {"success": True, "message": "User already registered!", "token": token}, 201

    # Register the user in the 'users' collection
    user_data = {
        "name": name,
        "username": username,
        "email": email,
        "google_id": google_id,
        "device_id": device_id,
        "registered_on": datetime.datetime.now()
    }
    users_collection.insert_one(user_data)

    # Generate JWT token
    token = jwt.encode({
        "email": email,
        "exp": datetime.datetime.utcnow() + datetime.timedelta(days=14)  # Expire in 14 days
    }, SECRET_KEY, algorithm="HS256")

    return {"success": True, "message": f"User {username} registered!", "token": token}, 201
    


# Function to fetch chat history for a specific user (registered or anonymous)
@app.route('/get_chat_history/<username>', methods=["GET"])
def get_chat_history(username: str):
    # Fetch chat history for registered user
    chat_history = chat_collection.find({"username": username})
    
    if chat_history.count() == 0:
        return jsonify({"message": "No chat history found!"}), 404

    return jsonify([{"user_data": chat["user_data"], "bot_data": chat["bot_data"], "timestamp": chat["timestamp"]} for chat in chat_history]), 200



# Protected route example
@app.route('/protected', methods=['GET'])
@token_required  # Apply the token_required decorator
def protected():
    return {"message": "This is a protected route."}, 200


@app.route('/validate_token', methods=['GET'])
@token_required
def validate_token(user_info):
    return {
        "message": "Token is valid!",
        "email": user_info["email"]  # Return the user info extracted from the token
    }, 200


@app.route('/generate_sanskrit_audio', methods=['POST'])
@limiter.limit("30/minute")
def generate_sanskrit_audio():
    data = request.get_json()
    text = data.get('cleanText', '')

    # Check if text is provided
    if not text:
        return {"error": "No text provided"}, 400

    try:
        # Generate audio using gTTS
        audio_file = f'output_{uuid.uuid4()}.mp3'  # Unique file name
        tts = gTTS(text=text, lang='hi', slow=False)  # 'sa' is for Sanskrit
        tts.save(audio_file)

        # Send the audio file back
        response = send_file(audio_file, mimetype='audio/mp3', as_attachment=True)

        # Cleanup: remove the audio file after sending
        os.remove(audio_file)

        return response

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

@app.route('/save_donation', methods=["POST"])
def save_donation():
    """
    Save donation details to the MongoDB collection.
    """
    data = request.get_json()

    # Extract user-provided data
    name = data.get('name')
    email = data.get('email')
    amount = data.get('amount')
    payment_method = data.get('payment_method', 'UPI')  # Default payment method is UPI

    if not name or not email or not amount:
        return {"error": "Name, email, and amount are required fields!"}, 400

    # Create the donation entry
    donation_entry = {
        "user_type": "anonymous",  # Assuming donations are anonymous by default
        "name": name,
        "email": email,
        "amount": amount,
        "payment_method": payment_method,
        "timestamp": datetime.datetime.now()
    }

    try:
        # Insert the donation entry into the MongoDB collection
        donation_collection.insert_one(donation_entry)
        return {"message": "Donation saved successfully!"}, 201
    except Exception as e:
        return {"error": f"An error occurred: {str(e)}"}, 500


@app.get('/chats/<username>')
def list_chats(username):
    items = list(chat_collection.find({'username': username}, {'_id':1,'user_data':1,'bot_data':1,'timestamp':1,'category':1}).sort('timestamp', -1))
    out = []
    for it in items:
        out.append({
            'id': str(it['_id']),
            'title': (it.get('user_data') or 'Chat')[:40],
            'timestamp': it.get('timestamp'),
            'category': it.get('category','General')
        })
    return jsonify(out), 200

@app.patch('/chat/<chat_id>/rename')
def rename_chat(chat_id):
    title = (request.get_json() or {}).get('title')
    if not title: return jsonify({'error':'title required'}), 400
    chat_collection.update_one({'_id': ObjectId(chat_id)}, {'$set': {'title': title}})
    return jsonify({'ok': True}), 200

@app.delete('/chat/<chat_id>')
def delete_chat(chat_id):
    chat_collection.delete_one({'_id': ObjectId(chat_id)})
    return jsonify({'ok': True}), 200

# Simple memory K/V
mem = client['user_memory_db']
mem_col = mem['memory']

@app.post('/memory/<username>')
def save_memory(username):
    data = request.get_json() or {}
    key, value = data.get('key'), data.get('value')
    if not key: return jsonify({'error':'key required'}), 400
    mem_col.update_one({'username': username, 'key': key},
                       {'$set': {'value': value, 'updated': datetime.datetime.now()}},
                       upsert=True)
    return jsonify({'ok': True}), 200

@app.get('/memory/<username>')
def get_memory(username):
    items = list(mem_col.find({'username': username}, {'_id':0}))
    return jsonify(items), 200

@app.delete('/memory/<username>/<key>')
def delete_memory(username, key):
    mem_col.delete_one({'username': username, 'key': key})
    return jsonify({'ok': True}), 200

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