"""Streamlit Research Agent using Perplexity sonar-reasoning-pro via OpenRouter.

A focused research application that leverages Perplexity's advanced reasoning model
for comprehensive research with real-time web search capabilities.
"""

import asyncio
import streamlit as st
from datetime import datetime, timezone
from typing import Optional
from dataclasses import asdict
import json

# Import project configuration and services
from src.config import config
from src.services.protocols import ResearchQuery, ResearchResult
from src.services.research.service import OpenRouterResearchService
from src.services.research.config import ResearchConfig


class PerplexityResearchAgent:
    """Focused research agent using Perplexity sonar-reasoning-pro."""
    
    def __init__(self):
        """Initialize the research agent."""
        self.setup_page_config()
        self.initialize_session_state()
        self.setup_research_service()
    
    def setup_page_config(self):
        """Configure Streamlit page settings."""
        st.set_page_config(
            page_title="Perplexity Research Agent",
            page_icon="🧠",
            layout="wide",
            initial_sidebar_state="expanded"
        )
    
    def initialize_session_state(self):
        """Initialize Streamlit session state."""
        if "research_history" not in st.session_state:
            st.session_state.research_history = []
        if "current_result" not in st.session_state:
            st.session_state.current_result = None
    
    def setup_research_service(self):
        """Setup the research service with OpenRouter."""
        if not config.OPENROUTER_API_KEY:
            st.error("❌ OPENROUTER_API_KEY not found in environment")
            st.stop()
        
        try:
            research_config = ResearchConfig.from_env(config.OPENROUTER_API_KEY)
            self.research_service = OpenRouterResearchService(research_config)
        except Exception as e:
            st.error(f"❌ Failed to initialize research service: {e}")
            st.stop()
    
    def render_header(self):
        """Render the application header."""
        st.title("🧠 Perplexity Research Agent")
        st.markdown("""
        **Powered by Perplexity sonar-reasoning-pro via OpenRouter**
        
        Advanced reasoning and research capabilities with real-time web search.
        """)
    
    def render_sidebar(self):
        """Render the sidebar with input controls."""
        with st.sidebar:
            st.header("🔍 Research Query")
            
            # Main query input
            query = st.text_area(
                "Enter your research question:",
                placeholder="What are the latest developments in AI reasoning models?",
                height=100
            )
            
            # Advanced options
            with st.expander("⚙️ Advanced Settings"):
                max_tokens = st.slider("Max Tokens", 1000, 4000, 4000, step=500)
                temperature = st.slider("Temperature", 0.0, 1.0, 0.1, step=0.1)
                
                context = st.text_area(
                    "Additional Context (optional):",
                    placeholder="Focus on specific aspects...",
                    height=80
                )
            
            # Research button
            start_research = st.button(
                "🚀 Start Research",
                type="primary",
                disabled=not query.strip(),
                use_container_width=True
            )
            
            # Example queries
            st.divider()
            st.subheader("💡 Example Queries")
            examples = [
                "What are the latest developments in AI reasoning models like o1 and o3?",
                "How do vector databases compare for RAG applications in 2025?",
                "What are the best practices for fine-tuning large language models?",
                "How is WebAssembly being used for edge computing in 2025?"
            ]
            
            for example in examples:
                if st.button(example, key=f"example_{hash(example)}", use_container_width=True):
                    st.session_state.example_query = example
                    st.rerun()
            
            return query, context, max_tokens, temperature, start_research
    
    async def conduct_research(self, query: str, context: str, max_tokens: int, temperature: float) -> ResearchResult:
        """Conduct research using Perplexity sonar-reasoning-pro."""
        research_query = ResearchQuery(
            query=query,
            context=context if context.strip() else None,
            max_tokens=max_tokens,
            temperature=temperature,
            model="perplexity/sonar-reasoning-pro"
        )
        
        return await self.research_service.research(research_query)
    
    def render_research_process(self, query: str, context: str, max_tokens: int, temperature: float):
        """Render the research process with progress indicators."""
        with st.spinner("🧠 Conducting research with Perplexity sonar-reasoning-pro..."):
            try:
                # Execute research
                result = asyncio.run(self.conduct_research(query, context, max_tokens, temperature))
                
                # Store results
                st.session_state.current_result = result
                st.session_state.research_history.append(result)
                
                st.success("✅ Research completed successfully!")
                return result
                
            except Exception as e:
                st.error(f"❌ Research failed: {str(e)}")
                return None
    
    def render_results(self, result: ResearchResult):
        """Render research results."""
        if not result:
            return
        
        # Main results display
        st.subheader("📊 Research Results")
        
        # Metrics row
        col1, col2, col3, col4 = st.columns(4)
        with col1:
            st.metric("Processing Time", f"{result.processing_time:.1f}s")
        with col2:
            st.metric("Confidence", f"{result.confidence_score:.1%}")
        with col3:
            st.metric("Sources", len(result.sources))
        with col4:
            st.metric("Tokens", result.token_usage.get("total_tokens", 0))
        
        # Answer display
        st.subheader("💡 Answer")
        st.markdown(result.answer)
        
        # Sources display
        if result.sources:
            st.subheader("🔗 Sources")
            for i, source in enumerate(result.sources, 1):
                st.markdown(f"{i}. {source}")
        
        # Download options
        st.subheader("📥 Export Options")
        col1, col2 = st.columns(2)
        
        with col1:
            # Markdown export
            markdown_content = self.format_markdown_report(result)
            st.download_button(
                "Download Markdown Report",
                markdown_content,
                file_name=f"research_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md",
                mime="text/markdown",
                use_container_width=True
            )
        
        with col2:
            # JSON export
            json_data = asdict(result)
            json_data['timestamp'] = datetime.now(timezone.utc).isoformat()
            st.download_button(
                "Download JSON Data",
                json.dumps(json_data, indent=2),
                file_name=f"research_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                mime="application/json",
                use_container_width=True
            )
    
    def format_markdown_report(self, result: ResearchResult) -> str:
        """Format research result as markdown report."""
        return f"""# Research Report: {result.query}

## Executive Summary
{result.answer[:300]}...

## Detailed Analysis
{result.answer}

## Sources
{chr(10).join(f"- {source}" for source in result.sources) if result.sources else "- Sources integrated in analysis"}

---
**Research Metadata:**
- Model: {result.model_used}
- Processing Time: {result.processing_time:.2f} seconds
- Confidence Score: {result.confidence_score:.1%}
- Tokens Used: {result.token_usage.get('total_tokens', 'N/A')}
- Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
"""
    
    def render_history(self):
        """Render research history."""
        if not st.session_state.research_history:
            return
        
        st.subheader("📚 Research History")
        
        for i, result in enumerate(reversed(st.session_state.research_history[-5:]), 1):
            with st.expander(f"Research #{len(st.session_state.research_history) - i + 1}: {result.query[:50]}..."):
                st.markdown(f"**Query:** {result.query}")
                st.markdown(f"**Answer:** {result.answer[:200]}...")
                st.markdown(f"**Confidence:** {result.confidence_score:.1%}")
                st.markdown(f"**Processing Time:** {result.processing_time:.1f}s")
                
                if st.button(f"View Full Results #{i}", key=f"view_{i}"):
                    st.session_state.current_result = result
                    st.rerun()
    
    def run(self):
        """Run the research agent application."""
        self.render_header()
        
        # Get user input
        query, context, max_tokens, temperature, start_research = self.render_sidebar()
        
        # Handle example query selection
        if hasattr(st.session_state, 'example_query'):
            query = st.session_state.example_query
            del st.session_state.example_query
            start_research = True
        
        # Main content area
        if start_research and query.strip():
            result = self.render_research_process(query, context, max_tokens, temperature)
            if result:
                self.render_results(result)
        
        # Show current result if available
        elif st.session_state.current_result:
            self.render_results(st.session_state.current_result)
        
        # Show history
        self.render_history()


def main():
    """Main entry point for the research agent."""
    agent = PerplexityResearchAgent()
    agent.run()


if __name__ == "__main__":
    main()
