from flask import Flask, render_template, request, jsonify
from functools import lru_cache
import math
import os
from dotenv import load_dotenv

from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Searcher

load_dotenv()

INDEX_NAME = os.getenv("INDEX_NAME")
INDEX_ROOT = os.getenv("INDEX_ROOT")
app = Flask(__name__)

# #searcher = Searcher(index=f"{INDEX_ROOT}/{INDEX_NAME}")
# searcher = Searcher(index=f"/home/icml01/multi_rag/RAG/Search-in-the-Chain/ColBERT/experiments/strategyqa/indexes/strategyqa.nbits=2")
import sys
sys.path.append(r"/home/icml01/multi_rag/RAG/Decompose_retrieval")
from multihop_ir import *
from transformers import AutoTokenizer
from openai import OpenAI

instruction = 'You are a query decomposition assistant. Please decompose one query Q into semantically coherent sub-queries.'
model_id = "/home/icml01/Models/Llama-3.1-8B-Instruct"
pipe = pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device="cuda:1",
)


def call_llama3_single_prompt(
    input_str, model="Llama-3.1-8B-Instruct", max_decode_steps=200, temperature=0.8
):

    if isinstance(input_str, str):
      messages = [
          {"role": "user", "content": input_str},
      ]
    else:
      messages = input_str
    
    if temperature > 0:
      outputs = pipe(
          messages,
          max_new_tokens=max_decode_steps,
          temperature=temperature,
          pad_token_id=pipe.tokenizer.eos_token_id,
      )
    else:
      outputs = pipe(
          messages,
          max_new_tokens=max_decode_steps,
          do_sample = False,
          pad_token_id=pipe.tokenizer.eos_token_id,
      )
 
    return outputs 

def call_llama3_func(
    inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=200, temperature=0.0
):
  outputs = []
  # for input_str in inputs:
  output = call_llama3_single_prompt(
      inputs,
      model=model,
      max_decode_steps=max_decode_steps,
      temperature=temperature
  )

  for item in output:
      outputs.append([item[0]["generated_text"][-1]["content"]])

  return outputs




def gen_prompt(# 分解结果产生
    query, 
    instruction):
    prompt = []

    if instruction:
        prompt.append({"role": "system", "content": instruction})
    
    
    prompt.append({"role": "system", "content":"Instruction: Split the question into key fragments separated by |. Do not generate sub-queries, rewrite the question, or include explanations. For example, if the input is 'What color is the Santa Anita Park logo?', output 'Santa Anita Park| logo'. Generating sub-queries like 'What color is the logo?' is incorrect."})
    # prompt.append({'role': 'user','content':'What color is the Santa Anita Park logo?'})
    # prompt.append({'role': 'assistant','content':'Santa Anita Park| logo'})
           
    prompt.append({"role": "user", "content": query})

    return prompt





def get_sub_query(query):
    assert isinstance(query,str)
    message = [gen_prompt(query, instruction)]
    
    print(message)
    
    tmp_ls = call_llama3_func(message)[0][0].replace("\n", "").split('|')
    
    # sub_q_ls = [[list(set([item.strip() for item in tmp_ls if item.strip() and item.strip() in query]))]]
    sub_q_ls = [[list(set([item.strip() for item in tmp_ls]))]]
    
    return sub_q_ls
    


counter = {"api" : 0}

# @lru_cache(maxsize=1000000)
# def api_search_query(query, k):
#     print(f"Query={query}")
#     if k == None: k = 10
#     k = min(int(k), 100)
#     pids, ranks, scores = searcher.search(query, k=100)
#     pids, ranks, scores = pids[:k], ranks[:k], scores[:k]
#     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 {"query" : query, "topk": topk}

# @lru_cache(maxsize=1000000)
def api_search_query(query, k):
    print(f"Query={query}")
    
    sub_query_str_l = get_sub_query(query)
    print(sub_query_str_l)
    
    return {"text" : get_ir_result(query, sub_query_str_l)}



@app.route("/api/update_instruction", methods=["POST"])
def update_instruction():
    global instruction  # 使用 global 来更新全局的 instruction 变量

    # 从请求的 JSON 数据中获取新的 instruction
    data = request.get_json()

    if "instruction" in data:
        instruction = data["instruction"]
        return jsonify({"message": "Instruction updated successfully.", "new_instruction": instruction}), 200
    else:
        return jsonify({"message": "Instruction not provided in the request."}), 400


    
@app.route("/api/search", methods=["GET"])
def api_search():
    if request.method == "GET":
        counter["api"] += 1
        print("API request count:", counter["api"])
        return api_search_query(request.args.get("query"), request.args.get("k"))
    else:
        return ('', 405)

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