import requests
import cohere
import os
import gradio as gr
from langchain.agents import Tool
from datetime import datetime

# Fetch the API keys from environment variables
cohere_api_key = os.getenv('COHERE_API_KEY')
news_api_key = os.getenv('NEWS_API_KEY')

# Initialize Cohere client
cohere_client = cohere.Client(cohere_api_key)

# Tool 1: Fetch AI news using NewsAPI
def fetch_ai_news(query):
    search_query = 'artificial intelligence AND (NLP OR "computer vision" OR "deep learning" OR "machine learning")'
    url = "https://newsapi.org/v2/everything"
    parameters = {
        'q': search_query,
        'sortBy': 'relevancy',
        'language': 'en',
        'apiKey': news_api_key
    }
    response = requests.get(url, params=parameters)
    if response.status_code == 200:
        data = response.json()
        articles = data['articles']
        
        # Filter out articles with future dates
        current_date = datetime.now().date()
        valid_articles = [article for article in articles if datetime.strptime(article['publishedAt'][:10], "%Y-%m-%d").date() <= current_date]
        
        news_summary = "\n".join([f"{i+1}. {article['title']} - {article['source']['name']}" for i, article in enumerate(valid_articles[:5])])
        return f"Here are the top 5 news articles:\n{news_summary}"
    else:
        return "Sorry, I couldn't fetch the latest AI news at the moment."

# Tool 2: Generate a response using Cohere
def generate_response(query):
    prompt = f"User's question: {query}\n\nYour response:"
    response = cohere_client.generate(
        model='command-r-plus-04-2024',
        prompt=prompt,
        max_tokens=500,
        temperature=0.7
    )
    return response.generations[0].text.strip()

# Tool definitions
news_tool = Tool(
    name="Fetch AI News",
    func=fetch_ai_news,
    description="Use this tool to fetch the latest news about AI."
)

llm_tool = Tool(
    name="Generate AI Response",
    func=generate_response,
    description="Use this tool to generate AI-related answers."
)

# Agent Logic: Choose between news or Cohere LLM based on input
def agent_logic(query):
    if "news" in query.lower() or "trends" in query.lower() or any(word in query.lower() for word in ["latest", "updates", "recent"]):
        return news_tool.run(query)
    else:
        return llm_tool.run(query)

# Gradio Interface: Connect the agent logic to a Gradio app
def chatbot(user_input):
    response = agent_logic(user_input)
    
    # Check if the response looks like code
    if response.count('\n') > 2 and ('def ' in response or 'class ' in response or '```' in response):
        return response, response
    else:
        return response, None

# Custom CSS for Gradio
custom_css = """
body {
    background-color: #f0f4f8;
    font-family: 'Arial', sans-serif;
}
.container {
    max-width: 800px;
    margin: 0 auto;
    padding: 20px;
    background-color: white;
    border-radius: 10px;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}
h1 {
    color: #2c3e50;
    text-align: center;
    margin-bottom: 20px;
}
.input-area, .output-area {
    margin-bottom: 20px;
}
.submit-btn, .copy-btn {
    background-color: #3498db;
    color: white;
    border: none;
    padding: 10px 20px;
    border-radius: 5px;
    cursor: pointer;
    transition: background-color 0.3s;
}
.submit-btn:hover, .copy-btn:hover {
    background-color: #2980b9;
}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
    with gr.Column(elem_classes="container"):
        gr.Markdown("# AI Conversational Agent")
        gr.Markdown("Ask questions about AI or request the latest AI news.")
        
        with gr.Column(elem_classes="input-area"):
            input_text = gr.Textbox(lines=3, placeholder="Ask me anything...", label="Your Question")
            submit_button = gr.Button("Submit", elem_classes="submit-btn")
        
        with gr.Column(elem_classes="output-area"):
            output_text = gr.Textbox(lines=10, label="Response")
            output_code = gr.Code(language="python", label="Code Output", visible=False)
            copy_button = gr.Button("Copy Response", elem_classes="copy-btn")
    
    def clear_input(input_text):
        return ""
    
    def copy_response(response, code):
        return gr.Textbox.update(value=code if code else response, visible=True)
    
    submit_button.click(
        fn=chatbot,
        inputs=input_text,
        outputs=[output_text, output_code]
    ).then(
        fn=clear_input,
        inputs=input_text,
        outputs=input_text
    )

    copy_button.click(
        fn=copy_response,
        inputs=[output_text, output_code],
        outputs=gr.Textbox(visible=False),
    )

# Launch the Gradio app
demo.launch()