main.py

import streamlit as st
import os
import json
from dotenv import load_dotenv
from src.audio_processor import AudioProcessor
from src.llm_engine import LLMEngine
from src.utils import save_uploaded_file, cleanup_temp_files

# Load environment variables
load_dotenv()

# Page Config
st.set_page_config(
    page_title="RelayInsights",
    page_icon="🎙️",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Load CSS
def local_css(file_name):
    with open(file_name) as f:
        st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)

if os.path.exists("assets/styles.css"):
    local_css("assets/styles.css")

# Load Prompts
def load_prompts():
    try:
        with open("config/prompts.json", "r") as f:
            return json.load(f)
    except FileNotFoundError:
        st.error("Config file not found. Please ensure config/prompts.json exists.")
        return []

prompts = load_prompts()

# Initialize Session State
if "transcript" not in st.session_state:
    st.session_state.transcript = ""
if "insights" not in st.session_state:
    st.session_state.insights = ""
if "processed_file" not in st.session_state:
    st.session_state.processed_file = None

# Sidebar
with st.sidebar:
    st.title("🎙️ RelayInsights")
    st.markdown("---")
    
    # API Key Management
    api_key_input = st.text_input("OpenAI API Key", type="password", placeholder="sk-...", help="Your OpenAI API Key is required.")
    
    # Prioritize input key, then env key
    api_key = api_key_input if api_key_input else os.getenv("OPENAI_API_KEY")
    
    if not api_key:
        st.warning("⚠️ Please provide an OpenAI API Key to proceed.")
        
    st.markdown("---")
    st.header("Configuration")
    
    # Model Selection
    model_option = st.selectbox(
        "Select Model",
        options=["gpt-4o", "gpt-4-turbo", "gpt-4", "gpt-3.5-turbo"],
        index=0
    )
    
    # Template Selection
    template_names = [p["name"] for p in prompts]
    selected_template_name = st.selectbox("Select Insight Template", options=template_names)
    
    # Find selected template object
    selected_template = next((p for p in prompts if p["name"] == selected_template_name), None)
    
    if selected_template:
        st.caption(f"**Description:** {selected_template['template'][:100]}...")

    st.markdown("---")
    if st.button("Reset App"):
        for key in list(st.session_state.keys()):
            del st.session_state[key]
        cleanup_temp_files()
        st.rerun()

# Main Interface
st.title("AI Audio Insights & Analytics")
st.markdown("Transform your audio files into actionable intelligence with AI-powered transcription and analysis.")

if not api_key:
    st.stop()

# Initialize Processors
try:
    audio_processor = AudioProcessor(api_key)
    llm_engine = LLMEngine(api_key)
except Exception as e:
    st.error(f"Failed to initialize AI engines: {e}")
    st.stop()

# File Upload Section
st.subheader("Input Source")
input_method = st.radio("Choose input method:", ["Upload File", "Record Audio"], horizontal=True)

if input_method == "Upload File":
    uploaded_file = st.file_uploader("Upload Audio File (MP3, WAV, M4A)", type=["mp3", "wav", "m4a", "txt"])

    if uploaded_file:
        # Check if we need to re-process (new file uploaded)
        if st.session_state.processed_file != uploaded_file.name:
            # Reset state for new file
            st.session_state.transcript = ""
            st.session_state.insights = ""
            st.session_state.processed_file = uploaded_file.name
            
            # Determine file type
            file_ext = uploaded_file.name.split('.')[-1].lower()
            
            if file_ext in ["mp3", "wav", "m4a"]:
                with st.status("Processing Audio...", expanded=True) as status:
                    st.write("Saving file...")
                    temp_path = save_uploaded_file(uploaded_file)
                    
                    st.write("Transcribing audio (this may take a moment)...")
                    try:
                        transcript_text = audio_processor.transcribe(temp_path)
                        st.session_state.transcript = transcript_text
                        st.success("Transcription Complete!")
                    except Exception as e:
                        st.error(f"Error during transcription: {e}")
                        status.update(label="Failed", state="error")
                    finally:
                        pass
                        
            elif file_ext == "txt":
                st.session_state.transcript = uploaded_file.read().decode("utf-8")

elif input_method == "Record Audio":
    from streamlit_mic_recorder import mic_recorder
    
    st.markdown("Click the microphone to start recording. Click again to stop.")
    audio = mic_recorder(
        start_prompt="Start Recording",
        stop_prompt="Stop Recording",
        key='recorder'
    )
    
    if audio:
        # Check if this is a new recording
        audio_id = f"recording_{len(audio['bytes'])}" 
        if st.session_state.processed_file != audio_id:
            st.session_state.transcript = ""
            st.session_state.insights = ""
            st.session_state.processed_file = audio_id
            
            with st.status("Processing Recording...", expanded=True) as status:
                st.write("Saving recording...")
                # Save bytes to a temp wav file
                temp_filename = "temp_recording.wav"
                with open(temp_filename, "wb") as f:
                    f.write(audio['bytes'])
                
                temp_path = os.path.abspath(temp_filename)
                
                st.write("Transcribing audio...")
                try:
                    transcript_text = audio_processor.transcribe(temp_path)
                    st.session_state.transcript = transcript_text
                    st.success("Transcription Complete!")
                except Exception as e:
                    st.error(f"Error during transcription: {e}")
                    status.update(label="Failed", state="error")

# Display Results using Tabs
if st.session_state.transcript:
    tab1, tab2 = st.tabs(["📝 Transcript", "🧠 AI Insights"])
    
    with tab1:
        st.markdown("### Transcript")
        st.text_area("Full Text", value=st.session_state.transcript, height=400)
        
        st.download_button(
            label="Download Transcript",
            data=st.session_state.transcript,
            file_name=f"transcript_{st.session_state.processed_file}.txt",
            mime="text/plain"
        )
        
    with tab2:
        st.markdown(f"### {selected_template['name']}")
        
        if st.session_state.insights and st.button("Regenerate Insights"):
             st.session_state.insights = "" # Clear to trigger gen
             
        if not st.session_state.insights:
            if st.button("Generate Insights", type="primary"):
                with st.spinner("Analyzing transcript..."):
                    try:
                        insights = llm_engine.generate_insights(
                            st.session_state.transcript, 
                            selected_template['template'],
                            model=model_option
                        )
                        st.session_state.insights = insights
                    except Exception as e:
                        st.error(f"Error generating insights: {e}")
        
        if st.session_state.insights:
            st.markdown("---")
            st.markdown(st.session_state.insights)
            st.markdown("---")
            
            st.download_button(
                label="Download Insights",
                data=st.session_state.insights,
                file_name=f"insights_{selected_template['id']}.md",
                mime="text/markdown"
            )

# Footer
st.markdown("---")
st.markdown(
    """
    <div style='text-align: center; color: #666;'>
        <p>Built with Streamlit & OpenAI • RelayInsights</p>
    </div>
    """,
    unsafe_allow_html=True
)