import json
import uuid


data = {
    "audio_url": "https://data.heartex.net/open-images/dictionary.wav",
    "transcript": "hello world this is a test",
    "words": [
        {"word": "hello", "start_time": 0.1, "end_time": 0.5},
        {"word": "world", "start_time": 0.6, "end_time": 1.0},
        {"word": "this",  "start_time": 1.1, "end_time": 1.4},
        {"word": "is",    "start_time": 1.5, "end_time": 1.7},
        {"word": "a",     "start_time": 1.8, "end_time": 1.9},
        {"word": "test",  "start_time": 2.0, "end_time": 2.5},
    ]
}

def generate_label_studio_json(data):
    transcript = data["transcript"]
    results = []
    current_char_idx = 0
    
    for word_info in data["words"]:
        word = word_info["word"]
        start_time = word_info["start_time"]
        end_time = word_info["end_time"]
        
        start_char = transcript.find(word, current_char_idx)
        if start_char == -1:
            continue 
            
        end_char = start_char + len(word)
        current_char_idx = end_char
        
        # Generate unique IDs for the text and audio regions
        text_id = f"text_{uuid.uuid4().hex[:8]}"
        audio_id = f"audio_{uuid.uuid4().hex[:8]}"
        
        # 1. Text Region
        text_region = {
            "id": text_id,
            "from_name": "text_labels", # Must match <Labels name="text_labels">
            "to_name": "transcript",    # Must match <Text name="transcript">
            "type": "labels",
            "value": {
                "start": start_char,
                "end": end_char,
                "text": word,
                "labels": ["Word"]      # Must match <Label value="Word">
            }
        }
        
        # 2. Audio Region
        audio_region = {
            "id": audio_id,
            "from_name": "labels",      # Must match <Labels name="labels">
            "to_name": "audio",         # Must match <Audio name="audio">
            "type": "labels",
            "value": {
                "start": start_time,
                "end": end_time,
                "labels": ["Word"]      # Must match <Label value="Word">
            }
        }
        
        # 3. Relation linking the Text Region and Audio Region
        relation = {
            "from_id": text_region["id"],
            "to_id": audio_region["id"],
            "type": "relation",
            "direction": "right"
        }
        
        # Add all three to the results array
        results.extend([text_region, audio_region, relation])
        
    # Construct the final Task object
    task = {
        "data": {
            "audio": data["audio_url"],
            "transcript": transcript
        },
        "predictions": [
            {
                "model_version": "auto_aligner_v1",
                "result": results
            }
        ]
    }
    
    return [task] # Label Studio expects an array of tasks

if __name__ == "__main__":
    ls_tasks = generate_label_studio_json(data)
    
    output_filename = "pre_annotations.json"
    with open(output_filename, "w") as f:
        json.dump(ls_tasks, f, indent=2)
        
    print(f"Successfully generated {output_filename}")
    print("You can now import this file into Label Studio!")