routes/query.py

from flask import Blueprint, request, jsonify, session
from datetime import datetime
from bson import ObjectId
import logging
import json

from utils.extensions import ext
from models.memory import SimpleConversationMemory
from utils.helpers import format_chat_history, preprocess_related_questions
from utils.query import can_make_query, update_query_count
from services.retriever import retrieve
from services.sambanova import sambanova_service
from services.prompts import create_main_json_prompt, create_judgment_json_prompt, create_related_questions_json_prompt

query_bp = Blueprint('query', __name__)
memory = SimpleConversationMemory(max_messages=10)

@query_bp.route("/query", methods=["POST"])
def query():
    if 'user_id' not in session:
        return jsonify({
            'error': 'Vui lòng đăng nhập để sử dụng tính năng này',
            'error_code': 'UNAUTHENTICATED'
        }), 401
    
    user_id = session['user_id']
    
    # Check query permission
    can_query, error_message, query_count, query_limit = can_make_query(user_id)
    if not can_query:
        return jsonify({
            'error': error_message,
            'error_code': 'QUERY_LIMIT_EXCEEDED',
            'query_count': query_count,
            'query_limit': query_limit,
            'upgrade_url': '/'
        }), 403
    elif error_message:
        logging.info(f"User {user_id} received warning: {error_message}")

    # Parse JSON input
    data = request.get_json(silent=True) or {}
    question = data.get('question', '').strip()
    if not question:
        return jsonify({
            'error': 'Câu hỏi không hợp lệ',
            'error_code': 'INVALID_QUESTION'
        }), 400

    try:
        # Update query count for limited accounts
        user = ext.db.users.find_one({'_id': ObjectId(user_id)})
        if not user:
            return jsonify({
                'error': 'Người dùng không tồn tại',
                'error_code': 'USER_NOT_FOUND'
            }), 404

        # Handle conversation
        conversation_id = data.get('conversation_id')
        if conversation_id:
            # Validate conversation_id
            if not ObjectId.is_valid(conversation_id):
                return jsonify({
                    'error': 'ID hội thoại không hợp lệ',
                    'error_code': 'INVALID_CONVERSATION_ID'
                }), 400
            
            conversation = ext.db.conversations.find_one({
                '_id': ObjectId(conversation_id),
                'user_id': user_id
            })
            if not conversation:
                return jsonify({
                    'error': 'Hội thoại không tồn tại hoặc không thuộc về người dùng',
                    'error_code': 'CONVERSATION_NOT_FOUND'
                }), 404
        else:
            # Create new conversation
            conversation = {
                'user_id': user_id,
                'title': question[:50] + ('...' if len(question) > 50 else ''),
                'timestamp': datetime.utcnow(),
                'created_at': datetime.utcnow(),
                'updated_at': datetime.utcnow()
            }
            result = ext.db.conversations.insert_one(conversation)
            conversation_id = str(result.inserted_id)

        # Save user message
        user_message = {
            'conversation_id': conversation_id,
            'type': 'user',
            'content': question,
            'timestamp': datetime.utcnow(),
            'created_at': datetime.utcnow()
        }
        ext.db.messages.insert_one(user_message)

        # Retrieve relevant legal documents
        banan_results = retrieve(question, collection_name="law_documents")

        # Format chat history
        chat_history_str = format_chat_history(memory)

        # Use JSON prompt for main response
        main_prompt = create_main_json_prompt(question, banan_results, chat_history_str)

        # Call SambaNova for main response
        answer = sambanova_service.generate_response(main_prompt, "Qwen3-235B")
        if not answer or answer == "SambaNova client not initialized.":
            answer = "Xin lỗi, hiện tại hệ thống đang bận. Vui lòng thử lại sau."

        # Generate related questions
        related_questions_prompt = create_related_questions_json_prompt(question)
        related_questions_raw = sambanova_service.generate_response(related_questions_prompt, "Qwen3-235B")
        related_questions = preprocess_related_questions(related_questions_raw)

        assistant_message = {
            'conversation_id': conversation_id,
            'type': 'assistant',
            'content': answer,
            'timestamp': datetime.utcnow(),
            'created_at': datetime.utcnow(),
            'sources': banan_results,
            'related_questions': related_questions
        }
        ext.db.messages.insert_one(assistant_message)

        # Update conversation timestamp and title if it's a new conversation
        ext.db.conversations.update_one(
            {'_id': ObjectId(conversation_id)},
            {'$set': {
                'updated_at': datetime.utcnow(),
                'title': question[:50] + ('...' if len(question) > 50 else '')
            }}
        )

        memory.save_context({'question': question}, {'answer': answer})
        
        # Update query count after successful query
        if user.get('account_type') == 'limited':
            update_query_count(user_id)
            # Get updated user data
            user = ext.db.users.find_one({'_id': ObjectId(user_id)})
            query_count = user.get('query_count', 0)
        
        return jsonify({
            'final_response': answer,
            'top_banan_documents': banan_results,
            'chat_history': chat_history_str,
            'sources': banan_results,
            'related_questions': related_questions,
            'conversation_id': conversation_id,
            'query_count': query_count,
            'query_limit': query_limit,
            'success': True
        }), 200

    except Exception as e:
        logging.error(f"Error in query endpoint: {e}")
        return jsonify({
            'error': 'Lỗi hệ thống',
            'error_code': 'SERVER_ERROR',
            'success': False
        }), 500

@query_bp.route("/draft_judgment", methods=["POST"])
def draft_judgment():
    if 'user_id' not in session:
        return jsonify({'error': 'Vui lòng đăng nhập'}), 401
    
    data = request.get_json(silent=True) or {}
    case_details = data.get('case_details', '').strip()
    
    if not case_details:
        return jsonify({'error': 'Chi tiết vụ án không hợp lệ'}), 400
    
    if len(case_details) < 50:
        return jsonify({'error': 'Vui lòng cung cấp chi tiết vụ án đầy đủ (ít nhất 50 ký tự)'}), 400
    
    # Check query permission
    user_id = session['user_id']
    can_query, error_message, query_count, query_limit = can_make_query(user_id)
    if not can_query:
        return jsonify({
            'error': error_message,
            'error_code': 'QUERY_LIMIT_EXCEEDED',
            'query_count': query_count,
            'query_limit': query_limit
        }), 403
    
    try:
        # Retrieve relevant legal documents
        banan_results = retrieve(case_details, collection_name="law_documents")
        
        # Format results for display
        top_banan_docs = []
        for result in banan_results[:3]:  # Limit to top 3
            top_banan_docs.append({
                'source': result['file'],
                'folder': result['folder'],
                'text_path': result['text_path'],
                'percentage_similarity': result['percentage_similarity']
            })
        
        chat_history_str = format_chat_history(memory)
        
        # Use JSON prompt for judgment drafting
        judgment_prompt = create_judgment_json_prompt(case_details, banan_results)
        judgment = sambanova_service.generate_response(judgment_prompt, "Qwen3-235B")
        
        if not judgment or judgment == "SambaNova client not initialized.":
            judgment = "Xin lỗi, hiện tại không thể tạo bản án. Vui lòng thử lại sau."
        
        # Create conversation for judgment drafting
        conversation = {
            'user_id': user_id,
            'title': 'Soạn thảo bản án - ' + case_details[:30] + ('...' if len(case_details) > 30 else ''),
            'timestamp': datetime.utcnow(),
            'created_at': datetime.utcnow(),
            'type': 'judgment_draft'
        }
        result = ext.db.conversations.insert_one(conversation)
        conversation_id = str(result.inserted_id)
        
        # Save user message
        user_message = {
            'conversation_id': conversation_id,
            'type': 'user',
            'content': f"Soạn thảo bản án với chi tiết: {case_details}",
            'timestamp': datetime.utcnow(),
            'created_at': datetime.utcnow()
        }
        ext.db.messages.insert_one(user_message)
        
        # Save assistant message
        assistant_message = {
            'conversation_id': conversation_id,
            'type': 'assistant',
            'content': judgment,
            'timestamp': datetime.utcnow(),
            'created_at': datetime.utcnow(),
            'sources': banan_results
        }
        ext.db.messages.insert_one(assistant_message)
        
        memory.save_context({'case_details': case_details}, {'judgment': judgment})
        
        # Update query count for limited accounts
        if session.get('account_type') == 'limited':
            update_query_count(user_id)
        
        return jsonify({
            'judgment': judgment,
            'top_banan_documents': top_banan_docs,
            'chat_history': chat_history_str,
            'conversation_id': conversation_id,
            'success': True
        }), 200
    except Exception as e:
        logging.error(f"Error in draft_judgment: {e}")
        return jsonify({'error': 'Lỗi hệ thống', 'success': False}), 500

@query_bp.route("/search_documents", methods=["POST"])
def search_documents():
    """API để tìm kiếm tài liệu pháp lý"""
    if 'user_id' not in session:
        return jsonify({'error': 'Vui lòng đăng nhập'}), 401
    
    data = request.get_json(silent=True) or {}
    query_text = data.get('query', '').strip()
    
    if not query_text:
        return jsonify({'error': 'Câu truy vấn không được để trống'}), 400
    
    try:
        # Retrieve documents
        results = retrieve(query_text, collection_name="law_documents", k=5)
        
        # Format results
        formatted_results = []
        for result in results:
            formatted_results.append({
                'file': result['file'],
                'folder': result['folder'],
                'text_path': result['text_path'],
                'content_preview': result['content'][:200] + '...' if len(result['content']) > 200 else result['content'],
                'percentage_similarity': round(result['percentage_similarity'], 2),
                'relevance_score': result['cosine_score']
            })
        
        return jsonify({
            'query': query_text,
            'results': formatted_results,
            'total_results': len(formatted_results),
            'success': True
        }), 200
    except Exception as e:
        logging.error(f"Error searching documents: {e}")
        return jsonify({'error': 'Lỗi hệ thống', 'success': False}), 500

@query_bp.route("/clear_memory", methods=["POST"])
def clear_memory():
    """API để xóa bộ nhớ hội thoại"""
    if 'user_id' not in session:
        return jsonify({'error': 'Vui lòng đăng nhập'}), 401
    
    try:
        memory.clear()
        return jsonify({'message': 'Đã xóa bộ nhớ hội thoại', 'success': True}), 200
    except Exception as e:
        logging.error(f"Error clearing memory: {e}")
        return jsonify({'error': 'Lỗi hệ thống', 'success': False}), 500

@query_bp.route("/get_memory", methods=["GET"])
def get_memory():
    """API để lấy nội dung bộ nhớ hội thoại"""
    if 'user_id' not in session:
        return jsonify({'error': 'Vui lòng đăng nhập'}), 401
    
    try:
        messages = memory.get_messages()
        return jsonify({
            'messages': messages,
            'total_messages': len(messages),
            'success': True
        }), 200
    except Exception as e:
        logging.error(f"Error getting memory: {e}")
        return jsonify({'error': 'Lỗi hệ thống', 'success': False}), 500