import os
import json
import time
import gradio as gr
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any, Optional, Union

# Import Groq - we'll install it in requirements.txt
from groq import Groq

class PersonalAIResearchAssistant:
    """
    Personal AI Research Assistant (PARA) using Groq's compound models with agentic capabilities.
    """
    
    def __init__(self, api_key: str, 
                 knowledge_base_path: str = "knowledge_base.json",
                 model: str = "compound-beta"):
        """
        Initialize the PARA system.
        
        Args:
            api_key: Groq API key
            knowledge_base_path: Path to store persistent knowledge
            model: Which Groq model to use ('compound-beta' or 'compound-beta-mini')
        """
        self.api_key = api_key
        if not self.api_key:
            raise ValueError("No API key provided")
        
        self.client = Groq(api_key=self.api_key)
        self.model = model
        self.knowledge_base_path = Path(knowledge_base_path)
        self.knowledge_base = self._load_knowledge_base()
        
    def _load_knowledge_base(self) -> Dict:
        """Load existing knowledge base or create a new one"""
        if self.knowledge_base_path.exists():
            with open(self.knowledge_base_path, 'r') as f:
                return json.load(f)
        else:
            # Initialize with empty collections
            kb = {
                "topics": {},
                "research_digests": [],
                "code_analyses": [],
                "concept_connections": [],
                "metadata": {
                    "created_at": datetime.now().isoformat(),
                    "last_updated": datetime.now().isoformat()
                }
            }
            self._save_knowledge_base(kb)
            return kb
    
    def _save_knowledge_base(self, kb: Dict = None) -> None:
        """Save the knowledge base to disk"""
        if kb is None:
            kb = self.knowledge_base
            
        # Update metadata
        kb["metadata"]["last_updated"] = datetime.now().isoformat()
        
        with open(self.knowledge_base_path, 'w') as f:
            json.dump(kb, f, indent=2)
    
    def _extract_tool_info(self, response) -> Dict:
        """
        Extract tool usage information in a JSON serializable format
        """
        tool_info = None
        if hasattr(response.choices[0].message, 'executed_tools'):
            # Convert ExecutedTool objects to dictionaries
            tools = response.choices[0].message.executed_tools
            if tools:
                tool_info = []
                for tool in tools:
                    # Extract only serializable data
                    tool_dict = {
                        "tool_type": getattr(tool, "type", "unknown"),
                        "tool_name": getattr(tool, "name", "unknown"),
                    }
                    # Add any other relevant attributes in a serializable form
                    if hasattr(tool, "input"):
                        tool_dict["input"] = str(tool.input)
                    if hasattr(tool, "output"):
                        tool_dict["output"] = str(tool.output)
                    tool_info.append(tool_dict)
        return tool_info
    
    def research_digest(self, topic: str, 
                        include_domains: List[str] = None,
                        exclude_domains: List[str] = None,
                        max_results: int = 5) -> Dict:
        """
        Generate a research digest on a specific topic
        
        Args:
            topic: The research topic to investigate
            include_domains: List of domains to include (e.g., ["arxiv.org", "*.edu"])
            exclude_domains: List of domains to exclude
            max_results: Maximum number of key findings to include
            
        Returns:
            Research digest including key findings and references
        """
        # Build the prompt
        prompt = f"""Generate a research digest on the topic: {topic}
        
        Please find the most recent and relevant information, focusing on:
        1. Key findings or breakthroughs
        2. Current trends and methodologies
        3. Influential researchers or organizations
        4. Practical applications
        
        Structure your response as a concise summary with {max_results} key points maximum.
        Include source information where possible.
        """
        
        # Set up API parameters
        params = {
            "messages": [
                {"role": "system", "content": "You are a research assistant tasked with finding and summarizing the latest information on specific topics."},
                {"role": "user", "content": prompt}
            ],
            "model": self.model
        }
        
        # Add domain filtering if specified
        if include_domains and include_domains[0].strip():
            params["include_domains"] = [domain.strip() for domain in include_domains]
        if exclude_domains and exclude_domains[0].strip():
            params["exclude_domains"] = [domain.strip() for domain in exclude_domains]
            
        # Make the API call
        response = self.client.chat.completions.create(**params)
        content = response.choices[0].message.content
        
        # Extract tool usage information in a serializable format
        tool_info = self._extract_tool_info(response)
        
        # Create digest entry
        digest = {
            "topic": topic,
            "timestamp": datetime.now().isoformat(),
            "content": content,
            "tool_usage": tool_info,
            "parameters": {
                "include_domains": include_domains,
                "exclude_domains": exclude_domains,
            }
        }
        
        # Add to knowledge base
        self.knowledge_base["research_digests"].append(digest)
        
        # Update topic entry in knowledge base
        if topic not in self.knowledge_base["topics"]:
            self.knowledge_base["topics"][topic] = {
                "first_researched": datetime.now().isoformat(),
                "research_count": 1,
                "related_topics": []
            }
        else:
            self.knowledge_base["topics"][topic]["research_count"] += 1
            self.knowledge_base["topics"][topic]["last_researched"] = datetime.now().isoformat()
        
        # Save updated knowledge base
        self._save_knowledge_base()
        
        return digest
    
    def evaluate_code(self, code_snippet: str, language: str = "python", 
                     analysis_type: str = "full") -> Dict:
        """
        Evaluate a code snippet for issues and suggest improvements
        
        Args:
            code_snippet: The code to evaluate
            language: Programming language of the code
            analysis_type: Type of analysis to perform ('full', 'security', 'performance', 'style')
            
        Returns:
            Analysis results including issues and suggestions
        """
        # Build the prompt
        prompt = f"""Analyze the following {language} code:
        
        ```{language}
        {code_snippet}
        ```
        
        Please perform a {analysis_type} analysis, including:
        1. Identifying any bugs or potential issues
        2. Security vulnerabilities (if applicable)
        3. Performance considerations
        4. Style and best practices
        5. Suggested improvements
        
        If possible, execute the code to verify functionality.
        """
        
        # Make the API call
        response = self.client.chat.completions.create(
            messages=[
                {"role": "system", "content": f"You are a code analysis expert specializing in {language}."},
                {"role": "user", "content": prompt}
            ],
            model=self.model
        )
        
        content = response.choices[0].message.content
        
        # Extract tool usage information in a serializable format
        tool_info = self._extract_tool_info(response)
        
        # Create code analysis entry
        analysis = {
            "code_snippet": code_snippet,
            "language": language,
            "analysis_type": analysis_type,
            "timestamp": datetime.now().isoformat(),
            "content": content,
            "tool_usage": tool_info
        }
        
        # Add to knowledge base
        self.knowledge_base["code_analyses"].append(analysis)
        self._save_knowledge_base()
        
        return analysis
    
    def connect_concepts(self, concept_a: str, concept_b: str) -> Dict:
        """
        Identify connections between two seemingly different concepts
        
        Args:
            concept_a: First concept
            concept_b: Second concept
            
        Returns:
            Analysis of connections between the concepts
        """
        # Build the prompt
        prompt = f"""Explore the connections between these two concepts:
        
        Concept A: {concept_a}
        Concept B: {concept_b}
        
        Please identify:
        1. Direct connections or shared principles
        2. Historical influences between them
        3. Common applications or use cases
        4. How insights from one field might benefit the other
        5. Potential for innovative combinations
        
        Search for the most up-to-date information that might connect these concepts.
        """
        
        # Make the API call
        response = self.client.chat.completions.create(
            messages=[
                {"role": "system", "content": "You are a cross-disciplinary research assistant specialized in finding connections between different fields and concepts."},
                {"role": "user", "content": prompt}
            ],
            model=self.model
        )
        
        content = response.choices[0].message.content
        
        # Extract tool usage information in a serializable format
        tool_info = self._extract_tool_info(response)
        
        # Create connection entry
        connection = {
            "concept_a": concept_a,
            "concept_b": concept_b,
            "timestamp": datetime.now().isoformat(),
            "content": content,
            "tool_usage": tool_info
        }
        
        # Add to knowledge base
        self.knowledge_base["concept_connections"].append(connection)
        
        # Update topic entries
        for concept in [concept_a, concept_b]:
            if concept not in self.knowledge_base["topics"]:
                self.knowledge_base["topics"][concept] = {
                    "first_researched": datetime.now().isoformat(),
                    "research_count": 1,
                    "related_topics": [concept_a if concept == concept_b else concept_b]
                }
            else:
                if concept_a if concept == concept_b else concept_b not in self.knowledge_base["topics"][concept]["related_topics"]:
                    self.knowledge_base["topics"][concept]["related_topics"].append(
                        concept_a if concept == concept_b else concept_b
                    )
        
        self._save_knowledge_base()
        
        return connection
    
    def ask_knowledge_base(self, query: str) -> Dict:
        """
        Query the accumulated knowledge base
        
        Args:
            query: Question about topics in the knowledge base
            
        Returns:
            Response based on accumulated knowledge
        """
        # Create a temporary context from the knowledge base
        context = {
            "topics_researched": list(self.knowledge_base["topics"].keys()),
            "research_count": len(self.knowledge_base["research_digests"]),
            "code_analyses_count": len(self.knowledge_base["code_analyses"]),
            "concept_connections_count": len(self.knowledge_base["concept_connections"]),
            "last_updated": self.knowledge_base["metadata"]["last_updated"]
        }
        
        # Add recent research digests (limited to keep context manageable)
        recent_digests = self.knowledge_base["research_digests"][-3:] if self.knowledge_base["research_digests"] else []
        context["recent_research"] = recent_digests
        
        # Build the prompt
        prompt = f"""Query: {query}
        
        Please answer based on the following knowledge base context:
        {json.dumps(context, indent=2)}
        
        If the knowledge base doesn't contain relevant information, please indicate this and suggest how we might research this topic.
        """
        
        # Make the API call
        response = self.client.chat.completions.create(
            messages=[
                {"role": "system", "content": "You are a research assistant with access to a personal knowledge base. Answer questions based on the accumulated knowledge."},
                {"role": "user", "content": prompt}
            ],
            model=self.model
        )
        
        content = response.choices[0].message.content
        
        return {
            "query": query,
            "timestamp": datetime.now().isoformat(),
            "response": content,
            "knowledge_base_state": context
        }
    
    def get_kb_stats(self):
        """Get statistics about the knowledge base"""
        return {
            "topics_count": len(self.knowledge_base["topics"]),
            "research_count": len(self.knowledge_base["research_digests"]),
            "code_analyses_count": len(self.knowledge_base["code_analyses"]),
            "concept_connections_count": len(self.knowledge_base["concept_connections"]),
            "created": self.knowledge_base["metadata"]["created_at"],
            "last_updated": self.knowledge_base["metadata"]["last_updated"],
            "topics": list(self.knowledge_base["topics"].keys())
        }

# Global variables for the Gradio app
para_instance = None
api_key_status = "Not Set"

# Helper functions for Gradio
def validate_api_key(api_key):
    """Validate Groq API key"""
    global para_instance, api_key_status
    
    if not api_key or len(api_key.strip()) < 10:
        return "❌ Please enter a valid API key"
    
    try:
        # Try to initialize with minimal actions
        client = Groq(api_key=api_key)
        # Create PARA instance
        para_instance = PersonalAIResearchAssistant(
            api_key=api_key,
            knowledge_base_path="para_knowledge.json"
        )
        api_key_status = "Valid ✅"
        
        # Get KB stats
        stats = para_instance.get_kb_stats()
        kb_info = f"**Knowledge Base Stats:**\n\n" \
                 f"- Topics: {stats['topics_count']}\n" \
                 f"- Research Digests: {stats['research_count']}\n" \
                 f"- Code Analyses: {stats['code_analyses_count']}\n" \
                 f"- Concept Connections: {stats['concept_connections_count']}\n" \
                 f"- Last Updated: {stats['last_updated'][:10]}\n\n" \
                 f"**Topics Explored:** {', '.join(stats['topics'][:10])}" + \
                 ("..." if len(stats['topics']) > 10 else "")
        
        return f"✅ API Key Valid! PARA is ready.\n\n{kb_info}"
    except Exception as e:
        api_key_status = "Invalid ❌"
        para_instance = None
        return f"❌ Error: {str(e)}"

def check_api_key():
    """Check if API key is set"""
    if para_instance is None:
        return "Please set your Groq API key first"
    return None

def update_model_selection(model_choice):
    """Update model selection"""
    global para_instance
    
    if para_instance:
        para_instance.model = model_choice
        return f"Model updated to: {model_choice}"
    else:
        return "Set API key first"

def research_topic(topic, include_domains, exclude_domains):
    """Research a topic with domain filters"""
    # Check if API key is set
    check_result = check_api_key()
    if check_result:
        return check_result
    
    if not topic:
        return "Please enter a topic to research"
    
    # Process domain lists
    include_list = [d.strip() for d in include_domains.split(",")] if include_domains else []
    exclude_list = [d.strip() for d in exclude_domains.split(",")] if exclude_domains else []
    
    try:
        # Perform research
        result = para_instance.research_digest(
            topic=topic,
            include_domains=include_list if include_list and include_list[0] else None,
            exclude_domains=exclude_list if exclude_list and exclude_list[0] else None
        )
        
        # Format response
        response = f"# Research: {topic}\n\n{result['content']}"
        
        # Add tool usage info if available
        if result.get("tool_usage"):
            response += f"\n\n*Tool Usage Information Available*"
            
        return response
    except Exception as e:
        return f"Error: {str(e)}"

def analyze_code(code_snippet, language, analysis_type):
    """Analyze code with Groq"""
    # Check if API key is set
    check_result = check_api_key()
    if check_result:
        return check_result
    
    if not code_snippet:
        return "Please enter code to analyze"
    
    try:
        # Perform analysis
        result = para_instance.evaluate_code(
            code_snippet=code_snippet,
            language=language,
            analysis_type=analysis_type
        )
        
        # Format response
        response = f"# Code Analysis ({language}, {analysis_type})\n\n{result['content']}"
        
        # Add tool usage info if available
        if result.get("tool_usage"):
            response += f"\n\n*Tool Usage Information Available*"
            
        return response
    except Exception as e:
        return f"Error: {str(e)}"

def connect_concepts_handler(concept_a, concept_b):
    """Connect two concepts"""
    # Check if API key is set
    check_result = check_api_key()
    if check_result:
        return check_result
    
    if not concept_a or not concept_b:
        return "Please enter both concepts"
    
    try:
        # Find connections
        result = para_instance.connect_concepts(
            concept_a=concept_a,
            concept_b=concept_b
        )
        
        # Format response
        response = f"# Connection: {concept_a} & {concept_b}\n\n{result['content']}"
        
        # Add tool usage info if available
        if result.get("tool_usage"):
            response += f"\n\n*Tool Usage Information Available*"
            
        return response
    except Exception as e:
        return f"Error: {str(e)}"

def query_knowledge_base(query):
    """Query the knowledge base"""
    # Check if API key is set
    check_result = check_api_key()
    if check_result:
        return check_result
    
    if not query:
        return "Please enter a query"
    
    try:
        # Query knowledge base
        result = para_instance.ask_knowledge_base(query=query)
        
        # Format response
        response = f"# Knowledge Base Query: {query}\n\n{result['response']}"
        
        # Add KB stats
        stats = result.get("knowledge_base_state", {})
        if stats:
            topics = stats.get("topics_researched", [])
            response += f"\n\n*Knowledge Base contains {len(topics)} topics: {', '.join(topics[:5])}" + \
                       ("..." if len(topics) > 5 else "") + "*"
            
        return response
    except Exception as e:
        return f"Error: {str(e)}"

# Create the Gradio interface
def create_gradio_app():
    # Define CSS for styling
    css = """
    .title-container {
        text-align: center;
        margin-bottom: 20px;
    }
    .container {
        margin: 0 auto;
        max-width: 1200px;
    }
    .tab-content {
        padding: 20px;
        border-radius: 10px;
        background-color: #f9f9f9;
    }
    """
    
    with gr.Blocks(css=css, title="PARA - Personal AI Research Assistant") as app:
        gr.Markdown(
            """
            <div class="title-container">
            # 🧠 PARA - Personal AI Research Assistant
            *Powered by Groq's Compound Beta models for intelligent research*
            </div>
            """
        )
        
        with gr.Row():
            with gr.Column(scale=4):
                api_key_input = gr.Textbox(
                    label="Groq API Key", 
                    placeholder="Enter your Groq API key here...",
                    type="password"
                )
            with gr.Column(scale=2):
                model_choice = gr.Radio(
                    ["compound-beta", "compound-beta-mini"],
                    label="Model Selection",
                    value="compound-beta"
                )
            with gr.Column(scale=1):
                validate_btn = gr.Button("Validate & Connect")
        
        api_status = gr.Markdown("### Status: Not connected")
        
        # Connect validation button
        validate_btn.click(
            fn=validate_api_key,
            inputs=[api_key_input],
            outputs=[api_status]
        )
        
        # Connect model selection
        model_choice.change(
            fn=update_model_selection,
            inputs=[model_choice],
            outputs=[api_status]
        )
        
        # Tabs for different features
        with gr.Tabs() as tabs:
            # Research Tab
            with gr.Tab("Research Topics"):
                with gr.Row():
                    with gr.Column(scale=1):
                        research_topic_input = gr.Textbox(
                            label="Research Topic",
                            placeholder="Enter a topic to research..."
                        )
                    with gr.Column(scale=1):
                        include_domains = gr.Textbox(
                            label="Include Domains (comma-separated)",
                            placeholder="arxiv.org, *.edu, example.com"
                        )
                        exclude_domains = gr.Textbox(
                            label="Exclude Domains (comma-separated)",
                            placeholder="wikipedia.org, twitter.com"
                        )
                research_btn = gr.Button("Research Topic")
                research_output = gr.Markdown("Results will appear here...")
                
                research_btn.click(
                    fn=research_topic,
                    inputs=[research_topic_input, include_domains, exclude_domains],
                    outputs=[research_output]
                )
                
                gr.Markdown("""
                ### Examples:
                - "Latest developments in quantum computing"
                - "Climate change mitigation strategies"
                - "Advancements in protein folding algorithms"
                
                *Include domains like "arxiv.org, *.edu" for academic sources*
                """)
            
            # Code Analysis Tab
            with gr.Tab("Code Analysis"):
                code_input = gr.Code(
                    label="Code Snippet",
                    language="python",
                    lines=10
                )
                with gr.Row():
                    language_select = gr.Dropdown(
                        ["python", "javascript", "java", "c++", "go", "rust", "typescript", "sql", "bash"],
                        label="Language",
                        value="python"
                    )
                    analysis_type = gr.Dropdown(
                        ["full", "security", "performance", "style"],
                        label="Analysis Type",
                        value="full"
                    )
                analyze_btn = gr.Button("Analyze Code")
                analysis_output = gr.Markdown("Results will appear here...")
                
                analyze_btn.click(
                    fn=analyze_code,
                    inputs=[code_input, language_select, analysis_type],
                    outputs=[analysis_output]
                )
                
                gr.Markdown("""
                ### Example Python Code:
                ```python
                def fibonacci(n):
                    if n <= 0:
                        return []
                    elif n == 1:
                        return [0]
                    else:
                        result = [0, 1]
                        for i in range(2, n):
                            result.append(result[i-1] + result[i-2])
                        return result
                
                print(fibonacci(10))
                ```
                """)
            
            # Concept Connections Tab
            with gr.Tab("Connect Concepts"):
                with gr.Row():
                    concept_a = gr.Textbox(
                        label="Concept A",
                        placeholder="First concept or field..."
                    )
                    concept_b = gr.Textbox(
                        label="Concept B",
                        placeholder="Second concept or field..."
                    )
                connect_btn = gr.Button("Find Connections")
                connection_output = gr.Markdown("Results will appear here...")
                
                connect_btn.click(
                    fn=connect_concepts_handler,
                    inputs=[concept_a, concept_b],
                    outputs=[connection_output]
                )
                
                gr.Markdown("""
                ### Example Concept Pairs:
                - "quantum computing" and "machine learning"
                - "blockchain" and "supply chain management"
                - "gene editing" and "ethics"
                """)
            
            # Knowledge Base Tab
            with gr.Tab("Knowledge Base"):
                kb_query = gr.Textbox(
                    label="Query Knowledge Base",
                    placeholder="Ask about topics in your knowledge base..."
                )
                kb_btn = gr.Button("Query Knowledge Base")
                kb_output = gr.Markdown("Results will appear here...")
                
                kb_btn.click(
                    fn=query_knowledge_base,
                    inputs=[kb_query],
                    outputs=[kb_output]
                )
                
                gr.Markdown("""
                ### Example Queries:
                - "What have we learned about quantum computing?"
                - "Summarize our research on AI safety"
                - "What connections exist between the topics we've studied?"
                """)
        
        gr.Markdown("""
        ## About PARA
        
        PARA (Personal AI Research Assistant) leverages Groq's compound models with agentic capabilities to help you research topics, analyze code, find connections between concepts, and build a personalized knowledge base.
        
        **How it works:**
        1. Set your Groq API key
        2. Choose between compound-beta (more powerful) and compound-beta-mini (faster)
        3. Use the tabs to access different features
        4. Your research is automatically saved to a knowledge base for future reference
        
        **Features:**
        - Web search with domain filtering
        - Code execution and analysis
        - Concept connections discovery
        - Persistent knowledge base
        """)
    
    return app

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