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

# 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"
DOCUMENT_FILE = 'removed_code.txt'  # Name of the document file

# 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 specified by DOCUMENT_FILE.
    """
    researcher = GPTResearcher(query=query, report_type=report_type, report_source='local')
    await researcher.conduct_research()
    report = await researcher.write_report()
    return report

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

# User input for the query using a text area for multi-line input
query = st.text_area(
    "Enter your research query:",
    "Extract all the information about how the ranking for internal links works.",
    height=150  # You can adjust the height as needed
)


# Button to generate report
if st.button("Generate Report"):
    if not query:
        st.warning("Please enter a query to generate a report.")
    else:
        # Collapsible expander to show progress
        with st.expander("See research progress"):
            with st_capture(st.write):
                report = run_async(fetch_report(query, "research_report"))
        
        if report:
            st.success("Report generated successfully!")
            st.write(report)  # Display the report in the app
            
            # Create a download button and provide the report as a downloadable file
            st.download_button(
                label="Download Report as Text File",
                data=report,
                file_name="research_report.txt",
                mime="text/plain"
            )
        else:
            st.error("Failed to generate the report.")