src/streamlit_app.py

import streamlit as st
import asyncio
from workflows_v2 import startup_validation_workflow
import time

# Real-time progress tracking class
class StreamlitProgressTracker:
    def __init__(self):
        self.progress_container = None
        self.status_text = None
        self.progress_bar = None
        self.log_container = None
        self.current_phase = 0
        self.total_phases = 4
        
    def setup_ui(self):
        self.progress_container = st.container()
        with self.progress_container:
            self.status_text = st.empty()
            self.progress_bar = st.progress(0)
            self.log_container = st.container()
    
    def update_progress(self, message):
        """Called by workflow for each progress update"""
        
        # Track phases based on message content
        if "PHASE 1:" in message:
            self.current_phase = 1
        elif "PHASE 2:" in message:
            self.current_phase = 2
        elif "PHASE 3:" in message:
            self.current_phase = 3
        elif "PHASE 4:" in message:
            self.current_phase = 4
        
        # Update progress bar
        progress = (self.current_phase / self.total_phases) * 100
        self.progress_bar.progress(int(progress))
        
        # Update status
        if "PHASE" in message and ":" in message:
            self.status_text.info(f"🔄 {message}")
        elif "✅" in message:
            self.status_text.success(f"{message}")
        else:
            self.status_text.info(f"🔄 {message}")
        
        # Add to log (skip the "=" separator lines)
        if message and message.strip() and not all(c == '=' for c in message.strip()):
            with self.log_container:
                st.text(message)
    
    def complete(self):
        self.progress_bar.progress(100)
        self.status_text.success("✅ Validation completed!")

def main():
    st.set_page_config(
        page_title="StartupScan",
        page_icon="🚀",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Custom CSS for better styling
    st.markdown("""
    <style>
    .main-header {
        font-size: 3rem;
        font-weight: bold;
        text-align: center;
        margin-bottom: 2rem;
        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
    }
    .stButton > button {
        background: linear-gradient(90deg, #667eea 0%, #764ba2 100%);
        color: white;
        border: none;
        padding: 0.5rem 2rem;
        border-radius: 10px;
        font-weight: bold;
    }
    .result-section {
        background-color: #f8f9fa;
        padding: 2rem;
        border-radius: 15px;
        border-left: 5px solid #667eea;
        margin: 2rem 0;
    }
    .log-container {
        background-color: #f1f3f4;
        padding: 1rem;
        border-radius: 8px;
        font-family: monospace;
        font-size: 0.9rem;
        max-height: 300px;
        overflow-y: auto;
        border: 1px solid #ddd;
    }
    .header-buttons {
        display: flex;
        justify-content: center;
        gap: 15px;
        margin: 1.5rem 0;
        flex-wrap: wrap;
    }
    .header-btn {
        display: inline-block;
        padding: 12px 24px;
        background: linear-gradient(135deg, #17a2b8 0%, #138496 100%);
        color: white !important;
        text-decoration: none !important;
        border-radius: 25px;
        font-weight: 600;
        font-size: 0.95rem;
        transition: all 0.3s ease;
        box-shadow: 0 4px 15px rgba(23, 162, 184, 0.3);
        border: none;
        cursor: pointer;
    }
    .header-btn:hover {
        transform: translateY(-2px);
        box-shadow: 0 6px 20px rgba(23, 162, 184, 0.4);
        color: white !important;
        text-decoration: none !important;
    }
    .header-btn:visited {
        color: white !important;
        text-decoration: none !important;
    }
    .header-btn:link {
        color: white !important;
        text-decoration: none !important;
    }
    .header-btn.secondary {
        background: linear-gradient(135deg, #fd7e14 0%, #e55b13 100%);
        box-shadow: 0 4px 15px rgba(253, 126, 20, 0.3);
        color: white !important;
    }
    .header-btn.secondary:hover {
        box-shadow: 0 6px 20px rgba(253, 126, 20, 0.4);
        color: white !important;
    }
    .header-btn.secondary:visited {
        color: white !important;
        text-decoration: none !important;
    }
    .header-btn.secondary:link {
        color: white !important;
        text-decoration: none !important;
    }
    .header-btn .emoji {
        margin-right: 8px;
    }
    </style>
    """, unsafe_allow_html=True)
    
    # Header
    st.markdown('<h1 class="main-header">🚀 StartupScan</h1>', unsafe_allow_html=True)

    # Subtitle with demo and guide links
    st.markdown("""
    <div class="header-buttons">
        <a href="https://demo.atla-ai.com/app/atla-demo-802mg?startDate=2025-08-05" target="_blank" class="header-btn">
            <span class="emoji">🔍</span>View your traces on Atla
        </a>
        <a href="https://www.agno.com/?utm_source=atla&utm_medium=partner-content&utm_campaign=partner-technical&utm_content=atla" target="_blank" class="header-btn secondary">
            <span class="emoji">🛠️</span>Built with Agno
        </a>
    </div>
    """, unsafe_allow_html=True)

    st.markdown("---")
    
    # Sidebar with information
    with st.sidebar:
        st.markdown("## 📋 About This Tool")
        st.markdown("""
        This agentic tool validates your startup idea through:
        
        🎯 **Idea Clarification Agent**
        - Originality assessment
        - Mission definition
        - Objective setting
        
        📊 **Market Research Agent**
        - TAM/SAM/SOM analysis
        - Customer segmentation
        - Market trends
        
        🏢 **Competitor Analysis Agent**
        - SWOT analysis
        - Positioning assessment
        - Market gaps
        
        📋 **Validation Report Agent**
        - Executive summary
        - Strategic recommendations
        - Next steps
        """)
        
        st.markdown("---")
        st.markdown("💡 **Tip:** Be as specific as possible about your startup idea for better results!")
    
    # Main content area
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.markdown("## 💡 Tell us about your startup idea")
        
        # Input form
        with st.form("startup_form", clear_on_submit=False):
            # Get default value from session state if an example was selected
            default_value = st.session_state.get('selected_example', '')
            
            startup_idea = st.text_area(
                "Describe your startup idea in detail:",
                value=default_value,
                height=150,
                placeholder="e.g., A marketplace for Christmas Ornaments made from leather that connects artisans with customers looking for unique holiday decorations...",
                help="The more detailed your description, the better the validation will be."
            )
            
            # Model selection toggle
            model_id = st.selectbox(
                "🤖 Model to use:",
                options=["gpt-4o", "gpt-4o-mini", "o1", "o3-mini"],
                index=0,  # Default to gpt-4o
                help="Choose which AI model to use for the validation analysis"
            )
            
            submitted = st.form_submit_button("🔍 Validate My Idea", use_container_width=True)
    
    with col2:
        st.markdown("## 🎯 Quick Examples")
        st.markdown("*Click any example to populate the idea field*")
        
        examples = [
            "AI-powered personal finance coach for millennials",
            "Sustainable packaging solutions for e-commerce",
            "Virtual reality fitness platform for remote workers",
            "Marketplace for local artisan food products",
            "Smart home energy optimization system"
        ]
        
        for i, example in enumerate(examples):
            if st.button(f"💡 {example}", key=f"example_{i}", help="Click to populate the startup idea field"):
                st.session_state.selected_example = example
                st.rerun()
    
    # Clear the selected example from session state after form renders
    if 'selected_example' in st.session_state and startup_idea:
        if startup_idea == st.session_state.selected_example:
            # Example has been loaded into the form, show success message
            st.success(f"✅ Example loaded: {startup_idea[:50]}...")
        # Keep the example in session state until form is submitted
    
    # Process the form submission
    if submitted and startup_idea and model_id:
        # Clear the selected example when submitting
        if 'selected_example' in st.session_state:
            del st.session_state.selected_example
            
        st.markdown("---")
        st.markdown("## 🔄 Real-time Validation Progress")
        
        # Initialize progress tracker
        tracker = StreamlitProgressTracker()
        tracker.setup_ui()
        
        try:
            # Prepare the message
            message = "Please validate this startup idea with comprehensive market research and competitive analysis"
            
            # Run the workflow with real-time progress
            async def run_validation():
                return await startup_validation_workflow.arun(
                    message=message,
                    startup_idea=startup_idea,
                    model_id=model_id,  # Pass the selected model
                    progress_callback=tracker.update_progress  # Pass the real-time callback
                )
            
            # Execute the async workflow
            result = asyncio.run(run_validation())
            
            # Complete progress
            tracker.complete()
            
            # Display results with improved formatting
            st.markdown("---")
            st.markdown("## 📊 Validation Results")
            
            # Extract clean content from WorkflowRunResponse
            validation_content = ""
            
            if hasattr(result, 'content') and result.content:
                validation_content = result.content
            elif hasattr(result, 'response') and result.response:
                validation_content = result.response
            else:
                validation_content = str(result)
            
            # Clean up content if needed
            if validation_content:
                validation_content = validation_content.replace('\\n', '\n')
                validation_content = validation_content.replace('\\"', '"')
            
            # Display in a formatted code block for consistent appearance
            st.code(validation_content, language="markdown")
            
            # Add download button for the report
            st.download_button(
                label="📥 Download Report",
                data=validation_content,
                file_name=f"startup_validation_{startup_idea[:30].replace(' ', '_')}.md",
                mime="text/markdown"
            )
            
        except Exception as e:
            st.error(f"❌ An error occurred during validation: {str(e)}")
            st.error("Please check your environment variables and try again.")
            
            # Display error details in expander
            with st.expander("🔍 Error Details"):
                st.code(str(e))
    
    elif submitted and not startup_idea:
        st.warning("⚠️ Please enter a startup idea before submitting.")
    
    # Footer
    st.markdown("---")
    st.markdown("""
    <div style='text-align: center; color: #666; font-size: 0.9rem;'>
        <p>🔬 Powered by AI agents and comprehensive market research</p>
        <p>⚠️ This validation is for informational purposes only. Conduct additional due diligence before making investment decisions.</p>
    </div>
    """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()