import os, sys
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', '..', 'document_retrieval', 'Decompose_retrieval'))
import ragqa_paths  # [ragqa] portable paths
from transformers import pipeline
import torch
import pandas as pd
from transformers import AutoTokenizer
from openai import OpenAI
import requests

client = OpenAI(api_key="0",base_url="http://0.0.0.0:50001/v1")
tokenizer = AutoTokenizer.from_pretrained(ragqa_paths.LLAMA_MODEL)


import re
import csv
from io import StringIO

def parse_string(s):
    s = s.strip()
    if s.startswith('[') and s.endswith(']'):
        # 处理列表结构
        inner = s[1:-1].strip()
        # 将单引号包裹的元素替换为双引号包裹，并转义内部双引号
        pattern = re.compile(r"'((?:[^'\\]|\\.)*?)'")
        def replace(match):
            content = match.group(1)
            content = content.replace('"', r'\"')
            return f'"{content}"'
        new_inner = pattern.sub(replace, inner)
        # 使用 csv.reader 解析处理后的内容
        csv_reader = csv.reader(
            StringIO(new_inner),
            quotechar='"',
            escapechar='\\',
            skipinitialspace=True
        )
        try:
            return next(csv_reader)
        except StopIteration:
            return []
    else:
        # 处理单个字符串，包裹为双引号并转义内部双引号
        content = s.replace('"', r'\"')
        csv_reader = csv.reader(
            StringIO(f'"{content}"'),
            quotechar='"',
            escapechar='\\'
        )
        try:
            return next(csv_reader)
        except StopIteration:
            return [s]

def cover_em(pred_ls, ans_ls):
    assert len(pred_ls) == len(ans_ls)
    cnt = 0
    for idx in range(len(pred_ls)):
        pred = pred_ls[idx].lower()
       
        if dataset_name in {"manyqa_text"}:
            ans = str(ans_ls[idx])
        else:
            ans = parse_string(ans_ls[idx])

        if dataset_name in {"manyqa_text"}:
            if ans.lower() in pred:
                cnt += 1
        else:
            for j in range(len(ans)):
                if ans[j].lower() in pred:
                    cnt += 1
                    break
                
    return cnt/len(pred_ls)
    

def get_vllm_llama(temperature, max_tokens, chats):
    url = "http://127.0.0.1:60000/ask"

    data = {
        "temperature": temperature,
        "max_tokens": max_tokens,
        "chats": chats,
    }
    
    response = requests.post(url, json=data, timeout=None)
    response_data = response.json()

    passage_ls = response_data.get("output", [])
    return passage_ls

# def call_llama3_single_prompt(
#     inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=20, temperature=0.8
# ):
      
#     inputs_ls = []
#     if isinstance(inputs, str):
#         messages = [
#             {"role": "user", "content": inputs},
#         ]
#         inputs_ls.append(tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True))

#     else:
#       for idx in range(len(inputs)):
#           inputs_ls.append(tokenizer.apply_chat_template(inputs[idx], tokenize=False, add_generation_prompt=True))
    
#     # ans = get_vllm_llama(temperature, max_decode_steps, inputs_ls)
#     pred_ls = [row[0] for row in call_llama3_func(inputs_ls, max_decode_steps=100)]
        
#     return pred_ls


def call_llama3_single_prompt(
    inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=20, temperature=0.0
):
      
    inputs_ls = []
    if isinstance(inputs, str):
        messages = [
            {"role": "user", "content": inputs},
        ]
        inputs_ls.append(tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True))

    else:
      for idx in range(len(inputs)):
          inputs_ls.append(tokenizer.apply_chat_template(inputs[idx], tokenize=False, add_generation_prompt=True))
    
    # ans = get_vllm_llama(temperature, max_decode_steps, inputs_ls)
    results = client.completions.create(
            model=ragqa_paths.LLAMA_MODEL,
            max_tokens=max_decode_steps,
            temperature=0,
            prompt=inputs_ls,
            timeout = None
    )
    ans = []
    for item in results.choices:
        ans.append([item.text.strip()])
        
    return ans


def call_llama3_func(
    inputs, model="Llama-3.1-8B-Instruct", max_decode_steps=20, temperature=0.0
):

  print(max_decode_steps, temperature)
  
  output = call_llama3_single_prompt(
      inputs,
      model=model,
      max_decode_steps=max_decode_steps,
      temperature=temperature
  )

  if isinstance(inputs, str):
    return output[0]
  else:
    return output


def get_lora_ans(queries_ls):
    prompts = []
    for i in range(len(queries_ls)):
        prompts.append([
            {"role": "user", "content": 'Question: ' + queries_ls[i]},
        ])
        
    prompts_tokened = [tokenizer.apply_chat_template(x, tokenize=False, add_generation_prompt=True) for x in prompts]
    
    results = client.completions.create(
            model="web_ft",
            max_tokens=100,
            temperature=0,
            prompt=prompts_tokened
    )
    
    ans = []
    for item in results.choices:
        ans.append(item.text.strip())
    
    return ans

def lora_evaluate(dataset_name):
    dataset_path = ragqa_paths.dataset_file(dataset_name, f"{dataset_name}_test.csv")

    raw_data = pd.read_csv(dataset_path, header=0)

    raw_queries = list(raw_data['question'])
    true_answers = list(raw_data['answers'])
    
    pred_ls = get_lora_ans(raw_queries)

    print(f"lora score: {cover_em(pred_ls, true_answers)}")


def evaluate(dataset_name):
    dataset_path = ragqa_paths.dataset_file(dataset_name, f"{dataset_name}_test.csv")

    raw_data = pd.read_csv(dataset_path, header=0)
    
    # raw_data = raw_data.head(5)
    
    raw_queries = [q.strip() + '?' if not q.strip().endswith('?') else q.strip() for q in raw_data['question']]
    true_answers = list(raw_data['answers'])
    
    prompts = []
    for i in range(len(raw_queries)):  # You are a helpful assistant. "You are a helpful assistant. Answer the question directly and concisely. Do not include explanations or extra information beyond the question's requirements."
        prompts.append([
            {"role": "system", "content": "You are a helpful assistant. answer the question."},
            {"role": "user", "content": 'Question: ' + raw_queries[i]},
        ])
    
    pred_ls = [row[0] for row in call_llama3_func(prompts, max_decode_steps=100)]

    print(f"score: {cover_em(pred_ls, true_answers)}")
    
if __name__ == '__main__':
    dataset_name = 'manyqa_text'

    # lora_evaluate(dataset_name)
    evaluate(dataset_name)