import streamlit as st
import torch

def check_password():
    def password_entered():
        if st.session_state["password"] == st.secrets["PASSWORD"]:
            st.session_state["password_correct"] = True
            del st.session_state["password"]  # удаляем пароль из сессии для безопасности
        else:
            st.session_state["password_correct"] = False
    if "password_correct" not in st.session_state:
        # Первоначальное состояние
        st.session_state["password_correct"] = False
    if not st.session_state["password_correct"]:
        # Показываем поле для ввода пароля и кнопку подтверждения
        st.title('Audio Transcription App')
        st.text_input("Введите пароль:", type="password", on_change=password_entered, key="password")
        return False
    else:
        return True

def get_gpu_memory_info():
    if torch.cuda.is_available():
        num_gpus = torch.cuda.device_count()
        memory_info = []
        for i in range(num_gpus):
            gpu_name = torch.cuda.get_device_name(i)
            gpu_memory_total = torch.cuda.get_device_properties(i).total_memory / (1024 ** 3)  # в гигабайтах
            gpu_memory_reserved = torch.cuda.memory_reserved(i) / (1024 ** 3)  # в гигабайтах
            gpu_memory_allocated = torch.cuda.memory_allocated(i) / (1024 ** 3)  # в гигабайтах
            gpu_memory_free = gpu_memory_reserved - gpu_memory_allocated

            memory_info.append((gpu_name, gpu_memory_total, gpu_memory_reserved, gpu_memory_allocated, gpu_memory_free))
        print(f"{'GPU Name':<25}{'Total Memory (GB)':<20}{'Reserved Memory (GB)':<22}{'Allocated Memory (GB)':<23}{'Free Memory (GB)':<20}")
        print("="*110)

    # Печать информации о каждой GPU
        for gpu_name, total, reserved, allocated, free in memory_info:
            print(f"{gpu_name:<25}{total:<20.2f}{reserved:<22.2f}{allocated:<23.2f}{free:<20.2f}")

    else:
        print("CUDA is not available.")

def convert_segments_object_to_text(data):
    result = []
    print(data)

    for segment in data['segments']:
        words = segment['words']
        segment_speaker = segment.get('speaker', None)
        segment_start = segment.get('start', None)
        segment_end = segment.get('end', None)
        current_speaker = None
        current_start = None
        current_end = None
        current_text = []

        # Forward fill speaker, start and end if missing
        for i, word_info in enumerate(words):
            if 'speaker' not in word_info:
                if i > 0 and 'speaker' in words[i - 1]:
                    word_info['speaker'] = words[i - 1]['speaker']
                elif i < len(words) - 1 and 'speaker' in words[i + 1]:
                    word_info['speaker'] = words[i + 1]['speaker']
                else:
                    word_info['speaker'] = segment_speaker
                    
            if 'start' not in word_info:
                if i > 0 and 'end' in words[i - 1]:
                    word_info['start'] = words[i - 1]['end']
                else:
                    word_info['start'] = segment_start
                    
            if 'end' not in word_info:
                if i < len(words) - 1 and 'start' in words[i + 1]:
                    word_info['end'] = words[i + 1]['start']
                elif i == len(words) - 1:
                    word_info['end'] = segment_end
                else:
                    word_info['end'] = word_info['start']

        for word_info in words:
            word = word_info.get('word', '')
            start = word_info.get('start', None)
            end = word_info.get('end', None)
            speaker = word_info.get('speaker', None)

            if current_speaker is None:
                current_speaker = speaker
                current_start = start

            if speaker == current_speaker:
                current_text.append(word)
                current_end = end
            else:
                # Finish current segment
                if current_start is not None and current_end is not None:
                    formatted_text = f'{current_speaker} ({current_start} : {current_end}) : {" ".join(current_text)}'
                else:
                    formatted_text = f'{current_speaker} : {" ".join(current_text)}'
                result.append(formatted_text)

                # Start new segment
                current_speaker = speaker
                current_start = start
                current_end = end
                current_text = [word]

        # Append the last segment
        if current_text:
            if current_start is not None and current_end is not None:
                formatted_text = f'{current_speaker} ({current_start} : {current_end}) : {" ".join(current_text)}'
            else:
                formatted_text = f'{current_speaker} : {" ".join(current_text)}'
            result.append(formatted_text)

    return '\n'.join(result)

def convert_segments_object_to_text_simple(data):
    result = []

    for segment in data['segments']:
        text = segment['text']
        current_speaker = segment.get('speaker', None)
        segment_start = segment.get('start', None)
        segment_end = segment.get('end', None)

        result.append(f'{current_speaker} ({segment_start} : {segment_end}) : {text}')
        
    return '\n'.join(result)