import streamlit as st
from gpt_researcher import GPTResearcher
import asyncio
import nest_asyncio
import os
from contextlib import contextmanager
from io import StringIO
import sys

# Access secrets
openai_api_key = st.secrets["OPENAI_API_KEY"]
tavily_api_key = st.secrets["TAVILY_API_KEY"]

# Apply the asyncio patch from nest_asyncio if required
nest_asyncio.apply()

# Set the document path environment variable
os.environ['DOC_PATH'] = './local'  # Path to the folder with documents

# Constants
REPORT_TYPE = "research_report"

# Function to capture output to the standard output
@contextmanager
def st_capture(output_func):
    old_out = sys.stdout
    sys.stdout = StringIO()
    try:
        yield
        output_func(sys.stdout.getvalue())
    finally:
        sys.stdout = old_out

# Function to handle asynchronous calls
def run_async(coroutine):
    loop = asyncio.get_event_loop()
    return loop.run_until_complete(coroutine)

# Define the asynchronous function to fetch the report
async def fetch_report(query, report_type):
    """
    Fetch a research report based on the provided query and report type.
    Research is conducted on a local document.
    """
    try:
        researcher = GPTResearcher(query=query, report_type=report_type, report_source='local')
        await researcher.conduct_research()
        return await researcher.write_report()
    except Exception as e:
        return f"Error during research: {str(e)}"

# Streamlit interface
st.title("Google Algo Leak Reporting Tool")

# User input for the query using a text area
query = st.text_area(
    "Enter your research query:",
    "Extract all the information about how the ranking for internal links works.",
    height=150  # Adjustable height
)

# Placeholder for the progress expander
progress_expander = st.expander("See research progress", expanded=True)
progress_placeholder = progress_expander.empty()

# Start the report generation process
if st.button("Generate Report"):
    if not query:
        st.warning("Please enter a query to generate a report.")
    else:
        # Display initial progress information
        progress_placeholder.info("Starting research...")
        # Run the research asynchronously and capture output
        with st_capture(progress_placeholder.code):
            report = run_async(fetch_report(query, REPORT_TYPE))
        
        if report and not report.startswith("Error"):
            st.success("Report generated successfully!")
            st.write(report)  # Display the report in the app
            # Create a download button for the report
            st.download_button(
                label="Download Report as Text File",
                data=report,
                file_name="research_report.txt",
                mime="text/plain"
            )
        else:
            st.error(report)  # Show the error message if any