import os
import gradio as gr
from crewai import Agent, Task, Crew, Process
from langchain_together import ChatTogether
from agents.web_crawler_agent import WebCrawlerAgent
from agents.summarizer_agent import SummarizerAgent
from agents.report_generator_agent import ReportGeneratorAgent
from dotenv import load_dotenv

load_dotenv()

model = os.getenv("MODEL_NAME")
os.environ['OPENAI_API_KEY'] = os.getenv("OPENAI_API_KEY")
os.environ['LANGCHAIN_TRACING_V2'] = "true"
os.environ['LANGCHAIN_API_KEY'] = os.getenv("LANGCHAIN_API_KEY")
os.environ['LANGCHAIN_PROJECT'] = "KnowledgeKeeper"

llm = ChatTogether(model=model)

web_crawler_agent = WebCrawlerAgent(
    llm=llm,
    role="web crawler",
    backstory="I search the web to gather information on a given topic.",
    goal="Fetch web search results for a given query.",
    serper_api_key=os.getenv("SERPER_API_KEY")
)


summarizer_agent = SummarizerAgent(llm)

report_generator_agent = ReportGeneratorAgent(llm=llm)

def generate_report(query, report_style):
    web_crawler_task = Task(
        description=query, 
        expected_output="A detailed response about the query.",
        agent=web_crawler_agent
    )

    summarizer_task = Task(
        description="Summarize the web data.", 
        expected_output="Concise summaries of the search results.", 
        agent=summarizer_agent
    )

    report_task = Task(
        description=f"Generate a final research report in '{report_style}' style by removing bias if any.",
        expected_output="A comprehensive research report.",
        agent=report_generator_agent
    )

    crew = Crew(
        agents=[web_crawler_agent, summarizer_agent, report_generator_agent],
        tasks=[web_crawler_task, summarizer_task, report_task],
        verbose=True,
        process=Process.sequential,
        context_flow={
            "summarizer_task": "web_crawler_task", 
            "report_task": "summarizer_task",     
        }
    )

    result = crew.kickoff()
    formatted_report = f"###### Report Generated in {report_style} Style:\n\n{result.raw}"

    return formatted_report

with gr.Blocks() as demo:
    gr.Markdown("## 🚀 Research Report Generation and Analysis Tool")
    gr.Markdown("#### Enter your query and choose the report style to generate a detailed, formatted report!")

    with gr.Row():
        query_input = gr.Textbox(
            label="Enter Your Query", 
            placeholder="e.g., Recent advancements in Indian Agriculture",
            lines=2,
            elem_id="query-input"
        )

    with gr.Row():
        report_style = gr.Radio(
            choices=[ "Detailed Analysis", "Brief Summary","Bullet Points"],
            label="Select Report Style",
            value="Detailed Analysis"
        )
    
    with gr.Row():
        generate_button = gr.Button("Generate Report", elem_id="generate-button", variant="primary")
    
    with gr.Row():
        report_output = gr.Markdown(
            label="Generated Report"
        )
    
    generate_button.click(
        generate_report, 
        inputs=[query_input, report_style], 
        outputs=[report_output]
    )

demo.launch()