import streamlit as st
import pandas as pd
from smolagents import CodeAgent, DuckDuckGoSearchTool
from typing import Union, List, Dict
from groq import Groq
import os

class GroqLLM:
    """Compatible LLM interface for smolagents CodeAgent"""
    def __init__(self, model_name="llama-3.1-8B-Instant"):
        self.client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
        self.model_name = model_name
    
    def __call__(self, prompt: Union[str, dict, List[Dict]]) -> str:
        try:
            # Handle different prompt formats
            if isinstance(prompt, (dict, list)):
                prompt_str = str(prompt)
            else:
                prompt_str = str(prompt)
            
            # Create a properly formatted message
            completion = self.client.chat.completions.create(
                model=self.model_name,
                messages=[{
                    "role": "user",
                    "content": prompt_str
                }],
                temperature=0.7,
                max_tokens=1024,
                stream=False
            )
            return completion.choices[0].message.content if completion.choices else "Error: No response generated"
        except Exception as e:
            error_msg = f"Error generating response: {str(e)}"
            print(error_msg)
            return error_msg

def create_analysis_prompt(topic: str, search_results: str) -> str:
    """Creates a structured prompt for news analysis"""
    return f"""Analyze the following news information about {topic}. 
    Search Results: {search_results}
    
    Please provide:
    1. Summary of key points
    2. Main stakeholders involved
    3. Potential implications
    4. Analysis of different perspectives
    5. Fact-check of major claims (if applicable)
    
    Format the analysis in a clear, journalistic style."""

def log_agent_activity(prompt: str, result: str, agent_name: str):
    """Logs agent activities in the Streamlit interface"""
    with st.expander("View Agent Activity Log"):
        st.write(f"### Agent Activity ({agent_name}):")
        st.write("**Input Prompt:**")
        st.code(prompt, language="text")
        st.write("**Analysis Output:**")
        st.code(result, language="text")

# Initialize Streamlit app
st.set_page_config(page_title="News Analysis Tool", layout="wide")

# Title and description
st.title("🔍 AI News Analysis Tool")
st.write("""
This tool uses advanced AI to analyze news topics, providing comprehensive insights 
and analysis using real-time data from the web. Powered by Groq's LLama 3.1 8B 
Instant model and DuckDuckGo search.
""")

# Initialize the agents
try:
    # Initialize LLM and tools
    llm = GroqLLM()
    search_tool = DuckDuckGoSearchTool()
    
    # Create the analysis agent
    news_agent = CodeAgent(
        tools=[search_tool],
        model=llm
    )
    
    # Input section
    news_topic = st.text_input(
        "Enter News Topic or Query:",
        placeholder="E.g., Recent developments in renewable energy"
    )
    
    # Analysis options
    col1, col2 = st.columns(2)
    with col1:
        search_depth = st.slider(
            "Search Depth (number of results)",
            min_value=3,
            max_value=10,
            value=5
        )
    with col2:
        analysis_type = st.selectbox(
            "Analysis Type",
            ["Comprehensive", "Quick Summary", "Technical", "Simplified"]
        )
    
    # Generate analysis button
    if st.button("Analyze News"):
        if news_topic:
            with st.spinner("Gathering information and analyzing..."):
                try:
                    # First, get search results
                    search_results = search_tool.run(
                        f"Latest news about {news_topic} last 7 days"
                    )
                    
                    # Create analysis prompt
                    analysis_prompt = create_analysis_prompt(news_topic, search_results)
                    
                    # Get analysis from the agent
                    analysis_result = news_agent.run(analysis_prompt)
                    
                    # Display results
                    st.subheader("📊 Analysis Results")
                    st.markdown(analysis_result)
                    
                    # Log the activity
                    log_agent_activity(
                        analysis_prompt,
                        analysis_result,
                        "News Analysis Agent"
                    )
                    
                except Exception as e:
                    st.error(f"An error occurred during analysis: {str(e)}")
        else:
            st.warning("Please enter a news topic to analyze.")

    # Add helpful tips
    with st.expander("💡 Tips for Better Results"):
        st.write("""
        - Be specific with your topic for more focused analysis
        - Use keywords related to recent events for timely information
        - Consider including timeframes in your query
        - Try different analysis types for various perspectives
        """)

except Exception as e:
    st.error(f"""
    Failed to initialize the application: {str(e)}
    
    Please ensure:
    1. Your GROQ_API_KEY is properly set in environment variables
    2. All required packages are installed
    3. You have internet connectivity for DuckDuckGo searches
    """)

# Footer
st.markdown("---")
st.caption(
    "Powered by Groq LLama 3.1 8B Instant, DuckDuckGo, and Streamlit | "
    "Created for news analysis and research purposes"
)