import json
import uuid

# Example data: transcript and corresponding SENTENCE timestamps
data = {
    "audio_url": "http://localhost:8000/my_audio.wav",
    "transcript": "Hello world this is a test. This is the second sentence.",
    "sentences": [
        {"text": "Hello world this is a test.", "start_time": 0.1, "end_time": 2.5},
        {"text": "This is the second sentence.", "start_time": 2.8, "end_time": 5.0}
    ]
}

def generate_sentence_annotations(data):
    transcript = data["transcript"]
    results = []
    current_char_idx = 0
    
    for sentence_info in data["sentences"]:
        sentence_text = sentence_info["text"]
        start_time = sentence_info["start_time"]
        end_time = sentence_info["end_time"]
        
        # Find the character offset for the full sentence
        start_char = transcript.find(sentence_text, current_char_idx)
        if start_char == -1:
            continue
            
        end_char = start_char + len(sentence_text)
        current_char_idx = end_char
        
        text_id = f"text_{uuid.uuid4().hex[:8]}"
        audio_id = f"audio_{uuid.uuid4().hex[:8]}"
        
        # 1. Text Region (Using "Sentence" label)
        text_region = {
            "id": text_id,
            "from_name": "text_labels", 
            "to_name": "transcript",    
            "type": "labels",
            "value": {
                "start": start_char,
                "end": end_char,
                "text": sentence_text,
                "labels": ["Sentence"]  # Changed from "Word" to "Sentence"
            }
        }
        
        # 2. Audio Region (Using "Sentence" label)
        audio_region = {
            "id": audio_id,
            "from_name": "labels",      
            "to_name": "audio",         
            "type": "labels",
            "value": {
                "start": start_time,
                "end": end_time,
                "labels": ["Sentence"]  # Changed from "Word" to "Sentence"
            }
        }
        
        # 3. Relation linking them
        relation = {
            "from_id": text_region["id"],
            "to_id": audio_region["id"],
            "type": "relation",
            "direction": "right"
        }
        
        results.extend([text_region, audio_region, relation])
        
    task = {
        "data": {
            "audio": data["audio_url"],
            "transcript": transcript
        },
        "predictions": [
            {
                "model_version": "auto_aligner_v1",
                "result": results
            }
        ]
    }
    
    return [task]

if __name__ == "__main__":
    ls_tasks = generate_sentence_annotations(data)
    with open("pre_annotations_sentences.json", "w") as f:
        json.dump(ls_tasks, f, indent=2)
    print("Successfully generated pre_annotations_sentences.json")