app.py

import streamlit as st
from dotenv import load_dotenv
import os
from agents import create_research_crew
from utils import process_crew_output, display_report
from report_generator import generate_direct_report, create_pdf_report  # Add this import

load_dotenv()

st.set_page_config(
    page_title="Market Research Generator",
    page_icon="📊",
    layout="wide"
)

def init_session_state():
    if 'report' not in st.session_state:
        st.session_state.report = None
    if 'generating' not in st.session_state:
        st.session_state.generating = False
    if 'logs' not in st.session_state:
        st.session_state.logs = []
    if 'direct_report' not in st.session_state:
        st.session_state.direct_report = None

def get_api_keys():
    try:
        openai_key = st.secrets["OPENAI_API_KEY"]
        serper_key = st.secrets["SERPER_API_KEY"]
        return openai_key, serper_key
    except Exception as e:
        st.error(f"Error loading API keys: {str(e)}")
        return None, None

def main():
    init_session_state()
    
    st.title("🤖 AI Market Research Generator")
    
    openai_key, serper_key = get_api_keys()
    if not openai_key or not serper_key:
        st.error("API keys not found in secrets.")
        return
        
    os.environ["OPENAI_API_KEY"] = openai_key
    os.environ["SERPER_API_KEY"] = serper_key
    
    tab1, tab2, tab3, tab4 = st.tabs([
        "Generate Report", 
        "View Report", 
        "Direct Research",
        "Agent Interaction"
    ])
    
    with tab1:
        st.subheader("Research Requirements")
        col1, col2 = st.columns([2, 1])
        
        with col1:
            topic = st.text_input(
                "Research Topic",
                placeholder="e.g., Electric Vehicles Market"
            )
            
            questions = st.text_area(
                "Key Research Questions (one per line)",
                placeholder="1. What is the market size?\n2. Who are the key competitors?\n3. What are the main trends?",
                height=150
            )
            
            if st.button("Generate Report", type="primary", disabled=st.session_state.generating):
                if not topic or not questions:
                    st.error("Please enter both topic and questions")
                    return
                
                st.session_state.generating = True
                st.session_state.logs = []
                
                try:
                    progress = st.progress(0)
                    status = st.empty()
                    
                    # Initialize research
                    status.text("🚀 Initializing research...")
                    st.session_state.logs.append("🚀 Starting market research...")
                    
                    crew = create_research_crew(topic, questions.split('\n'))
                    progress.progress(25)
                    
                    # Research phase
                    status.text("🔍 Gathering market data...")
                    st.session_state.logs.append("👨‍🔬 Researcher collecting data...")
                    results = crew.kickoff()
                    progress.progress(50)
                    
                    # Analysis and report generation
                    status.text("📊 Processing insights...")
                    st.session_state.logs.append("📊 Analyzing market data...")
                    enhanced_report = process_crew_output(results, topic)
                    progress.progress(75)
                    
                    # Create PDF
                    status.text("📄 Generating PDF report...")
                    pdf_content = create_pdf_report(enhanced_report['content'], topic)
                    enhanced_report['pdf_content'] = pdf_content
                    progress.progress(100)
                    
                    # Completion
                    status.text("✅ Report ready!")
                    st.session_state.logs.append("✅ Report generation complete!")
                    
                    st.session_state.report = enhanced_report
                    st.success("Report generated successfully! View in Report tab")
                    
                except Exception as e:
                    st.error(f"Error: {str(e)}")
                    st.session_state.logs.append(f"❌ Error: {str(e)}")
                
                finally:
                    st.session_state.generating = False
        
        with col2:
            st.subheader("Progress")
            for log in st.session_state.logs[-5:]:
                st.markdown(log)
    
    with tab2:
        if st.session_state.report:
            st.warning("⚠️ AI-generated report. Please verify critical information.")
            display_report(st.session_state.report)
            
            # Download buttons with styling
            st.markdown('<div class="section-box">', unsafe_allow_html=True)
            col1, col2 = st.columns(2)
            with col1:
                st.download_button(
                    "📥 Download Markdown Report",
                    st.session_state.report['raw'],
                    "market_research_report.md",
                    "text/markdown"
                )
            with col2:
                st.download_button(
                    "📥 Download PDF Report",
                    st.session_state.report['pdf_content'],
                    f"market_research_{topic.lower().replace(' ', '_')}.pdf",
                    "application/pdf"
                )
            st.markdown('</div>', unsafe_allow_html=True)
        else:
            st.info("Generate a report to view results")

    with tab3:
        st.subheader("Direct Research Report")
        if topic:
            try:
                with st.spinner("Generating direct research report..."):
                    # Automatically generate direct report when topic is available
                    if 'direct_report' not in st.session_state or st.session_state.direct_report is None:
                        direct_report = generate_direct_report(topic, questions)
                        st.session_state.direct_report = direct_report
                    
                    # Display the report
                    st.markdown("## 📊 Market Research Report")
                    st.markdown(st.session_state.direct_report['final_report'])
                    
                    # Create and offer PDF download
                    pdf_content = create_pdf_report(st.session_state.direct_report['final_report'], topic)
                    st.download_button(
                        "📥 Download PDF Report",
                        pdf_content,
                        f"direct_research_{topic.lower().replace(' ', '_')}.pdf",
                        "application/pdf"
                    )
            except Exception as e:
                st.error(f"Error generating direct report: {str(e)}")
        else:
            st.info("Enter a topic in the 'Generate Report' tab to view direct research")
    
    with tab4:
        if st.session_state.report and 'agent_outputs' in st.session_state.report:
            st.subheader("🤖 Behind the Scenes: Agent Interactions")
            
            st.markdown("""
            ### Research & Analysis Process
            Watch how our AI agents work together to create your market research report:
            """)
            
            # Timeline visualization
            for idx, (agent, data) in enumerate(st.session_state.report['agent_outputs'].items()):
                with st.container():
                    col1, col2 = st.columns([1, 3])
                    with col1:
                        st.markdown(f"**Step {idx+1}**")
                        st.markdown(f"**{agent.title()}**")
                        st.markdown(f"*{data['timestamp']}*")
                    
                    with col2:
                        st.markdown(f"**Task:** {data['analysis_type']}")
                        with st.expander("View Detailed Analysis"):
                            st.markdown("#### Key Findings")
                            st.markdown(data['raw_output'])
                            if 'processed_data' in data:
                                st.markdown("#### Processed Data")
                                st.json(data['processed_data'])
                
                # Add connector line except for last item
                if idx < len(st.session_state.report['agent_outputs']) - 1:
                    st.markdown("↓")
        else:
            st.info("Generate a report to view agent interactions")

    with tab5:
        if st.session_state.report:
            st.subheader("📽️ Presentation View")
            
            # Initialize slide index in session state if not present
            if 'current_slide' not in st.session_state:
                st.session_state.current_slide = 0
            
            # Create slides from report data
            slides = create_presentation_slides(st.session_state.report)
            
            # Navigation buttons callbacks
            def next_slide():
                if st.session_state.current_slide < len(slides) - 1:
                    st.session_state.current_slide += 1
            
            def prev_slide():
                if st.session_state.current_slide > 0:
                    st.session_state.current_slide -= 1
            
            # Display current slide
            with st.container():
                st.markdown("---")
                display_presentation_slide(
                    slides[st.session_state.current_slide],
                    st.session_state.current_slide,
                    len(slides)
                )
                st.markdown("---")
            
            # Add download presentation button
            if st.button("📥 Download Presentation"):
                presentation_pdf = create_presentation_pdf(slides)
                st.download_button(
                    "Download PDF",
                    presentation_pdf,
                    f"market_research_presentation_{topic.lower().replace(' ', '_')}.pdf",
                    "application/pdf"
                )
        else:
            st.info("Generate a report first to view the presentation")

if __name__ == "__main__":
    main()