import gradio as gr
import asyncio
from groq import Groq
import json
import time
from typing import Dict, List, Tuple, Optional
import threading
from datetime import datetime

class ReasoningOrchestra:
    def __init__(self):
        self.client = None
        self.is_api_key_set = False
        
    def set_api_key(self, api_key: str) -> str:
        """Set the Groq API key and test connection"""
        if not api_key.strip():
            return "❌ Please enter a valid API key"
        
        try:
            self.client = Groq(api_key=api_key.strip())
            # Test the connection with a simple request
            test_completion = self.client.chat.completions.create(
                model="qwen/qwen3-32b",
                messages=[{"role": "user", "content": "Hello"}],
                max_completion_tokens=10
            )
            self.is_api_key_set = True
            return "✅ API key validated successfully! You can now use the Reasoning Orchestra."
        except Exception as e:
            self.is_api_key_set = False
            return f"❌ API key validation failed: {str(e)}"
    
    def deep_thinker_analyze(self, problem: str, context: str = "") -> Dict:
        """DeepSeek R1 - The Deep Thinker"""
        if not self.is_api_key_set:
            return {"error": "API key not set"}
        
        prompt = f"""You are the Deep Thinker in a collaborative reasoning system. Your role is to provide thorough, methodical analysis with extensive step-by-step reasoning.

Problem: {problem}
{f"Additional Context: {context}" if context else ""}

Please provide a comprehensive analysis with deep reasoning. Think through all implications, consider multiple angles, and provide detailed step-by-step logic."""

        try:
            completion = self.client.chat.completions.create(
                model="deepseek-r1-distill-llama-70b",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.6,
                max_completion_tokens=2048,
                top_p=0.95,
                reasoning_format="raw"
            )
            
            response_content = completion.choices[0].message.content
            
            return {
                "model": "DeepSeek R1 (Deep Thinker)",
                "role": "🎭 The Philosopher & Deep Analyzer",
                "reasoning": response_content,
                "timestamp": datetime.now().strftime("%H:%M:%S"),
                "tokens_used": completion.usage.total_tokens if hasattr(completion, 'usage') else "N/A"
            }
        except Exception as e:
            return {"error": f"Deep Thinker error: {str(e)}"}
    
    def quick_strategist_analyze(self, problem: str, context: str = "") -> Dict:
        """Qwen3 32B - The Quick Strategist"""
        if not self.is_api_key_set:
            return {"error": "API key not set"}
        
        prompt = f"""You are the Quick Strategist in a collaborative reasoning system. Your role is to provide fast, efficient strategic analysis with clear action plans.

Problem: {problem}
{f"Additional Context: {context}" if context else ""}

Please provide a strategic analysis with:
1. Key insights and patterns
2. Practical solutions
3. Implementation priorities
4. Risk assessment
5. Clear next steps

Be decisive and solution-focused."""

        try:
            completion = self.client.chat.completions.create(
                model="qwen/qwen3-32b",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.6,
                top_p=0.95,
                max_completion_tokens=1536,
                reasoning_effort="default"
            )
            
            return {
                "model": "Qwen3 32B (Quick Strategist)",
                "role": "🚀 The Strategic Decision Maker",
                "reasoning": completion.choices[0].message.content,
                "timestamp": datetime.now().strftime("%H:%M:%S"),
                "tokens_used": completion.usage.total_tokens if hasattr(completion, 'usage') else "N/A"
            }
        except Exception as e:
            return {"error": f"Quick Strategist error: {str(e)}"}
    
    def detail_detective_analyze(self, problem: str, context: str = "") -> Dict:
        """QwQ 32B - The Detail Detective"""
        if not self.is_api_key_set:
            return {"error": "API key not set"}
        
        prompt = f"""You are the Detail Detective in a collaborative reasoning system. Your role is to provide meticulous investigation and comprehensive fact-checking.

Problem: {problem}
{f"Additional Context: {context}" if context else ""}

Please conduct a thorough investigation including:
1. Detailed analysis of all aspects
2. Potential edge cases and considerations
3. Verification of assumptions
4. Historical context or precedents
5. Comprehensive pros and cons
6. Hidden connections or implications

Be extremely thorough and leave no stone unturned."""

        try:
            completion = self.client.chat.completions.create(
                model="qwen-qwq-32b",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.6,
                top_p=0.95,
                max_completion_tokens=2048,
                reasoning_format="parsed"
            )
            
            return {
                "model": "QwQ 32B (Detail Detective)",
                "role": "🔍 The Meticulous Investigator",
                "reasoning": completion.choices[0].message.content,
                "timestamp": datetime.now().strftime("%H:%M:%S"),
                "tokens_used": completion.usage.total_tokens if hasattr(completion, 'usage') else "N/A"
            }
        except Exception as e:
            return {"error": f"Detail Detective error: {str(e)}"}
    
    def synthesize_orchestra(self, deep_result: Dict, strategic_result: Dict, detective_result: Dict, original_problem: str) -> str:
        """Synthesize all three perspectives into a final orchestrated solution"""
        if not self.is_api_key_set:
            return "API key not set"
        
        synthesis_prompt = f"""You are the Orchestra Conductor. You have received three different analytical perspectives on the same problem. Your job is to synthesize these into a comprehensive, unified solution.

ORIGINAL PROBLEM: {original_problem}

DEEP THINKER ANALYSIS:
{deep_result.get('reasoning', 'No analysis available')}

STRATEGIC ANALYSIS:
{strategic_result.get('reasoning', 'No analysis available')}

DETECTIVE INVESTIGATION:
{detective_result.get('reasoning', 'No analysis available')}

Please create a unified synthesis that:
1. Combines the best insights from all three perspectives
2. Resolves any contradictions between the analyses
3. Provides a comprehensive final recommendation
4. Highlights where the different reasoning styles complement each other
5. Gives a clear, actionable conclusion

Format your response as a well-structured final solution that leverages all three reasoning approaches."""

        try:
            completion = self.client.chat.completions.create(
                model="qwen/qwen3-32b",
                messages=[{"role": "user", "content": synthesis_prompt}],
                temperature=0.7,
                max_completion_tokens=2048,
                top_p=0.9
            )
            
            return completion.choices[0].message.content
        except Exception as e:
            return f"Synthesis error: {str(e)}"

# Initialize the orchestra
orchestra = ReasoningOrchestra()

def validate_api_key(api_key: str) -> str:
    """Validate the API key and return status"""
    return orchestra.set_api_key(api_key)

def run_single_model(problem: str, model_choice: str, context: str = "") -> Tuple[str, str]:
    """Run a single model analysis"""
    if not orchestra.is_api_key_set:
        return "❌ Please set your API key first", ""
    
    if not problem.strip():
        return "❌ Please enter a problem to analyze", ""
    
    start_time = time.time()
    
    if model_choice == "Deep Thinker (DeepSeek R1)":
        result = orchestra.deep_thinker_analyze(problem, context)
    elif model_choice == "Quick Strategist (Qwen3 32B)":
        result = orchestra.quick_strategist_analyze(problem, context)
    elif model_choice == "Detail Detective (QwQ 32B)":
        result = orchestra.detail_detective_analyze(problem, context)
    else:
        return "❌ Invalid model selection", ""
    
    elapsed_time = time.time() - start_time
    
    if "error" in result:
        return f"❌ {result['error']}", ""
    
    formatted_output = f"""## {result['role']} - {result['model']}
**Analysis Time:** {elapsed_time:.2f} seconds | **Timestamp:** {result['timestamp']} | **Tokens:** {result['tokens_used']}

---

{result['reasoning']}
"""
    
    return formatted_output, ""

def run_full_orchestra(problem: str, context: str = "") -> Tuple[str, str, str, str]:
    """Run the full collaborative reasoning orchestra"""
    if not orchestra.is_api_key_set:
        error_msg = "❌ Please set your API key first"
        return error_msg, error_msg, error_msg, error_msg
    
    if not problem.strip():
        error_msg = "❌ Please enter a problem to analyze"
        return error_msg, error_msg, error_msg, error_msg
    
    status_updates = []
    
    # Phase 1: Deep Thinker
    status_updates.append("🎭 Deep Thinker is analyzing the problem...")
    deep_result = orchestra.deep_thinker_analyze(problem, context)
    
    # Phase 2: Quick Strategist  
    status_updates.append("🚀 Quick Strategist is developing strategies...")
    strategic_result = orchestra.quick_strategist_analyze(problem, context)
    
    # Phase 3: Detail Detective
    status_updates.append("🔍 Detail Detective is investigating thoroughly...")
    detective_result = orchestra.detail_detective_analyze(problem, context)
    
    # Phase 4: Synthesis
    status_updates.append("🎼 Orchestra Conductor is synthesizing all perspectives...")
    synthesis = orchestra.synthesize_orchestra(deep_result, strategic_result, detective_result, problem)
    
    # Format outputs
    def format_result(result: Dict) -> str:
        if "error" in result:
            return f"❌ {result['error']}"
        
        return f"""## {result['role']} - {result['model']}
**Timestamp:** {result['timestamp']} | **Tokens:** {result['tokens_used']}

---

{result['reasoning']}
"""
    
    deep_output = format_result(deep_result)
    strategic_output = format_result(strategic_result)
    detective_output = format_result(detective_result)
    
    synthesis_output = f"""## 🎼 Orchestra Conductor - Final Synthesis

---

{synthesis}
"""
    
    return deep_output, strategic_output, detective_output, synthesis_output

# Custom CSS for better styling
custom_css = """
.gradio-container {
    max-width: 1200px !important;
}
.api-key-section {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    padding: 20px;
    border-radius: 10px;
    margin-bottom: 20px;
}
.model-section {
    border: 2px solid #e1e5e9;
    border-radius: 10px;
    padding: 15px;
    margin: 10px 0;
}
.orchestra-header {
    text-align: center;
    background: linear-gradient(45deg, #f093fb 0%, #f5576c 100%);
    padding: 20px;
    border-radius: 15px;
    margin-bottom: 20px;
}
.status-box {
    background-color: #f8f9fa;
    border-left: 4px solid #28a745;
    padding: 15px;
    margin: 10px 0;
    border-radius: 5px;
}
"""

# Build the Gradio interface
with gr.Blocks(css=custom_css, title="Reasoning Orchestra") as app:
    # Header
    gr.HTML("""
    <div class="orchestra-header">
        <h1>🎼 The Collaborative Reasoning Orchestra</h1>
        <p><em>Where AI models collaborate like musicians in an orchestra to solve complex problems</em></p>
    </div>
    """)
    
    # API Key Section
    with gr.Group():
        gr.HTML('<div class="api-key-section"><h3 style="color: white; margin-top: 0;">🔑 API Configuration</h3></div>')
        with gr.Row():
            api_key_input = gr.Textbox(
                label="Enter your Groq API Key",
                type="password",
                placeholder="gsk_...",
                info="Get your free API key from https://console.groq.com/keys"
            )
            api_status = gr.Textbox(
                label="API Status",
                interactive=False,
                placeholder="Enter API key to validate..."
            )
        
        validate_btn = gr.Button("🔍 Validate API Key", variant="primary")
        validate_btn.click(
            fn=validate_api_key,
            inputs=[api_key_input],
            outputs=[api_status]
        )
    
    # Main Interface Tabs
    with gr.Tabs() as tabs:
        
        # Single Model Tab
        with gr.TabItem("🎯 Single Model Analysis"):
            gr.Markdown("### Test individual reasoning models")
            
            with gr.Row():
                with gr.Column():
                    single_problem = gr.Textbox(
                        label="Problem Statement",
                        placeholder="Enter the problem you want to analyze...",
                        lines=4
                    )
                    single_context = gr.Textbox(
                        label="Additional Context (Optional)",
                        placeholder="Any additional context or constraints...",
                        lines=2
                    )
                    model_choice = gr.Dropdown(
                        label="Choose Model",
                        choices=[
                            "Deep Thinker (DeepSeek R1)",
                            "Quick Strategist (Qwen3 32B)", 
                            "Detail Detective (QwQ 32B)"
                        ],
                        value="Deep Thinker (DeepSeek R1)"
                    )
                    single_analyze_btn = gr.Button("🚀 Analyze", variant="primary")
                
                with gr.Column():
                    single_output = gr.Markdown(label="Analysis Result")
            
            single_analyze_btn.click(
                fn=run_single_model,
                inputs=[single_problem, model_choice, single_context],
                outputs=[single_output, gr.Textbox(visible=False)]
            )
        
        # Full Orchestra Tab  
        with gr.TabItem("🎼 Full Orchestra Collaboration"):
            gr.Markdown("### Run all three models collaboratively for comprehensive analysis")
            
            with gr.Row():
                with gr.Column(scale=1):
                    orchestra_problem = gr.Textbox(
                        label="Problem Statement",
                        placeholder="Enter a complex problem that benefits from multiple reasoning perspectives...",
                        lines=6
                    )
                    orchestra_context = gr.Textbox(
                        label="Additional Context (Optional)",
                        placeholder="Background information, constraints, or specific requirements...",
                        lines=3
                    )
                    orchestra_analyze_btn = gr.Button("🎼 Start Orchestra Analysis", variant="primary", size="lg")
                
                with gr.Column(scale=2):
                    gr.Markdown("### 🎭 Deep Thinker Analysis")
                    deep_output = gr.Markdown()
                    
                    gr.Markdown("### 🚀 Quick Strategist Analysis") 
                    strategic_output = gr.Markdown()
                    
                    gr.Markdown("### 🔍 Detail Detective Analysis")
                    detective_output = gr.Markdown()
                    
                    gr.Markdown("### 🎼 Final Orchestrated Solution")
                    synthesis_output = gr.Markdown()
            
            orchestra_analyze_btn.click(
                fn=run_full_orchestra,
                inputs=[orchestra_problem, orchestra_context],
                outputs=[deep_output, strategic_output, detective_output, synthesis_output]
            )
        
        # Examples Tab
        with gr.TabItem("💡 Example Problems"):
            gr.Markdown("""
            ### Try these example problems to see the Orchestra in action:
            
            **🏢 Business Strategy:**
            "Our tech startup has limited funding and needs to decide between focusing on product development or marketing. We have a working MVP but low user adoption."
            
            **🤖 Ethical AI:**  
            "Should autonomous vehicles prioritize passenger safety over pedestrian safety in unavoidable accident scenarios? Consider the ethical, legal, and practical implications."
            
            **🌍 Environmental Policy:**
            "Design a policy framework to reduce carbon emissions in urban areas while maintaining economic growth and social equity."
            
            **🧬 Scientific Research:**
            "We've discovered a potential breakthrough in gene therapy, but it requires human trials. How should we proceed given the risks, benefits, and regulatory requirements?"
            
            **🎓 Educational Innovation:**
            "How can we redesign traditional university education to better prepare students for the rapidly changing job market of the 2030s?"
            
            **🏠 Urban Planning:**
            "A city wants to build affordable housing but faces opposition from current residents, environmental concerns, and budget constraints. Develop a comprehensive solution."
            """)
    
    # Footer
    gr.HTML("""
    <div style="text-align: center; margin-top: 30px; padding: 20px; background-color: #f8f9fa; border-radius: 10px;">
        <h4>🎼 How the Orchestra Works</h4>
        <p><strong>Deep Thinker (DeepSeek R1):</strong> Provides thorough philosophical and theoretical analysis</p>
        <p><strong>Quick Strategist (Qwen3 32B):</strong> Delivers practical strategies and action plans</p>
        <p><strong>Detail Detective (QwQ 32B):</strong> Conducts comprehensive investigation and fact-checking</p>
        <p><strong>Orchestra Conductor:</strong> Synthesizes all perspectives into a unified solution</p>
        <br>
        <p><em>Built with ❤️ using Groq's lightning-fast inference and Gradio</em></p>
    </div>
    """)

# Launch the app
if __name__ == "__main__":
    app.launch(
        share=True
    )