from smolagents import CodeAgent, HfApiModel, load_tool, tool
import yaml
import feedparser
import gradio as gr

@tool
def fetch_latest_arxiv_papers(keywords: list, num_results: int = 1) -> list:
    """Fetches the latest research papers from arXiv based on provided keywords.
    Args:
        keywords: A list of keywords to search for relevant papers.
        num_results: The number of papers to fetch (default is 5).
    """
    try:
        print(f"DEBUG: Searching arXiv papers with keywords: {keywords}")  # Debug input
        query = "+".join(keywords)
        url = f"http://export.arxiv.org/api/query?search_query=all:{query}&start=0&max_results={num_results}&sortBy=submittedDate&sortOrder=descending"
        feed = feedparser.parse(url)

        papers = []
        for entry in feed.entries:
            papers.append({
                "title": entry.title,
                "authors": ", ".join(author.name for author in entry.authors),
                "year": entry.published[:4],  # Extract year
                "abstract": entry.summary,
                "link": entry.link
            })

        return papers

    except Exception as e:
        print(f"ERROR: {str(e)}")  # Debug errors
        return [f"Error fetching research papers: {str(e)}"]


model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

agent = CodeAgent(
    model=model,
    tools=[fetch_latest_arxiv_papers],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="ScholarAgent",
    description="An AI agent that fetches the latest research papers from arXiv based on user-defined keywords and filters.",
    prompt_templates=prompt_templates
)

def search_papers(user_input):
    keywords = [kw.strip() for kw in user_input.split(",") if kw.strip()]  # Ensure valid keywords
    print(f"DEBUG: Received input keywords - {keywords}")  # Debug user input
    
    if not keywords:
        print("DEBUG: No valid keywords provided.")
        return "Error: Please enter at least one valid keyword."
    
    results = fetch_latest_arxiv_papers(keywords, num_results=3)  # Fetch 3 results
    print(f"DEBUG: Results received - {results}")  # Debug function output
    
    if isinstance(results, list) and results and isinstance(results[0], dict):
        return "\n\n".join([
            f"**Title:** {paper['title']}\n**Authors:** {paper['authors']}\n**Year:** {paper['year']}\n**Abstract:** {paper['abstract']}\n[Read More]({paper['link']})"
            for paper in results
        ])
    
    print("DEBUG: No results found.")
    return "No results found. Try different keywords."


# Create a simple Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# arXiv Research Paper Fetcher")
    keyword_input = gr.Textbox(label="Enter keywords (comma-separated)", placeholder="e.g., deep learning, reinforcement learning")
    output_display = gr.Markdown()
    search_button = gr.Button("Search")

    search_button.click(search_papers, inputs=[keyword_input], outputs=[output_display])

    print("DEBUG: Gradio UI is running. Waiting for user input...")

demo.launch()