import asyncio # This will helpus to handle tasks without blocking execution
import streamlit as st  # Build the web application
from typing import Dict, Any, List  
from agents import Agent, Runner, trace
from agents import set_default_openai_key
from firecrawl import FirecrawlApp
from agents.tool import function_tool

# Setup the page configuration (done using streamlit)
st.set_page_config(
    page_title = "OpenAI based Deep Research Agent",
    page_icon = "📚",
    layout = "wide"
)

# Initialize session state for API Key if don't exist
if "openai_api_key" not in st.session_state:
    st.session_state.openai_api_key = ""
if "firecrawl_api_key" not in st.session_state:
    st.session_state.firecrawl_api_key = ""

# Sidebar for API Key
with st.sidebar:
    st.title("API Configuration")
    openai_api_key = st.text_input(
        "OpenAI API Key",
        value = st.session_state.openai_api_key,
        type = "password"
    )

    firecrawl_api_key = st.text_input(
        "Firecrawl API Key",
        value = st.session_state.firecrawl_api_key,
        type = "password"
    )

    if openai_api_key:
        st.session_state.openai_api_key = openai_api_key
        set_default_openai_key(openai_api_key)
    if firecrawl_api_key:
        st.session_state.firecrawl_api_key = firecrawl_api_key

# Main Application and Input Field
st.title("🔍 OpenAI Deep Research Agent")
st.markdown("This OpenAI Agent from OpenAI Agent SDK performs deep research on any topic using Firecrawl")

# This takes the input from the user for the specific research related concern
research_topic = st.text_input("Enter research topic: ", placeholder = "e.g., Latest Development in AI")

# This function tools help us to register this function
# as a tool to the agent
@function_tool
async def deep_research(query: str, max_depth: int, time_limit: int, max_urls: int):
    """
    Perform comprehensive web research using Firecrawl's deep research endpoint.
    """
    try:
        # Initialize the firecrawl using the saved API key
        firecrawl_app = FirecrawlApp(api_key =  st.session_state.firecrawl_api_key)
        params = {
            "maxDepth": max_depth,
            "timeLimit": time_limit,
            "maxUrls": max_urls
        }
        # Callback Setup for real-time update
        def on_activity(activity):
            st.write(f"[{activity['type']}]{activity['message']}")

        # Run the deep research using firecrawl
        with st.spinner("Performing Deep Research..."):
            resp = firecrawl_app.deep_research(
                query = query,
                params = params,
                on_activity = on_activity
            )
            
        return {
            "success" : True,
            "final_analysis" : resp["data"]["finalAnalysis"]
            "sources_count": len(resp["data"]["sources"]),
            "sources":resp["data"]["sources"]
        }
    except Exception as e:
        st.error(f"Deep Research Error: {str(e)}")
        return {
            "error" : str(e),
            "success" : False
        }