app.py

from flask import Flask, request, jsonify, Response, render_template
from flask_cors import CORS
import os
import logging
import functools
import pandas as pd
import threading
import time
import tempfile
import shutil
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Custom Imports
from rag_system import initialize_and_get_rag_system
from config import (
    API_USERNAME, API_PASSWORD, RAG_SOURCES_DIR, RAG_STORAGE_PARENT_DIR,
    FAQ_SOURCES_DIR, FAQ_STORAGE_PARENT_DIR, FAQ_CSV_FILENAME,
    FAQ_INITIAL_FETCH_K, FAQ_FINAL_K, FAQ_CONFIDENCE_THRESHOLD,
    GDRIVE_INDEX_ENABLED, GDRIVE_INDEX_ID_OR_URL,
    GDRIVE_USERS_CSV_ENABLED, GDRIVE_USERS_CSV_ID_OR_URL,
    ADMIN_USERNAME, ADMIN_PASSWORD, RAG_RERANKER_K, RAG_INITIAL_FETCH_K,
    EXTERNAL_URL, URL_UPDATE_PERIOD_MINUTES, URL_FETCH_ENABLED,
    RAG_CSV_INITIAL_FETCH_K, RAG_CSV_FINAL_K,
    RAG_CSV_MAX_RESULTS, RAG_CSV_CONFIDENCE_THRESHOLD
)
from utils import download_and_unzip_gdrive_file, download_gdrive_file, fetch_and_clean_url

# Logging Setup
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Flask Init
app = Flask(__name__, static_folder='static', template_folder='templates')
CORS(app)

# Global State
rag_system = None
faq_rag_system = None
user_df = None
_APP_BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# --- Helper: Load Users ---
def load_users_from_csv():
    global user_df
    assets_folder = os.path.join(_APP_BASE_DIR, 'assets')
    os.makedirs(assets_folder, exist_ok=True)
    users_csv_path = os.path.join(assets_folder, 'users.csv')
    
    try:
        if os.path.exists(users_csv_path):
            user_df = pd.read_csv(users_csv_path)
            # Normalize email
            if 'email' in user_df.columns:
                user_df['email'] = user_df['email'].str.lower().str.strip()
            logger.info(f"Loaded {len(user_df)} users from CSV.")
        else:
            logger.warning("users.csv not found in assets folder.")
            user_df = None
    except Exception as e:
        logger.error(f"Failed to load users.csv: {e}")
        user_df = None

# --- Helper: Auth Decorators ---
def require_api_auth(f):
    """Protects the N8N Webhook endpoint"""
    @functools.wraps(f)
    def decorated(*args, **kwargs):
        auth = request.authorization
        if not auth or auth.username != API_USERNAME or auth.password != API_PASSWORD:
            return Response('Unauthorized', 401, {'WWW-Authenticate': 'Basic realm="API Login Required"'})
        return f(*args, **kwargs)
    return decorated

def require_admin_auth(f):
    """Protects Admin Rebuild/Update endpoints"""
    @functools.wraps(f)
    def decorated(*args, **kwargs):
        auth = request.authorization
        if not auth:
            return jsonify({"error": "Unauthorized"}), 401

        if user_df is not None:
            user_email = auth.username.lower().strip()
            user_record = user_df[user_df['email'] == user_email]
            if not user_record.empty:
                user_data = user_record.iloc[0]
                if str(user_data['password']) == auth.password and user_data['role'] == 'admin':
                    return f(*args, **kwargs)

        if auth.username == ADMIN_USERNAME and auth.password == ADMIN_PASSWORD:
            return f(*args, **kwargs)

        return jsonify({"error": "Unauthorized"}), 401
    return decorated

# --- URL Zero-Downtime Updater ---
def trigger_url_update():
    global rag_system
    if not URL_FETCH_ENABLED or not EXTERNAL_URL: 
        return {"error": "External URL fetching is disabled or not configured"}
    
    logger.info(f"[URL_UPDATE] Starting zero-downtime fetch from {EXTERNAL_URL}")
    
    # 1. Create temporary staging folders
    temp_staging_sources = tempfile.mkdtemp(prefix="rag_sources_temp_")
    temp_index = tempfile.mkdtemp(prefix="rag_index_temp_")
    
    try:
        # 2. COMBINE SOURCES: Copy existing GDrive/Local sources to staging first
        if os.path.exists(RAG_SOURCES_DIR):
            shutil.copytree(RAG_SOURCES_DIR, temp_staging_sources, dirs_exist_ok=True)
            
        # 3. Fetch URL data — saved to <app_root>/tmp/ for persistence & inspection
        tmp_dir = os.path.join(_APP_BASE_DIR, 'tmp')
        os.makedirs(tmp_dir, exist_ok=True)
        url_out_path = os.path.join(tmp_dir, "url_data.txt")
        success = fetch_and_clean_url(EXTERNAL_URL, url_out_path)
        
        if not success:
            return {"error": "Failed to fetch or parse the URL."}

        # Copy from tmp/ into staging so it gets indexed alongside other sources
        shutil.copy2(url_out_path, os.path.join(temp_staging_sources, "url_data.txt"))
        
        # 4. Build a brand new RAG instance isolated in the temp directories
        new_rag = initialize_and_get_rag_system(
            force_rebuild=True, 
            source_dir_override=temp_staging_sources,
            storage_dir_override=temp_index
        )
        
        if new_rag is None:
            raise Exception("Failed to build new RAG index from parsed text.")
            
        # 5. Atomic Swap (Now incoming requests hit the new DB immediately)
        rag_system = new_rag
        
        # 6. Backup/Replace persistent INDEX directory ONLY
        os.makedirs(RAG_STORAGE_PARENT_DIR, exist_ok=True)
        shutil.copytree(temp_index, RAG_STORAGE_PARENT_DIR, dirs_exist_ok=True)
        
        rag_system.index_storage_dir = RAG_STORAGE_PARENT_DIR
        
        logger.info("[URL_UPDATE] Success! RAG database updated combining Local, GDrive, and URL sources.")
        return {"status": "success", "message": "Database successfully updated using combined sources."}
        
    except Exception as e:
        logger.error(f"[URL_UPDATE] Error during update: {e}", exc_info=True)
        return {"error": str(e)}
        
    finally:
        shutil.rmtree(temp_staging_sources, ignore_errors=True)
        shutil.rmtree(temp_index, ignore_errors=True)

def url_periodic_loop():
    if not URL_FETCH_ENABLED or not EXTERNAL_URL or URL_UPDATE_PERIOD_MINUTES <= 0:
        logger.info("Periodic URL updates disabled.")
        return
        
    logger.info(f"[URL_UPDATE] Background thread started for: {EXTERNAL_URL}")
    trigger_url_update()
    
    while True:
        time.sleep(URL_UPDATE_PERIOD_MINUTES * 60)
        logger.info(f"[URL_UPDATE] Triggering scheduled periodic update...")
        trigger_url_update()

# --- Startup Logic ---
def run_startup_tasks():
    global rag_system, faq_rag_system
    logger.info("--- Executing Startup Tasks ---")
    
    if GDRIVE_USERS_CSV_ENABLED and GDRIVE_USERS_CSV_ID_OR_URL:
        target = os.path.join(_APP_BASE_DIR, 'assets', 'users.csv')
        download_gdrive_file(GDRIVE_USERS_CSV_ID_OR_URL, target)
    
    load_users_from_csv()

    if GDRIVE_INDEX_ENABLED and GDRIVE_INDEX_ID_OR_URL:
        download_and_unzip_gdrive_file(GDRIVE_INDEX_ID_OR_URL, os.getcwd())
    
    # Main knowledgebase: unchanged behavior, indexes the regular sources directory.
    rag_system = initialize_and_get_rag_system()

    # App FAQs: separate source folder and separate index/storage.
    # Put faqs.csv in sources/app_faqs/ by default, or override FAQ_SOURCES_DIR.
    faq_rag_system = initialize_and_get_rag_system(
        source_dir_override=FAQ_SOURCES_DIR,
        storage_dir_override=FAQ_STORAGE_PARENT_DIR
    )
    
    if URL_FETCH_ENABLED and EXTERNAL_URL:
        threading.Thread(target=url_periodic_loop, daemon=True).start()
        
    logger.info("--- Startup Tasks Complete ---")

with app.app_context():
    run_startup_tasks()

def _csv_row_is_inactive(content: str) -> bool:
    """Blank/missing active is active; only explicit false values are hidden."""
    inactive_values = {"0", "false", "no", "n", "inactive", "disabled", "off"}
    for line in content.splitlines():
        key, sep, value = line.partition(":")
        if sep and key.strip().lower() == "active":
            return value.strip().lower() in inactive_values
    return False


def _search_rag_system(
    target_rag,
    label,
    default_initial_fetch_k,
    default_final_k,
    csv_initial_fetch_k,
    csv_final_k,
    csv_max_results,
    csv_confidence_threshold,
    allowed_source_name=None,
):
    if not target_rag or not target_rag.retriever:
        return jsonify({"error": f"{label} RAG not initialized. Check server logs and source files."}), 503

    data = request.json or {}

    query = data.get('query')
    if not query:
        return jsonify({"error": "Query field is required"}), 400

    final_k = data.get('final_k', default_final_k)
    initial_fetch_k = data.get('initial_fetch_k', default_initial_fetch_k)
    csv_initial_fetch_k = data.get('csv_initial_fetch_k', csv_initial_fetch_k)
    csv_final_k = data.get('csv_final_k', csv_final_k)
    use_reranker = data.get('use_reranker', True)
    cleaned = data.get('cleaned', False)
    min_score = data.get('min_score')

    try:
        final_k = int(final_k)
        initial_fetch_k = int(initial_fetch_k)
        csv_initial_fetch_k = int(csv_initial_fetch_k)
        csv_final_k = int(csv_final_k)

        if final_k <= 0 or initial_fetch_k <= 0 or csv_initial_fetch_k <= 0 or csv_final_k <= 0:
            return jsonify({"error": "final_k, initial_fetch_k, csv_initial_fetch_k, and csv_final_k must be positive integers"}), 400

        if min_score is not None:
            min_score = float(min_score)

    except (TypeError, ValueError):
        return jsonify({"error": "Invalid final_k, initial_fetch_k, csv_initial_fetch_k, csv_final_k, or min_score"}), 400

    # initial_fetch_k = normal candidate retrieval.
    # csv_initial_fetch_k = wider candidate retrieval for row-by-row CSVs.
    # final_k = max total results returned.
    # csv_final_k = max CSV rows allowed inside final results.
    effective_initial_fetch_k = max(initial_fetch_k, csv_initial_fetch_k)
    candidate_k = max(effective_initial_fetch_k, final_k, csv_final_k)

    original_reranker = None
    original_initial_fetch_k = None

    if target_rag.retriever:
        original_reranker = target_rag.retriever.reranker
        original_initial_fetch_k = target_rag.retriever.initial_fetch_k

        if not use_reranker:
            target_rag.retriever.reranker = None
        elif use_reranker and target_rag.reranker:
            target_rag.retriever.reranker = target_rag.reranker

        target_rag.retriever.initial_fetch_k = effective_initial_fetch_k

    try:
        raw_results = target_rag.search_knowledge_base(query, top_k=candidate_k)

        # Normalize scores so clients can always treat higher as better.
        # With reranker enabled: score is reranker_score, already higher-is-better.
        # With reranker disabled: score is FAISS L2 distance, lower-is-better, so convert it.
        for res in raw_results:
            score = res.get("score")
            if score is None:
                continue

            if not use_reranker:
                res.setdefault("metadata", {})["raw_retrieval_score"] = score
                res["score"] = 1 / (1 + score)

        filtered_results = []
        csv_count = 0

        for res in raw_results:
            metadata = res.get("metadata", {})
            source_name = metadata.get("source_document_name", "")
            is_csv = metadata.get("source_type") == "csv" or source_name.endswith(".csv")
            score = res.get("score")

            # FAQ endpoint must only return rows from the app FAQ CSV.
            if allowed_source_name and source_name != allowed_source_name:
                continue

            # FAQ CSV can keep inactive rows for future use; they will not be returned.
            if allowed_source_name == FAQ_CSV_FILENAME and _csv_row_is_inactive(res.get("content", "")):
                continue

            # Optional per-request global score filter.
            if min_score is not None and (score is None or score < min_score):
                continue

            # Existing CSV-specific filter, scoped per endpoint.
            if is_csv:
                csv_result_cap = min(csv_max_results, csv_final_k)
                if score is not None and score >= csv_confidence_threshold and csv_count < csv_result_cap:
                    filtered_results.append(res)
                    csv_count += 1
            else:
                filtered_results.append(res)

            if len(filtered_results) >= final_k:
                break

        if cleaned:
            filtered_results = [{"content": r["content"]} for r in filtered_results]

        return jsonify({"results": filtered_results, "count": len(filtered_results), "status": "success"})

    except Exception as e:
        logger.error(f"{label} Search API Error: {e}", exc_info=True)
        return jsonify({"error": str(e)}), 500

    finally:
        if target_rag.retriever:
            target_rag.retriever.reranker = original_reranker
            target_rag.retriever.initial_fetch_k = original_initial_fetch_k


# ===========================
#       API ROUTES
# ===========================

@app.route('/webhook/search', methods=['POST'])
@require_api_auth
def search_knowledgebase_api():
    # Current regular knowledgebase endpoint: same URL, same auth, same response shape.
    return _search_rag_system(
        target_rag=rag_system,
        label="Knowledgebase",
        default_initial_fetch_k=RAG_INITIAL_FETCH_K,
        default_final_k=RAG_RERANKER_K,
        csv_initial_fetch_k=RAG_CSV_INITIAL_FETCH_K,
        csv_final_k=RAG_CSV_FINAL_K,
        csv_max_results=RAG_CSV_MAX_RESULTS,
        csv_confidence_threshold=RAG_CSV_CONFIDENCE_THRESHOLD,
    )


@app.route('/webhook/search_faqs', methods=['POST'])
@require_api_auth
def search_app_faqs_api():
    # Dedicated app FAQ endpoint: searches only sources/app_faqs/faqs.csv.
    return _search_rag_system(
        target_rag=faq_rag_system,
        label="App FAQ",
        default_initial_fetch_k=FAQ_INITIAL_FETCH_K,
        default_final_k=FAQ_FINAL_K,
        csv_initial_fetch_k=FAQ_INITIAL_FETCH_K,
        csv_final_k=FAQ_FINAL_K,
        csv_max_results=FAQ_FINAL_K,
        csv_confidence_threshold=FAQ_CONFIDENCE_THRESHOLD,
        allowed_source_name=FAQ_CSV_FILENAME,
    )

@app.route('/user-login', methods=['POST'])
def user_login():
    if user_df is None:
        return jsonify({"error": "User database not available."}), 503
        
    data = request.json
    email = data.get('email', '').lower().strip()
    password = data.get('password')
    
    if not email or not password:
        return jsonify({"error": "Email and password required"}), 400

    user_record = user_df[user_df['email'] == email]
    if not user_record.empty:
        u_data = user_record.iloc[0]
        if str(u_data['password']) == str(password):
            resp = u_data.to_dict()
            if 'password' in resp: 
                del resp['password']
            return jsonify(resp), 200
            
    return jsonify({"error": "Invalid credentials"}), 401

@app.route('/')
def index_route():
    return render_template('chat-bot.html')

@app.route('/admin/login', methods=['POST'])
@require_admin_auth
def admin_login():
    return jsonify({"status": "success", "message": "Authenticated"}), 200

@app.route('/admin/update_faiss_index', methods=['POST'])
@require_admin_auth
def update_faiss_index():
    if not rag_system: 
        return jsonify({"error": "RAG system not initialized"}), 503
    
    data = request.json or {}
    max_files = data.get('max_new_files')
    
    try:
        result = rag_system.update_index_with_new_files(RAG_SOURCES_DIR, max_files)
        return jsonify(result), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route('/admin/rebuild_index', methods=['POST'])
@require_admin_auth
def rebuild_index():
    global rag_system
    try:
        if URL_FETCH_ENABLED and EXTERNAL_URL:
            result = trigger_url_update()
            if "error" in result:
                return jsonify(result), 500
            return jsonify({"status": "Index rebuilt successfully using combined local & URL sources"}), 200
        else:
            rag_system = initialize_and_get_rag_system(force_rebuild=True)
            return jsonify({"status": "Index rebuilt successfully"}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/admin/rebuild_faq_index', methods=['POST'])
@require_admin_auth
def rebuild_faq_index():
    global faq_rag_system
    try:
        faq_rag_system = initialize_and_get_rag_system(
            force_rebuild=True,
            source_dir_override=FAQ_SOURCES_DIR,
            storage_dir_override=FAQ_STORAGE_PARENT_DIR
        )
        if not faq_rag_system or not faq_rag_system.retriever:
            return jsonify({"error": "FAQ index rebuild failed. Check sources/app_faqs/faqs.csv."}), 500
        return jsonify({"status": "FAQ index rebuilt successfully"}), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/admin/update_faq_index', methods=['POST'])
@require_admin_auth
def update_faq_index():
    if not faq_rag_system:
        return jsonify({"error": "FAQ RAG system not initialized"}), 503

    data = request.json or {}
    max_files = data.get('max_new_files')

    try:
        result = faq_rag_system.update_index_with_new_files(FAQ_SOURCES_DIR, max_files)
        return jsonify(result), 200
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# Retained specific endpoint name to ensure the frontend doesn't break
@app.route('/admin/fetch_rentry', methods=['POST'])
@require_admin_auth
def api_fetch_url():
    result = trigger_url_update()
    if "error" in result:
        return jsonify(result), 500
    return jsonify(result), 200

def _rag_ready(target_rag):
    return target_rag is not None and getattr(target_rag, "retriever", None) is not None


@app.route('/status', methods=['GET'])
def status_route():
    main_rag_ready = _rag_ready(rag_system)
    faq_rag_ready = _rag_ready(faq_rag_system)

    return jsonify({
        "status": "online",
        "rag_initialized": main_rag_ready or faq_rag_ready,
        "main_rag_initialized": main_rag_ready,
        "faq_rag_initialized": faq_rag_ready,
        "faq_sources_dir": FAQ_SOURCES_DIR,
        "faq_csv_filename": FAQ_CSV_FILENAME,
        "users_loaded": user_df is not None
    })

if __name__ == '__main__':
    port = int(os.environ.get("PORT", 7860)) 
    app.run(host='0.0.0.0', port=port)