from flask import Flask, render_template, request, jsonify
from functools import lru_cache
import math
import os
from dotenv import load_dotenv
os.environ['CUDA_VISIBLE_DEVICES'] = "1"
from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Searcher
import torch
import base64
import io

app = Flask(__name__)

dataset_name = "webq"

searcher = Searcher(index=f"/data1/liuyaoyang/Papers/icml2025/multi_rag/RAG/Search-in-the-Chain/ColBERT/experiments/{dataset_name}/indexes/{dataset_name}.nbits=2")


def api_search_query(query, mask, k=100):
    pids, ranks, scores = searcher.search_with_mask(query, mask, k=100)
    # pids, ranks, scores = searcher.search(query, k=100)
    pids, ranks, scores = pids[:100], ranks[:100], scores[:100]
    passages = [searcher.collection[pid] for pid in pids]
    probs = [math.exp(score) for score in scores]
    probs = [prob / sum(probs) for prob in probs]
    topk = []
    for pid, rank, score, prob in zip(pids, ranks, scores, probs):
        text = searcher.collection[pid]            
        d = {'text': text, 'pid': pid, 'rank': rank, 'score': score, 'prob': prob}
        topk.append(d)
    topk = list(sorted(topk, key=lambda p: (-1 * p['score'], p['pid'])))
    
    return {
        "topk": topk,
    }

@app.route('/api/search', methods=['POST'])
def handle_search():
    # 获取JSON数据
    data = request.get_json()
    
    # 校验必要参数
    if not data or 'query' not in data or 'mask' not in data:
        return jsonify({"error": "Missing 'query' or 'mask' parameter"}), 400
    
    try:
        # 解析参数
        query = data['query']
        mask_b64 = data['mask']
        k = data.get('k', 100)  # 默认值100
        
        # Base64解码并加载张量
        mask_bytes = base64.b64decode(mask_b64)
        buffer = io.BytesIO(mask_bytes)
        mask_tensor = torch.load(buffer)

        # print(mask_tensor)
        
        # 执行处理函数
        result = api_search_query(query, mask_tensor, k)
        return jsonify(result)
    
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == '__main__':
    app.run("0.0.0.0", 8110)