# app.py
from core.pineconeqa import PineconeQA
from core.voice_processor import AudioRecorder
import streamlit as st
from config import get_settings
import os

def initialize_qa():
    """Initialize QA system with API keys from settings"""
    try:
        settings = get_settings()
        return PineconeQA(
            pinecone_api_key=settings.PINECONE_API_KEY,
            openai_api_key=settings.OPENAI_API_KEY,
            index_name=settings.INDEX_NAME
        )
    except Exception as e:
        st.error(f"Error initializing QA system: {str(e)}")
        return None

def handle_audio_input():
    """Handle audio recording, saving, and transcription"""
    settings = get_settings()
    
    if 'audio_recorder' not in st.session_state:
        st.session_state.audio_recorder = AudioRecorder(settings.OPENAI_API_KEY)
    if 'transcribed_text' not in st.session_state:
        st.session_state.transcribed_text = ""
    
    col1, col2 = st.columns([1, 1])
    
    with col1:
        if st.button("🎤 Record", 
                    help="Click to start recording",
                    type="primary" if not st.session_state.get('recording', False) else "secondary"):
            st.session_state.audio_recorder.start_recording()
            st.session_state.recording = True
            st.info("Recording... Click 'Stop' when finished.")
    
    with col2:
        if st.button("⏹️ Stop", 
                    help="Click to stop recording",
                    disabled=not st.session_state.get('recording', False)):
            audio_file = st.session_state.audio_recorder.stop_recording()
            st.session_state.recording = False
            
            if audio_file:
                with st.spinner("Transcribing audio..."):
                    # Transcribe audio
                    transcription = st.session_state.audio_recorder.transcribe_audio(audio_file)
                    if transcription:
                        # Store transcription in session state
                        st.session_state.transcribed_text = transcription
                        st.success("Audio transcribed successfully!")
                        # Play audio for verification
                        st.audio(audio_file)
                    else:
                        st.error("Failed to transcribe audio.")

def process_question(question):
    st.session_state.chat_history = []
    with st.spinner("Thinking..."):
        response = st.session_state.qa_system.ask(question)
        
        if "error" in response:
            st.error(f"Error: {response['error']}")
        else:
            # Display answer
            st.markdown("### Answer:")
            st.write(response["answer"])
            
            # Display sources
            with st.expander("View Sources"):
                for i, doc in enumerate(response["context"], 1):
                    st.markdown(f"**Source {i}:**")
                    st.write(f"Content: {doc.page_content}")
                    st.write(f"Source: {doc.metadata.get('source', 'Unknown')}")
                    st.write(f"Title: {doc.metadata.get('title', 'Unknown')}")
                    st.write(f"Keywords: {doc.metadata.get('keywords', 'N/A')}")
                    st.write(f"Publication Date: {doc.metadata.get('publication_date', 'Unknown')}")
                    st.markdown("---")
            
            # Add to chat history
            # st.session_state.chat_history.append((question, response["answer"]))

def main():
    st.title("Scientific Paper Q&A System")
    
    # Initialize session state variables
    if 'qa_system' not in st.session_state:
        st.session_state.qa_system = initialize_qa()
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []
    if 'recording' not in st.session_state:
        st.session_state.recording = False
    if 'transcribed_text' not in st.session_state:
        st.session_state.transcribed_text = ""
    
    # Main chat interface
    if st.session_state.qa_system:
        # Chat container
        chat_container = st.container()
        
        with chat_container:
            # Display chat history
            for i, (question, answer) in enumerate(st.session_state.chat_history):
                st.markdown(f"**You:** {question}")
                st.markdown(f"**Assistant:** {answer}")
                st.markdown("---")
        
        # Input container at the bottom
        with st.container():
            st.markdown("### Ask a Question")
            
            # Text input area with transcribed text as default
            question = st.text_area(
                "Type your question or use voice input:",
                value=st.session_state.transcribed_text,
                height=100,
                key="question_input",
                label_visibility="collapsed"
            )
            
            # Audio recording interface
            handle_audio_input()
            
            # Ask button
            if st.button("Send Question", type="primary"):
                if question:
                    process_question(question)
                    # Clear transcribed text after sending
                    st.session_state.transcribed_text = ""
                else:
                    st.warning("Please enter a question or record your voice.")
            
            # Clear chat button
            if st.button("Clear Chat"):
                st.session_state.chat_history = []
                st.session_state.transcribed_text = ""
                st.rerun()
    
    else:
        st.error("Could not initialize QA system. Please check your environment variables.")


if __name__ == "__main__":
    main()