import os
import sys
import re
import jsonlines
import json
import time
import os 
import random
from openai import OpenAI
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from utils.set_random_seed import set_random_seed
from utils.json_reader import jsonl_loader, json_loader
from module_06_evaluation.prompt_template import *
client = OpenAI(
    api_key=""
)

def process_output_text(text):
    lines = text.split("\n")
    cleaned_lines = []
    for line in lines:
        cleaned_line = re.sub(r'^\d+\.\s*', '', line.strip())
        cleaned_lines.append(cleaned_line)
    
    return cleaned_lines
    
def openai_inference(jsonl_input, jsonl_output, num_questions):
    start = time.time()

    json_list = json_loader(jsonl_input)
    selected_list = random.sample(json_list, num_questions) 

    with jsonlines.open(jsonl_output, mode='w') as w:

        for line in selected_list:
            system_prompt = system_prompt_template
            user_prompt = question_rewrite_user_prompt_template.format(question = line["new_question"])
            stream = client.chat.completions.create(
            messages=[{"role": "system", "content": system_prompt},
                      {"role": "user", "content": user_prompt}
            ],
            model="gpt-4o",
            stream=True,
            seed = 42)

            output = ""
            for chunk in stream:
                if chunk.choices[0].delta.content is not None:
                    output += chunk.choices[0].delta.content

            w.write({
                    "id": line["id"],
                    "question": line["new_question"],
                    "modified_question": output,
                    "original_answer": line["original_answer"],
                })
    w.close
    
    end = time.time()
    print(f"Time taken: {end-start}")

if __name__ == '__main__':
    jsonl_input = "dataset/lifestyle/lifestyle_04.jsonl"
    jsonl_output = "module_06_evaluation/dataset/lifestyle_rewrite_question.jsonl"
    num_questions = 200
    openai_inference(jsonl_input, jsonl_output, num_questions)