from agents import Runner, gen_trace_id, trace
from search_agent import search_agent
from planner_agent import planner_agent, WebSearchItem, WebSearchPlan
from writer_agent import writer_agent, ReportData
from email_agent import email_agent
import asyncio

class ResearchManager:

    async def run(self, query: str):
        """Run the deep research process, yielding the status updates and the final report"""
        trace_id = gen_trace_id()

        with trace("Research Trace", trace_id=trace_id):
            print(f"View Trace: https://platform.openai.com/logs/trace?trace_id={trace_id}")
            yield f"View Trace: https://platform.openai.com/logs/trace?trace_id={trace_id}"
            
            print("Starting research...")

            search_plan = await self.plan_searches(query)
            yield "Searches planned, starting searching process..."

            search_results = await self.perform_searches(search_plan)
            yield "Searches complete, writing detailed report..."

            report = await self.write_report(query, search_results)
            yield "Report complete, sending email..."

            await self.send_email(report)
            yield "Email sent..."

            yield "Research completed. Displaying research results."
            yield report.markdown_report


    async def plan_searches(self, query: str) -> WebSearchPlan:
        """Plan the searches to perform for the query"""
        print("Planning searches...")

        result = await Runner.run(planner_agent, f"Query: {query}")

        print(f"Will perform {len(result.final_output.searches)} searches.")

        return result.final_output_as(WebSearchPlan)
    
    async def perform_searches(self, search_plan: WebSearchPlan) -> list[str]:
        """Perform the searches to perform for the query"""
        print("Searching...")

        tasks = [asyncio.create_task(self.search(item)) for item in search_plan.searches]
        
        num_completed = 0
        results = []
        for task in asyncio.as_completed(tasks):
            result = await task
            if result is not None:
                results.append(result)
                num_completed += 1
            print(f"Searching... {num_completed}/{len(tasks)} completed.")
        
        print("Finished searching.")
        return results

    async def search(self, item: WebSearchItem) -> str | None:
        """Perform a search for the query"""

        input = f"Search term: {item.query}\nReason for searching: {item.reason}"

        try:
            result = await Runner.run(search_agent, input)

            return result.final_output

        except Exception:
            return None
    
    async def write_report(self, query: str, search_results: list[str]) -> ReportData:
        """Write a detailed report for the query"""
        
        print("Thinking about a detailed report...")

        input = f"Original query: {query}\nSummarized search results: {search_results}"

        result = await Runner.run(writer_agent, input)

        print("Finished writing report.")
        return result.final_output_as(ReportData)
    
    async def send_email(self, report: ReportData) -> None:
        print("Writing email...")

        await Runner.run(email_agent, report.markdown_report)

        print("Email sent.")