import os
import sys
import re
import jsonlines
import json
import time
import os 
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
from module_03_constraint.prompt_template import constraint_rewrite_user_prompt, constraint_rewrite_system_prompt

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_constraints):
    start = time.time()

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

        for line in json_list:
            system_prompt = constraint_rewrite_system_prompt
            user_prompt = constraint_rewrite_user_prompt.format(constraint = line["constraint"], num = num_constraints)
            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

            final_output = process_output_text(output)
            w.write(line)
            try:
                for item in final_output:
                    if item != "":
                        line["constraint"] = item
                        w.write(line)
            except:
                print("Error")
    w.close
    
    end = time.time()
    print(f"Time taken: {end-start}")

if __name__ == '__main__':
    jsonl_input = "module_03_constraint/constraint_seeds/test.jsonl"
    jsonl_output = "module_03_constraint/constraint_seed_gpt/test.jsonl"
    num_constraints = 2
    openai_inference(jsonl_input, jsonl_output, num_constraints)