import os
import json
from tqdm import tqdm
import openai
import traceback
import threading
from concurrent.futures import ThreadPoolExecutor
import json_repair

lock = threading.Lock()

client = openai.Client(
    api_key = os.environ.get("DEEPSEEK_API_KEY", "[DEEPSEEK_API_KEY]"),
    base_url = "https://api.deepseek.com/v1",
)

with open("output1.json", "r") as f:
    data = json.load(f)

new_data = {}
if os.path.exists("output.json"):
    with open("output.json", "r") as f:
        new_data = json.load(f)

progress = tqdm(data.items(), total=len(data), leave=True, position=0)
with ThreadPoolExecutor(max_workers=8) as executor:
    for key, value in data.items():
        if key in new_data and new_data[key]:
            progress.update(1)
            continue
        def process_and_update(key, value):
            try:
                prompt = "Given the following description of objects in an image, format the information as a JSON array where each object has 'name', 'bbox', and 'confidence' fields. The 'bbox' should be a list of four integers representing [x_min, y_min, width, height]. Here is the description: " + value + "\n\nOutput Format:\nOutput in JSON format: [{\"name\": object_name1, \"bbox\": bounding_box1, \"confidence\": confidence1}, {\"name\": object_name2, \"bbox\": bounding_box2, \"confidence\": confidence2}].\nDO NOT OUTPUT any other content besides JSON. If the name or bbox of certain objects cannot be determined, skip them in the output."
                response = client.chat.completions.create(
                    model="deepseek-chat",
                    messages=[
                        {
                            "role": "user",
                            "content": prompt
                        }
                    ],
                    temperature=0.1,
                )
                resp_json = response.choices[0].message.content
                resp_dict = json_repair.loads(resp_json) if resp_json else ""
                new_data[key] = resp_dict
            except Exception as e:
                print(f"Error processing key {key}: {e}")
                print(traceback.format_exc())
                new_data[key] = ""
            finally:
                progress.update(1)
                with lock:
                    with open("output.json", "w") as f:
                        json.dump(new_data, f, indent=4, ensure_ascii=False)
        executor.submit(process_and_update, key, value)