import streamlit as st
from pydub import AudioSegment
from io import BytesIO
import numpy as np
import plotly.graph_objs as go
import plotly.express as px

# Function to plot the waveform
def plot_waveform(audio, start_time, end_time):
    samples = np.array(audio.get_array_of_samples())
    sample_rate = audio.frame_rate
    duration = len(samples) / sample_rate

    times = np.linspace(0, duration, num=len(samples))

    fig = go.Figure()
    fig.add_trace(go.Scatter(x=times, y=samples, mode='lines', name='Waveform'))
    fig.update_layout(title='Audio Waveform', xaxis_title='Time (s)', yaxis_title='Amplitude')

    fig.add_vrect(x0=start_time, x1=end_time, fillcolor="LightSalmon", opacity=0.5, layer="below", line_width=0)
    return fig

# Streamlit UI
st.title("Audio Cutter")

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

if uploaded_file is not None:
    audio = AudioSegment.from_file(uploaded_file)
    st.audio(uploaded_file, format='audio/wav')
    
    duration = len(audio) / 1000.0
    st.write(f"Audio Duration: {duration:.2f} seconds")

    # Initial slider values
    start_time = st.slider("Start time (in seconds)", 0.0, duration, 0.0)
    end_time = st.slider("End time (in seconds)", 0.0, duration, duration)

    if start_time >= end_time:
        st.error("End time must be greater than start time.")
    else:
        fig = plot_waveform(audio, start_time, end_time)
        st.plotly_chart(fig, use_container_width=True)

        if st.button("Cut Audio"):
            clipped_audio = audio[start_time * 1000:end_time * 1000]

            buf = BytesIO()
            clipped_audio.export(buf, format='wav')
            buf.seek(0)

            st.audio(buf, format='audio/wav')

            st.download_button(
                label="Download Clipped Audio",
                data=buf,
                file_name="clipped_audio.wav",
                mime="audio/wav"
            )