"""
Multi-Speaker Audio Analyzer
A Streamlit application that performs speaker diarization, transcription, and summarization on audio files.

Author: [Your Name]
Date: January 2025
"""

import streamlit as st
from pyannote.audio import Pipeline
import whisper
import tempfile
import os
import torch
from transformers import pipeline as tf_pipeline, BartTokenizer
from pydub import AudioSegment
import io
import pickle

class SpeakerDiarizer:
   def __init__(self, token):
       self.pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization", use_auth_token=token)

   def process(self, audio_file):
       return self.pipeline(audio_file)

class Transcriber:
   def __init__(self):
       self.model = whisper.load_model("base")

   def process(self, audio_file):
       return self.model.transcribe(audio_file)["text"]

class Summarizer:
   def __init__(self, model_path='bart_ami_finetuned.pkl'):
       self.tokenizer = BartTokenizer.from_pretrained('facebook/bart-base')
       with open(model_path, 'rb') as f:
           self.model = pickle.load(f)

   def process(self, text):
       inputs = self.tokenizer(text, return_tensors="pt", max_length=1024, truncation=True)
       summary_ids = self.model.generate(inputs["input_ids"], max_length=150, min_length=40)
       return self.tokenizer.decode(summary_ids[0], skip_special_tokens=True)

@st.cache_resource
def load_models():
   try:
       diarizer = SpeakerDiarizer(st.secrets["hf_token"])
       transcriber = Transcriber()
       summarizer = Summarizer()
       return diarizer, transcriber, summarizer
   except Exception as e:
       st.error(f"Error loading models: {str(e)}")
       return None, None, None

def process_audio(audio_file):
   try:
       audio_bytes = io.BytesIO(audio_file.getvalue())
       
       with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
           if audio_file.name.lower().endswith('.mp3'):
               audio = AudioSegment.from_mp3(audio_bytes)
           else:
               audio = AudioSegment.from_wav(audio_bytes)
           
           audio = audio.set_frame_rate(16000).set_channels(1).set_sample_width(2)
           audio.export(tmp.name, format="wav", parameters=["-ac", "1", "-ar", "16000"])
           tmp_path = tmp.name

           diarizer, transcriber, summarizer = load_models()
           if not all([diarizer, transcriber, summarizer]):
               return "Model loading failed"

           with st.spinner("Processing..."):
               diarization = diarizer.process(tmp_path)
               transcription = transcriber.process(tmp_path)
               summary = summarizer.process(transcription)

           os.unlink(tmp_path)
           
           return {
               "diarization": diarization,
               "transcription": transcription,
               "summary": summary
           }
           
   except Exception as e:
       st.error(f"Error: {str(e)}")
       return None

def main():
   st.title("Multi-Speaker Audio Analyzer")
   st.write("Upload an audio file (MP3/WAV) up to 5 minutes long")

   uploaded_file = st.file_uploader("Choose a file", type=["mp3", "wav"])

   if uploaded_file:
       file_size = len(uploaded_file.getvalue()) / (1024 * 1024)
       st.write(f"File size: {file_size:.2f} MB")
       
       st.audio(uploaded_file, format='audio/wav')
       
       if st.button("Analyze Audio"):
           if file_size > 200:
               st.error("File size exceeds 200MB limit")
           else:
               results = process_audio(uploaded_file)
               
               if results:
                   tab1, tab2, tab3 = st.tabs(["Speakers", "Transcription", "Summary"])
                   
                   with tab1:
                       st.write("Speaker Timeline:")
                       for turn, _, speaker in results["diarization"].itertracks(yield_label=True):
                           col1, col2, col3 = st.columns([2,3,5])
                           with col1:
                               speaker_num = int(speaker.split('_')[1])
                               colors = ['🔵', '🔴']
                               st.write(f"{colors[speaker_num % 2]} {speaker}")
                           with col2:
                               st.write(f"{format_timestamp(turn.start)} → {format_timestamp(turn.end)}")
                   
                   with tab2:
                       st.write("Transcription:")
                       st.write(results["transcription"])
                   
                   with tab3:
                       st.write("Summary:")
                       st.write(results["summary"])

if __name__ == "__main__":
   main()