# coding=utf8
import os
import shutil
import sys
import subprocess
import xml.etree.ElementTree as ET
from transformers import pipeline
from environment import DEFAULT_MODEL_LANGUAGE, DEFAULT_MODEL, DEFAULT_LANGUAGE, DEVICE

# class for annotation segments
class Segment:
    def __init__(self, segment_id: str, start: int, end: int):
        self.segment_id:str = segment_id
        self.start:int = start
        self.end:int = end
        self.transcription:str = ''
        #self.ts_start:str = ts_start
        #self.ts_end:str = ts_end

def getTimeSlots(eaf):
    time_slot_dic = {}

    order = eaf.find('TIME_ORDER')
    for slot in order:
        time_slot_dic[slot.get('TIME_SLOT_ID')] = slot.get('TIME_VALUE')
    
    return time_slot_dic

def getAnnotationSegments(eaf, tier_type):
    segment_list = []

    time_slot_dic = getTimeSlots(eaf)

    for tier in eaf.findall('TIER'):
        if tier.get('LINGUISTIC_TYPE_REF') == tier_type:
            for annotation in tier:
                alignable_annotation = annotation.find('ALIGNABLE_ANNOTATION')
                segment_id = alignable_annotation.get('ANNOTATION_ID')
                start = time_slot_dic[alignable_annotation.get('TIME_SLOT_REF1')]
                end = time_slot_dic[alignable_annotation.get('TIME_SLOT_REF2')]
                segment_list.append(Segment(segment_id, start, end))

    return segment_list

def splice_audio(audio_path, segment_id, start, end, temp_dir):
    file_path = f"{temp_dir}/{segment_id}.wav"

    if os.path.exists(file_path):
        os.remove(file_path)

    subprocess.call([
        "ffmpeg", 
        "-loglevel", "fatal", 
        "-hide_banner", 
        "-nostdin",
        "-i", audio_path,
        "-ss", f"{int(start)/1000}",
        "-to", f"{int(end)/1000}",
        file_path
    ])
    return f"{temp_dir}/{segment_id}.wav"


# transcribes a single and returns the transcription
def transcribe_audio(model_id, audio_path):
    transcribe = pipeline(
        task = "automatic-speech-recognition",
        model = model_id,
        chunk_length_s = 30,
        device = DEVICE,
    )
    
    transcribe.model.config.forced_decoder_ids = transcribe.tokenizer.get_decoder_prompt_ids(language='bengali', task="transcribe")
    
    result = transcribe(audio_path, max_new_tokens=448)
    
    transcription = result["text"].strip()
    
    #print(f"Transcription for {audio_path}: {transcription}")
    
    return transcription


def transcribe_eaf(eaf_path, audio_path, tier_type):

    eaf_tree = ET.parse(eaf_path)
    eaf_root = eaf_tree.getroot()
    
    segment_list = getAnnotationSegments(eaf_root, tier_type)

    if not os.path.exists('temp_dir'):
        os.makedirs('temp_dir')

    for segment in segment_list:
        # get the audio segment
        segment_audio_file = splice_audio(audio_path, segment.segment_id, segment.start, segment.end, 'temp_dir')
        segment.transcription = transcribe_audio(DEFAULT_MODEL, segment_audio_file)
        os.remove(segment_audio_file)

        print(f'{segment.segment_id}\t{segment.transcription}')
    shutil.rmtree('temp_dir')

    for segment in segment_list:
        for e in eaf_root.iter():
            if e.tag == 'ALIGNABLE_ANNOTATION' and e.get('ANNOTATION_ID') == segment.segment_id:
                e.find('ANNOTATION_VALUE').text = segment.transcription

    new_eaf_path = f'{eaf_path[:-4]}_autotranscribed.eaf'
    eaf_tree.write(new_eaf_path, encoding='utf-8', xml_declaration=True)

    return new_eaf_path