app.py

import os
import sys
import json
import logging
from functools import wraps
from flask import Flask, request, jsonify, Response, session
from flask_cors import CORS
import secrets
import time
from collections import defaultdict

ENVIRONMENT = os.getenv('ENVIRONMENT', 'prod')

PORT_MAP = {'prod': 8001, 'dev': 8002, 'test': 8003}

current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.append(current_dir)

try:
    from rag_system_v2_2 import RAGSystem
    from config import *
except ImportError as e:
    print(f"Error importing RAGSystem or config: {e}")
    sys.exit(1)

class Config:
    HOST = os.getenv('FLASK_HOST', '0.0.0.0')
    PORT = int(os.getenv('FLASK_PORT', PORT_MAP.get(ENVIRONMENT, 8001)))
    DEBUG = os.getenv('FLASK_DEBUG', 'True').lower() == 'true'
    AUTO_LOAD_METADATA = os.getenv('AUTO_LOAD_METADATA', 'False').lower() == 'true'

# Class to watch conversation memory
class SessionMemoryManager:
    def __init__(self):
        self.sessions = {}  # {session_id: conversation_memory as a list}
        self.last_accessed = {}  # {session_id: timestamp}
        self.cleanup_interval = 1800
        self.max_inactive_time = 3600

    def get_session_id(self):
        if 'session_id' not in session:
            session['session_id'] = secrets.token_hex(16)
        return session['session_id']

    def get_memory(self, session_id):
        self.last_accessed[session_id] = time.time()
        return self.sessions.get(session_id, [])

    def set_memory(self, session_id, memory):
        self.sessions[session_id] = memory
        self.last_accessed[session_id] = time.time()

    def clear_memory(self, session_id):
        if session_id in self.sessions:
            self.sessions[session_id] = []

    def cleanup_old_sessions(self):
        current_time = time.time()
        to_remove = []

        for session_id, last_time in self.last_accessed.items():
            if current_time - last_time > self.max_inactive_time:
                to_remove.append(session_id)

        for session_id in to_remove:
            if session_id in self.sessions:
                del self.sessions[session_id]
            del self.last_accessed[session_id]

        return len(to_remove)

memory_manager = SessionMemoryManager()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = Flask(__name__)
CORS(app, supports_credentials=True, resources={r"/*": {'origins': ['https://harts.rc.ufl.edu', 'https://devharts.rc.ufl.edu', 'https://testharts.rc.ufl.edu']}})

app.config.update(
    SECRET_KEY=os.getenv('FLASK_SECRET_KEY', secrets.token_hex(32)),
    SESSION_TYPE='filesystem',
    SESSION_FILE_DIR=os.path.join(os.getcwd(), 'flask_session'),
    SESSION_PERMANENT=False,
    SESSION_USE_SIGNER=True,
    SESSION_KEY_PREFIX='rag_session:',
    SESSION_COOKIE_HTTPONLY=True,
    SESSION_COOKIE_SAMESITE='None',
    SESSION_COOKIE_SECURE=True,
    PERMANENT_SESSION_LIFETIME=3600
)
os.makedirs(os.path.join(os.getcwd(), 'flask_session'), exist_ok=True)

rag_system = None
try:
    rag_system = RAGSystem(
        memory_enabled=True,
        memory_turns=3,
        query_expansion_enabled=False,
        auto_load_metadata=Config.AUTO_LOAD_METADATA
    )
    logger.info("RAG system initialized successfully")
except Exception as e:
    logger.error(f"Failed to initialize RAG system: {e}")
    sys.exit(1)


def validate_rag_system(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if not rag_system:
            return jsonify({"error": "RAG system not initialized"}), 500
        return f(*args, **kwargs)

    return decorated_function


def validate_json_request(required_fields=None):
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            data = request.get_json()
            if not data:
                return jsonify({"error": "No data provided"}), 400

            if required_fields:
                for field in required_fields:
                    if field not in data:
                        return jsonify({"error": f"'{field}' field is missing"}), 400

            return f(data, *args, **kwargs)

        return decorated_function

    return decorator


def handle_api_error(e, operation="operation"):
    error_msg = f"Error during {operation}: {str(e)}"
    logger.error(error_msg)
    return jsonify({"error": error_msg}), 500


def format_document(doc, rank):
    return {
        "rank": rank + 1,
        "filename": doc["filename"],
        "text_preview": doc["text"][:300] + ("..." if len(doc["text"]) > 300 else ""),
        "full_text": doc["text"],
        "score": doc.get("score", 0.0),
        "metadata": doc.get("metadata", {})
    }


def build_retrieve_response(query, expanded_query, documents):
    return {
        "query": query,
        "expanded_query": expanded_query if expanded_query != query else None,
        "total_results": len(documents),
        "documents": [format_document(doc, i) for i, doc in enumerate(documents)]
    }


@app.route('/')
def home():
    return "RAG System Backend is running! Access the frontend via frontend.html."


@app.route('/stream-query', methods=['POST'])
@validate_rag_system
@validate_json_request(['query'])
def handle_streaming_query(data):

    session_id = memory_manager.get_session_id()
    session_memory = memory_manager.get_memory(session_id)
    logger.info(f"[MEMORY DEBUG] Session ID: {session_id}, Memory: {session_memory}")

    user_query = data['query']
    generate_report = data.get('generate_report', False)
    temporal_analysis = data.get('temporal_analysis', False)
    use_wos_search = data.get('use_wos_search', False)

    # Get explicit temporal parameters
    temporal_start_year = data.get('temporal_start_year')
    temporal_end_year = data.get('temporal_end_year')

    logger.info(
        f"Streaming query: '{user_query}' (Report: {generate_report}, Temporal: {temporal_analysis}, WoS: {use_wos_search})")
    if temporal_start_year and temporal_end_year:
        logger.info(f"Temporal range: {temporal_start_year}-{temporal_end_year}")

    def generate_stream():
        wos_timespan = data.get('wos_timespan', '1400-01-01+1700-12-31')
        nonlocal session_memory, session_id

        try:
            yield f"data: {json.dumps({'type': 'status', 'message': 'Processing query...'})}\n\n"

            use_temporal = (
                    (temporal_analysis and (temporal_start_year or temporal_end_year)) or
                    (temporal_start_year and temporal_end_year)
            )

            if use_temporal and hasattr(rag_system, 'temporal_analyzer'):
                yield f"data: {json.dumps({'type': 'status', 'message': 'Analyzing temporal query...'})}\n\n"

                expanded_query = rag_system.expand_query(
                    user_query) if rag_system.query_expansion_enabled else user_query

                # Use explicit years
                if temporal_start_year and temporal_end_year:
                    date_range = (temporal_start_year, temporal_end_year)
                    temporal_info = {'date_range': date_range}
                    # Update WoS timespan for temporal queries
                    wos_timespan = f"{temporal_start_year}-01-01+{temporal_end_year}-12-31"
                else:
                    temporal_info = rag_system.temporal_analyzer.parse_date_query(user_query)
                    date_range = None
                    if temporal_info['date_range']:
                        date_range = temporal_info['date_range']
                        # Update WoS timespan based on parsed temporal info
                        wos_timespan = f"{date_range[0]}-01-01+{date_range[1]}-12-31"
                    elif temporal_info['specific_year']:
                        year = temporal_info['specific_year']
                        date_range = (year, year)
                        wos_timespan = f"{year}-01-01+{year}-12-31"
                    elif temporal_info['decade']:
                        date_range = temporal_info['decade']
                        wos_timespan = f"{date_range[0]}-01-01+{date_range[1]}-12-31"
                    elif temporal_info['century']:
                        date_range = temporal_info['century']
                        wos_timespan = f"{date_range[0]}-01-01+{date_range[1]}-12-31"

                if date_range:
                    start_year, end_year = date_range
                    yield f"data: {json.dumps({'type': 'status', 'message': f'Retrieving documents from {start_year}-{end_year}...'})}\n\n"

                    retrieved_docs = rag_system.temporal_analyzer.retrieve_with_temporal_filter(
                        query=expanded_query,
                        start_year=start_year,
                        end_year=end_year,
                        limit=rag_system.search_limit
                    )
                else:
                    yield f"data: {json.dumps({'type': 'status', 'message': 'No temporal constraints found, using regular search...'})}\n\n"
                    retrieved_docs = rag_system.retrieve(query=expanded_query, limit=rag_system.search_limit)

                # Web searches
                web_results = []
                wos_results = []

                # Brave Search
                if rag_system.use_web_search and rag_system.brave_api_key:
                    yield f"data: {json.dumps({'type': 'status', 'message': 'Searching the web...'})}\n\n"
                    web_results = rag_system.web_search(user_query, limit=2)

                # Web of Science Search
                if use_wos_search and hasattr(rag_system, 'wos_api_key') and rag_system.wos_api_key:
                    yield f"data: {json.dumps({'type': 'status', 'message': 'Searching Web of Science...'})}\n\n"
                    wos_results = rag_system.web_of_science_search(
                        user_query,
                        limit=3,
                        timespan=wos_timespan
                    )

                combined_docs = retrieved_docs + web_results + wos_results
                combined_docs.sort(key=lambda x: x.get("score", 0), reverse=True)

                if rag_system.search_limit:
                    combined_docs = combined_docs[:rag_system.search_limit]

                if not combined_docs:
                    if date_range:
                        yield f"data: {json.dumps({'type': 'error', 'message': f'No relevant documents found for {start_year}-{end_year}'})}\n\n"
                    else:
                        yield f"data: {json.dumps({'type': 'error', 'message': 'No relevant documents found'})}\n\n"
                    yield f"data: {json.dumps({'type': 'done'})}\n\n"
                    return

                yield f"data: {json.dumps({'type': 'status', 'message': f'Generating answer from {len(combined_docs)} documents...'})}\n\n"

                # Send context information
                context_refs = {}
                for i, doc in enumerate(combined_docs):
                    context_key = f"context_{i + 1}"
                    source_type = "Web of Science" if doc["filename"].startswith("wos:") else \
                        "Web Search" if doc["filename"].startswith("web:") else "Knowledge Base"

                    context_refs[context_key] = {
                        "filename": doc["filename"],
                        "text": doc["text"][:500] + ("..." if len(doc["text"]) > 500 else ""),
                        "score": doc.get("score", 0.0),
                        "metadata": doc.get("metadata", {}),
                        "source_type": source_type
                    }

                yield f"data: {json.dumps({'type': 'context', 'contexts': context_refs, 'temporal_analysis': True, 'wos_search_used': use_wos_search and len(wos_results) > 0})}\n\n"

                yield f"data: {json.dumps({'type': 'status', 'message': 'Generating answer...'})}\n\n"
                logger.info(f"[MEMORY DEBUG] Calling generate_streaming_answer with {len(session_memory)} memory items")
                for chunk in rag_system.generate_streaming_answer(user_query, combined_docs, conversation_memory=session_memory):
                    if chunk['type'] == 'content':
                        yield f"data: {json.dumps(chunk)}\n\n"
                    elif chunk['type'] == 'done':

                        if 'full_answer' in chunk:
                            logger.info(
                                f"[MEMORY DEBUG] Saving to memory: Q='{user_query}' A='{chunk['full_answer'][:100]}...'")
                            session_memory.append((user_query, chunk['full_answer']))
                            if len(session_memory) > rag_system.memory_turns:
                                session_memory = session_memory[-rag_system.memory_turns:]
                            memory_manager.set_memory(session_id, session_memory)
                            logger.info(f"[MEMORY DEBUG] Memory now contains {len(session_memory)} items")

                        if 'context_data' not in chunk:
                            chunk['context_data'] = {}

                        chunk['context_data'].update({
                            "query": user_query,
                            "expanded_query": expanded_query if expanded_query != user_query else None,
                            "date_range": date_range,
                            "temporal_info": temporal_info,
                            "temporal_analysis": True,
                            "wos_search_used": use_wos_search and len(wos_results) > 0,
                            "wos_results_count": len(wos_results),
                            "wos_timespan": wos_timespan if use_wos_search else None,
                            "sources": [
                                {
                                    "filename": doc["filename"],
                                    "score": doc.get("score", 0),
                                    "metadata": doc.get("metadata", {}),
                                    "source_type": "wos" if doc["filename"].startswith("wos:") else
                                    "web" if doc["filename"].startswith("web:") else "knowledge_base"
                                }
                                for doc in combined_docs
                            ]
                        })

                        # report generation
                        if generate_report:
                            yield f"data: {json.dumps({'type': 'status', 'message': 'Generating report...'})}\n\n"
                            try:
                                report_path, conversation_id = rag_system.generate_report(
                                    query=user_query,
                                    expanded_query=expanded_query,
                                    answer=chunk.get('full_answer', ''),
                                    contexts=combined_docs,
                                    context_data=chunk.get('context_data', {})
                                )
                                chunk['context_data']['report_path'] = report_path
                                chunk['context_data']['conversation_id'] = conversation_id
                            except Exception as e:
                                logger.error(f"Error generating report: {e}")

                        yield f"data: {json.dumps(chunk)}\n\n"
                        break
            else:
                # Use streaming logic
                expanded_query = rag_system.expand_query(
                    user_query) if rag_system.query_expansion_enabled else user_query
                yield f"data: {json.dumps({'type': 'status', 'message': 'Thinking...'})}\n\n"

                retrieved_docs = rag_system.retrieve(query=expanded_query, limit=rag_system.search_limit)

                # Web searches
                web_results = []
                wos_results = []

                # Brave Search
                if rag_system.use_web_search and rag_system.brave_api_key:
                    yield f"data: {json.dumps({'type': 'status', 'message': 'Searching the web...'})}\n\n"
                    web_results = rag_system.web_search(user_query, limit=2)

                # Web of Science Search
                if use_wos_search and hasattr(rag_system, 'wos_api_key') and rag_system.wos_api_key:
                    yield f"data: {json.dumps({'type': 'status', 'message': 'Searching Web of Science...'})}\n\n"
                    wos_results = rag_system.web_of_science_search(
                        user_query,
                        limit=3,
                        timespan=wos_timespan
                    )

                combined_docs = retrieved_docs + web_results + wos_results
                combined_docs.sort(key=lambda x: x.get("score", 0), reverse=True)

                if rag_system.search_limit:
                    combined_docs = combined_docs[:rag_system.search_limit]

                if not combined_docs:
                    yield f"data: {json.dumps({'type': 'error', 'message': 'No relevant documents found'})}\n\n"
                    yield f"data: {json.dumps({'type': 'done'})}\n\n"
                    return

                yield f"data: {json.dumps({'type': 'status', 'message': 'Generating answer...'})}\n\n"

                # Send context information
                context_refs = {}
                for i, doc in enumerate(combined_docs):
                    context_key = f"context_{i + 1}"
                    source_type = "Web of Science" if doc["filename"].startswith("wos:") else \
                        "Web Search" if doc["filename"].startswith("web:") else "Knowledge Base"

                    context_refs[context_key] = {
                        "filename": doc["filename"],
                        "text": doc["text"][:500] + ("..." if len(doc["text"]) > 500 else ""),
                        "score": doc.get("score", 0.0),
                        "metadata": doc.get("metadata", {}),
                        "source_type": source_type
                    }

                yield f"data: {json.dumps({'type': 'context', 'contexts': context_refs, 'wos_search_used': use_wos_search and len(wos_results) > 0})}\n\n"

                # Use existing streaming generator
                for chunk in rag_system.generate_streaming_answer(user_query, combined_docs, conversation_memory=session_memory):
                    if chunk['type'] in ['content']:
                        yield f"data: {json.dumps(chunk)}\n\n"
                    elif chunk['type'] == 'done':
                        if 'full_answer' in chunk:
                            session_memory.append((user_query, chunk['full_answer']))
                            if len(session_memory) > rag_system.memory_turns:
                                session_memory = session_memory[-rag_system.memory_turns:]
                            memory_manager.set_memory(session_id, session_memory)

                        if 'context_data' not in chunk:
                            chunk['context_data'] = {}

                        chunk['context_data'].update({
                            'wos_search_used': use_wos_search and len(wos_results) > 0,
                            'wos_results_count': len(wos_results),
                            'wos_timespan': wos_timespan if use_wos_search else None,
                            'sources': [
                                {
                                    "filename": doc["filename"],
                                    "score": doc.get("score", 0),
                                    "metadata": doc.get("metadata", {}),
                                    "source_type": "wos" if doc["filename"].startswith("wos:") else
                                    "web" if doc["filename"].startswith("web:") else "knowledge_base"
                                }
                                for doc in combined_docs
                            ]
                        })

                        if generate_report:
                            yield f"data: {json.dumps({'type': 'status', 'message': 'Generating report...'})}\n\n"
                            try:
                                report_path, conversation_id = rag_system.generate_report(
                                    query=user_query,
                                    expanded_query=expanded_query,
                                    answer=chunk.get('full_answer', ''),
                                    contexts=combined_docs,
                                    context_data=chunk.get('context_data', {})
                                )
                                chunk['context_data']['report_path'] = report_path
                                chunk['context_data']['conversation_id'] = conversation_id
                            except Exception as e:
                                logger.error(f"Error generating report: {e}")

                        yield f"data: {json.dumps(chunk)}\n\n"
                        break

        except Exception as e:
            logger.error(f"Error in streaming query: {e}")
            yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
            yield f"data: {json.dumps({'type': 'done'})}\n\n"

    response = Response(
        generate_stream(),
        mimetype='text/plain',
        headers={
            'Cache-Control': 'no-cache',
            'Connection': 'keep-alive',
            'X-Accel-Buffering': 'no',
        }
    )

    session.modified = True
    return response


@app.route('/retrieve', methods=['POST'])
@validate_rag_system
@validate_json_request(['query'])
def handle_retrieve(data):
    user_query = data['query']
    limit = data.get('limit', rag_system.search_limit)
    text_match_filter = data.get('text_match_filter', None)

    logger.info(f"Retrieve request: '{user_query}' (Limit: {limit})")

    try:
        expanded_query = rag_system.expand_query(user_query) if rag_system.query_expansion_enabled else user_query
        retrieved_docs = rag_system.retrieve(
            query=expanded_query,
            text_match_filter=text_match_filter,
            limit=limit
        )

        response_data = build_retrieve_response(user_query, expanded_query, retrieved_docs)
        logger.info(f"Retrieved {len(retrieved_docs)} documents")
        return jsonify(response_data)

    except Exception as e:
        return handle_api_error(e, "document retrieval")


@app.route('/text_search', methods=['POST'])
@validate_rag_system
@validate_json_request(['term'])
def handle_text_search(data):
    search_term = data['term']
    limit = data.get('limit', rag_system.search_limit)

    logger.info(f"Text search: '{search_term}' (Limit: {limit})")

    try:
        results = rag_system.retrieve(query="", text_match_filter=search_term, limit=limit)

        response_data = {
            "search_term": search_term,
            "total_results": len(results),
            "documents": []
        }

        for i, doc in enumerate(results):
            text = doc["text"]
            search_term_lower = search_term.lower()
            text_lower = text.lower()

            first_pos = text_lower.find(search_term_lower)
            if first_pos != -1:
                start_preview = max(0, first_pos - 150)
                end_preview = min(len(text), first_pos + len(search_term) + 150)
                preview = text[start_preview:end_preview]

                if start_preview > 0:
                    preview = "..." + preview
                if end_preview < len(text):
                    preview = preview + "..."

                occurrences = text_lower.count(search_term_lower)
            else:
                preview = text[:300] + ("..." if len(text) > 300 else "")
                occurrences = 0

            doc_data = format_document(doc, i)
            doc_data.update({
                "text_preview": preview,
                "occurrences": occurrences
            })
            response_data["documents"].append(doc_data)

        logger.info(f"Found {len(results)} documents with text matches")
        return jsonify(response_data)

    except Exception as e:
        return handle_api_error(e, "text search")


@app.route('/settings', methods=['GET'])
@validate_rag_system
def get_settings():
    settings = {
        "search_limit": rag_system.search_limit,
        "dense_weight": rag_system.dense_weight,
        "sparse_weight": rag_system.sparse_weight,
        "temperature": rag_system.temperature,
        "max_tokens": rag_system.max_tokens,
        "use_web_search": rag_system.use_web_search,
        "use_wos_search": getattr(rag_system, 'use_wos_search', True),
        "wos_api_available": bool(getattr(rag_system, 'wos_api_key', False)),
        "memory_enabled": rag_system.memory_enabled,
        "memory_turns": rag_system.memory_turns,
        "query_expansion_enabled": rag_system.query_expansion_enabled,
        "query_expansion_mode": rag_system.query_expansion_mode,
        "current_collection": rag_system.collection_name
    }
    return jsonify(settings)


@app.route('/settings', methods=['POST'])
@validate_rag_system
@validate_json_request()
def update_settings(data):
    try:
        if 'search_limit' in data:
            rag_system.search_limit = int(data['search_limit'])

        if 'dense_weight' in data and 'sparse_weight' in data:
            dense = float(data['dense_weight'])
            sparse = float(data['sparse_weight'])
            total = dense + sparse
            if total > 0:
                rag_system.dense_weight = dense / total
                rag_system.sparse_weight = sparse / total
            else:
                rag_system.dense_weight = 0.5
                rag_system.sparse_weight = 0.5
        elif 'dense_weight' in data:
            dense = float(data['dense_weight'])
            rag_system.dense_weight = dense
            rag_system.sparse_weight = 1.0 - dense
        elif 'sparse_weight' in data:
            sparse = float(data['sparse_weight'])
            rag_system.sparse_weight = sparse
            rag_system.dense_weight = 1.0 - sparse

        if 'temperature' in data:
            rag_system.temperature = float(data['temperature'])
        if 'max_tokens' in data:
            rag_system.max_tokens = int(data['max_tokens'])
        if 'use_web_search' in data:
            rag_system.use_web_search = bool(data['use_web_search'])
        if 'memory_enabled' in data:
            rag_system.memory_enabled = bool(data['memory_enabled'])
        if 'memory_turns' in data:
            rag_system.memory_turns = int(data['memory_turns'])
        if 'query_expansion_enabled' in data:
            rag_system.query_expansion_enabled = bool(data['query_expansion_enabled'])
        if 'query_expansion_mode' in data:
            rag_system.query_expansion_mode = str(data['query_expansion_mode'])

        logger.info("Settings updated successfully")
        return jsonify({"status": "success", "message": "Settings updated successfully"})

    except ValueError as e:
        return jsonify({"error": f"Invalid value: {e}"}), 400
    except Exception as e:
        return handle_api_error(e, "settings update")


@app.route('/clear_memory', methods=['POST'])
@validate_rag_system
def clear_memory():
    try:
        if 'session_id' in session:
            session_id = session['session_id']
            memory_manager.clear_memory(session_id)
            logger.info(f"Memory cleared for session: {session_id}")
            session.modified = True
            return jsonify({"status": "success", "message": "Memory cleared"})
        else:
            logger.info("No active session to clear")
            return jsonify({"status": "success", "message": "No active session to clear"})
    except Exception as e:
        return handle_api_error(e, "memory clearing")


@app.route('/collections', methods=['GET'])
@validate_rag_system
def get_collections():
    try:
        collections = rag_system.milvus_client.list_collections()
        return jsonify({"collections": collections})
    except Exception as e:
        return handle_api_error(e, "collection retrieval")


@app.route('/switch_collection', methods=['POST'])
@validate_rag_system
@validate_json_request(['new_collection_name'])
def switch_collection(data):
    new_collection_name = data['new_collection_name']
    try:
        success = rag_system.switch_collection(new_collection_name, auto_load_metadata=Config.AUTO_LOAD_METADATA)
        if success:
            logger.info(f"Switched to collection: {new_collection_name}")
            return jsonify({"status": "success", "message": f"Switched to collection '{new_collection_name}'"})
        else:
            return jsonify({"status": "failed", "message": f"Collection '{new_collection_name}' not found"}), 404
    except Exception as e:
        return handle_api_error(e, "collection switching")

# Clustering code
@app.route('/find_similar', methods=['POST'])
@validate_rag_system
@validate_json_request(['filename'])
def find_similar_documents(data):
    filename = data['filename']
    limit = data.get('limit', 5)

    try:
        results = rag_system.milvus_client.query(
            collection_name=rag_system.collection_name,
            filter=f"filename == '{filename.replace(chr(39), chr(39) + chr(39))}'",
            output_fields=["dense"],
            limit=1
        )

        if not results:
            return jsonify({"error": "Document not found"}), 404

        source_embedding = results[0]["dense"]

        similar_results = rag_system.milvus_client.search(
            collection_name=rag_system.collection_name,
            data=[source_embedding],
            anns_field="dense",
            search_params={"metric_type": "COSINE", "params": {"ef": 50}},
            limit=limit + 1,
            output_fields=["filename", "text"]
        )

        similar_docs = []
        seen_files = set()

        for hit in similar_results[0]:
            hit_filename = hit["entity"]["filename"]
            base_filename = hit_filename.split('_chunk_')[0] if '_chunk_' in hit_filename else hit_filename
            source_base = filename.split('_chunk_')[0] if '_chunk_' in filename else filename

            if base_filename != source_base and base_filename not in seen_files and len(similar_docs) < limit:
                seen_files.add(base_filename)
                document_metadata = rag_system.metadata.get(base_filename, {})

                similar_docs.append({
                    "filename": hit_filename,
                    "similarity_score": hit["distance"],
                    "text_preview": hit["entity"]["text"][:300] + ("..." if len(hit["entity"]["text"]) > 300 else ""),
                    "metadata": document_metadata
                })

        return jsonify({
            "source_document": filename,
            "similar_documents": similar_docs,
            "total_results": len(similar_docs)
        })

    except Exception as e:
        return handle_api_error(e, "similarity search")


@app.route('/cluster_documents', methods=['POST'])
@validate_rag_system
def cluster_documents():
    data = request.get_json() or {}
    n_clusters = data.get('n_clusters', 5)
    sample_size = data.get('sample_size', 100)

    try:
        results = rag_system.milvus_client.query(
            collection_name=rag_system.collection_name,
            filter="",
            output_fields=["filename", "dense", "text"],
            limit=sample_size
        )

        if len(results) < n_clusters:
            return jsonify({"error": "Not enough documents for clustering"}), 400

        embeddings = [doc["dense"] for doc in results]
        filenames = [doc["filename"] for doc in results]

        from sklearn.cluster import KMeans
        import numpy as np

        kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
        cluster_labels = kmeans.fit_predict(np.array(embeddings))

        clusters = {}
        for i, (filename, label) in enumerate(zip(filenames, cluster_labels)):
            if label not in clusters:
                clusters[label] = []

            base_filename = filename.split('_chunk_')[0] if '_chunk_' in filename else filename
            document_metadata = rag_system.metadata.get(base_filename, {})

            clusters[label].append({
                "filename": filename,
                "text_preview": results[i]["text"][:200] + "...",
                "metadata": document_metadata
            })

        cluster_info = []
        for cluster_id, docs in clusters.items():
            topics = []
            dates = []
            authors = set()

            for doc in docs:
                meta = doc.get("metadata", {})
                if meta.get("topic"):
                    topics.append(meta["topic"])
                if meta.get("date"):
                    dates.append(meta["date"])
                if meta.get("author"):
                    authors.add(meta["author"])

            cluster_info.append({
                "cluster_id": int(cluster_id),
                "document_count": len(docs),
                "documents": docs[:10],
                "common_topics": list(set(topics))[:3],
                "date_range": f"{min(dates)} - {max(dates)}" if dates else "Unknown",
                "authors": list(authors)[:5]
            })

        return jsonify({
            "clusters": cluster_info,
            "total_documents": len(results),
            "n_clusters": n_clusters
        })

    except ImportError:
        return jsonify({"error": "scikit-learn not installed. Run: pip install scikit-learn"}), 500
    except Exception as e:
        return handle_api_error(e, "document clustering")


@app.route('/get_document', methods=['POST'])
@validate_rag_system
@validate_json_request(['filename'])
def get_full_document(data):
    filename = data['filename']

    try:
        results = rag_system.milvus_client.query(
            collection_name=rag_system.collection_name,
            filter=f"filename == '{filename.replace(chr(39), chr(39) + chr(39))}'",
            output_fields=["filename", "text"],
            limit=1
        )

        if not results:
            return jsonify({"error": "Document not found"}), 404

        document = results[0]
        base_filename = filename.split('_chunk_')[0] if '_chunk_' in filename else filename
        document_metadata = rag_system.metadata.get(base_filename, {})

        if '_chunk_' in filename:
            base_name = filename.split('_chunk_')[0]
            chunk_results = rag_system.milvus_client.query(
                collection_name=rag_system.collection_name,
                filter=f"filename like '{base_name}_chunk_%'",
                output_fields=["filename", "text"],
                limit=50
            )

            sorted_chunks = sorted(chunk_results, key=lambda x: int(x['filename'].split('_chunk_')[1].split('of')[0]))
            combined_text = '\n\n'.join([chunk['text'] for chunk in sorted_chunks])

            return jsonify({
                "filename": filename,
                "text": combined_text,
                "metadata": document_metadata,
                "chunks": len(sorted_chunks),
                "is_chunked": True
            })
        else:
            return jsonify({
                "filename": filename,
                "text": document["text"],
                "metadata": document_metadata,
                "chunks": 1,
                "is_chunked": False
            })

    except Exception as e:
        return handle_api_error(e, "document retrieval")


# Temporal Analysis
@app.route('/temporal_trend_analysis', methods=['POST'])
@validate_rag_system
@validate_json_request(['query', 'start_year', 'end_year'])
def handle_temporal_trend_analysis(data):
    query = data['query']
    start_year = int(data['start_year'])
    end_year = int(data['end_year'])
    interval_years = data.get('interval_years', 10)

    logger.info(f"Temporal trend analysis: '{query}' from {start_year}-{end_year}")

    try:
        if not hasattr(rag_system, 'temporal_analyzer'):
            return jsonify({"error": "Temporal analysis not available"}), 500

        results = rag_system.temporal_analyzer.temporal_trend_analysis(
            query=query,
            date_range=(start_year, end_year),
            interval_years=interval_years
        )

        logger.info(f"Generated trend analysis with {len(results['periods'])} periods")
        return jsonify(results)

    except Exception as e:
        return handle_api_error(e, "temporal trend analysis")

@app.route('/temporal_search', methods=['POST'])
@validate_rag_system
@validate_json_request(['query'])
def handle_temporal_search(data):
    query = data['query']
    start_year = data.get('start_year')
    end_year = data.get('end_year')
    limit = data.get('limit', rag_system.search_limit)

    logger.info(f"Temporal search: '{query}' {start_year}-{end_year if end_year else 'present'}")

    try:
        if not hasattr(rag_system, 'temporal_analyzer'):
            return jsonify({"error": "Temporal analysis not available"}), 500

        # Use metadata based temporal filtering
        if start_year or end_year:
            results = rag_system.temporal_analyzer.retrieve_with_temporal_filter(
                query=query,
                start_year=start_year,
                end_year=end_year,
                limit=limit
            )
        else:
            results = rag_system.retrieve(query=query, limit=limit)

        response_data = build_retrieve_response(query, query, results)
        response_data['temporal_filter_applied'] = bool(start_year or end_year)
        response_data['date_range'] = [start_year, end_year] if start_year or end_year else None

        logger.info(f"Temporal search found {len(results)} documents")
        return jsonify(response_data)

    except Exception as e:
        return handle_api_error(e, "temporal search")


@app.route('/temporal_timeline', methods=['POST'])
@validate_rag_system
@validate_json_request(['query'])
def handle_temporal_timeline(data):
    query = data['query']
    start_year = data.get('start_year', 1400)
    end_year = data.get('end_year', 1700)

    logger.info(f"Generating timeline for: '{query}' from {start_year}-{end_year}")

    try:
        if not hasattr(rag_system, 'temporal_analyzer'):
            return jsonify({"error": "Temporal analysis not available"}), 500

        # Get all documents matching query, then filter by date
        documents = rag_system.temporal_analyzer.retrieve_with_temporal_filter(
            query=query,
            start_year=start_year,
            end_year=end_year,
            limit=100
        )

        # Group documents by year
        timeline_data = defaultdict(list)
        for doc in documents:
            metadata = doc.get('metadata', {})
            if metadata.get('date'):
                try:
                    year = int(metadata['date'])
                    timeline_data[year].append({
                        'filename': doc['filename'],
                        'title': metadata.get('title', 'Untitled'),
                        'author': metadata.get('author', 'Unknown'),
                        'score': doc.get('score', 0),
                        'text_preview': doc['text'][:200] + "..."
                    })
                except (ValueError, TypeError):
                    continue

        # Sort timeline by year
        sorted_timeline = sorted(timeline_data.items())

        response_data = {
            'query': query,
            'date_range': [start_year, end_year],
            'total_documents': len(documents),
            'timeline': [
                {
                    'year': year,
                    'document_count': len(docs),
                    'documents': docs
                }
                for year, docs in sorted_timeline
            ]
        }

        logger.info(f"Generated timeline with {len(sorted_timeline)} years")
        return jsonify(response_data)

    except Exception as e:
        return handle_api_error(e, "temporal timeline generation")


@app.route('/wos_search', methods=['POST'])
@validate_rag_system
@validate_json_request(['query'])
def handle_wos_search(data):
    """Handle Web of Science search requests"""
    query = data['query']
    limit = data.get('limit', 5)
    timespan = data.get('timespan', '1400-01-01+1700-12-31')

    logger.info(f"Web of Science search: '{query}' (limit: {limit}, timespan: {timespan})")

    try:
        results = rag_system.web_of_science_search(
            query=query,
            limit=limit,
            timespan=timespan
        )

        response_data = {
            "query": query,
            "timespan": timespan,
            "total_results": len(results),
            "documents": []
        }

        for i, doc in enumerate(results):
            doc_data = {
                "rank": i + 1,
                "filename": doc["filename"],
                "text_preview": doc["text"][:300] + ("..." if len(doc["text"]) > 300 else ""),
                "full_text": doc["text"],
                "score": doc.get("score", 0.0),
                "metadata": doc.get("metadata", {}),
                "source": "Web of Science"
            }
            response_data["documents"].append(doc_data)

        logger.info(f"Found {len(results)} Web of Science results")
        return jsonify(response_data)

    except Exception as e:
        return handle_api_error(e, "Web of Science search")

# Agentic Pattern Recognition Implementation ---------------------------------------------------------
@app.route('/analyze_concept_correlation', methods=['POST'])
@validate_rag_system
@validate_json_request(['concept_a', 'concept_b', 'start_year', 'end_year'])
def handle_concept_correlation(data):
    concept_a = data['concept_a']
    concept_b = data['concept_b']
    start_year = int(data['start_year'])
    end_year = int(data['end_year'])
    sample_size = data.get('sample_size', 500)
    confidence_level = data.get('confidence_level', 0.95)
    effect_size_threshold = data.get('effect_size_threshold', 0.2)

    logger.info(f"Enhanced concept correlation: '{concept_a}' vs '{concept_b}' ({start_year}-{end_year})")

    try:
        # Check if enhanced method exists, fallback to original
        if hasattr(rag_system.pattern_discovery, 'analyze_concept_correlation_enhanced'):
            results = rag_system.pattern_discovery.analyze_concept_correlation_enhanced(
                concept_a=concept_a,
                concept_b=concept_b,
                date_range=(start_year, end_year),
                sample_size=sample_size,
                confidence_level=confidence_level,
                effect_size_threshold=effect_size_threshold
            )
        else:
            # Fallback to original method
            results = rag_system.pattern_discovery.analyze_concept_correlation(
                concept_a=concept_a,
                concept_b=concept_b,
                date_range=(start_year, end_year),
                sample_size=sample_size
            )

        logger.info(f"Correlation analysis completed")
        return jsonify(results)

    except Exception as e:
        return handle_api_error(e, "concept correlation analysis")


@app.route('/detect_temporal_anomalies', methods=['POST'])
@validate_rag_system
@validate_json_request(['concept', 'start_year', 'end_year'])
def handle_temporal_anomalies(data):
    concept = data['concept']
    start_year = int(data['start_year'])
    end_year = int(data['end_year'])
    window_size = data.get('window_size', 10)
    sample_size = data.get('sample_size', 100)
    confidence_level = data.get('confidence_level', 0.95)
    min_documents_per_window = data.get('min_documents_per_window', 5)

    logger.info(f"Enhanced temporal anomaly detection: '{concept}' ({start_year}-{end_year})")

    try:
        # Check if enhanced method exists
        if hasattr(rag_system.pattern_discovery, 'detect_temporal_anomalies_enhanced'):
            results = rag_system.pattern_discovery.detect_temporal_anomalies_enhanced(
                concept=concept,
                date_range=(start_year, end_year),
                window_size=window_size,
                sample_size=sample_size,
                confidence_level=confidence_level,
                min_documents_per_window=min_documents_per_window
            )
        else:
            # Fallback to original method
            results = rag_system.pattern_discovery.detect_temporal_anomalies(
                concept=concept,
                date_range=(start_year, end_year),
                window_size=window_size,
                sample_size=sample_size
            )

        logger.info(f"Anomaly detection completed")
        return jsonify(results)

    except Exception as e:
        return handle_api_error(e, "temporal anomaly detection")


@app.route('/statistical_options', methods=['GET'])
@validate_rag_system
def get_statistical_options():
    """Get available statistical options and their defaults"""
    options = {
        "correlation_analysis": {
            "confidence_levels": [0.90, 0.95, 0.99],
            "effect_size_thresholds": [0.1, 0.2, 0.3, 0.5],
            "sample_sizes": [200, 500, 1000, 2000],
            "default_confidence": 0.95,
            "default_effect_threshold": 0.2,
            "default_sample_size": 500
        },
        "anomaly_detection": {
            "confidence_levels": [0.90, 0.95, 0.99],
            "window_sizes": [5, 10, 15, 20, 25],
            "min_documents_per_window": [3, 5, 10, 15],
            "detection_methods": ["z_score", "mad_robust", "iqr", "modified_z"],
            "default_confidence": 0.95,
            "default_window_size": 10,
            "default_min_docs": 5
        },
        "available_enhancements": {
            "correlation": hasattr(rag_system.pattern_discovery, 'analyze_concept_correlation_enhanced'),
            "anomaly": hasattr(rag_system.pattern_discovery, 'detect_temporal_anomalies_enhanced')
        }
    }

    return jsonify(options)


def cleanup_sessions_periodically():
    """Clean up old sessions"""
    import glob
    while True:
        time.sleep(memory_manager.cleanup_interval)

        removed = memory_manager.cleanup_old_sessions()
        if removed > 0:
            logger.info(f"Cleaned up {removed} inactive sessions")

        session_dir = os.path.join(os.getcwd(), 'flask_session')
        if os.path.exists(session_dir):
            current_time = time.time()
            cleaned_files = 0
            for f in glob.glob(f'{session_dir}/*'):
                try:
                    if current_time - os.path.getmtime(f) > 7200:  # 2 hours
                        os.remove(f)
                        cleaned_files += 1
                except Exception as e:
                    pass
            if cleaned_files > 0:
                logger.info(f"Cleaned up {cleaned_files} old session files")

import threading
cleanup_thread = threading.Thread(target=cleanup_sessions_periodically, daemon=True)
cleanup_thread.start()

if __name__ == '__main__':
    app.run(host=Config.HOST, port=Config.PORT, debug=Config.DEBUG)