import streamlit as st
import os
import json
from dataclasses import dataclass, field
from typing import List
from groq import Groq
from json.decoder import JSONDecodeError
from pydantic_settings import BaseSettings
from tavily import TavilyClient

# Configuration class for API keys and settings
class Config(BaseSettings):
    GROQ_API_KEY: str
    LLM_REASONING: str = "deepseek-r1-distill-llama-70b"  # Example Groq model for reasoning tasks
    LLM_REGULAR: str = "deepseek-r1-distill-llama-70b"     # Example Groq model for regular tasks
    TAVILY_API_KEY: str

# Load configuration from environment variables
config = Config(
    GROQ_API_KEY=os.environ["GROQ_API_KEY"],
    TAVILY_API_KEY=os.environ["TAVILY_API_KEY"]
)

# Data classes for state management
@dataclass
class Search:
    url: str = ""
    content: str = ""

@dataclass
class Research:
    search_history: List[Search] = field(default_factory=list)
    latest_summary: str = ""
    reflection_iteration: int = 0

@dataclass
class Paragraph:
    title: str = ""
    content: str = ""
    research: Research = field(default_factory=Research)

@dataclass
class State:
    report_title: str = ""
    paragraphs: List[Paragraph] = field(default_factory=list)

# Helper functions
def remove_reasoning_from_output(output):
    """Remove reasoning tags from output if present."""
    return output.split("</think>")[-1].strip()

def clean_json_tags(text):
    """Remove JSON markdown tags from text."""
    return text.replace("```json\n", "").replace("\n```", "")

def clean_markdown_tags(text):
    """Remove markdown tags from text."""
    return text.replace("```markdown\n", "").replace("\n```", "")

def tavily_search(query, include_raw_content=True, max_results=3):
    """Perform a web search using Tavily."""
    tavily_client = TavilyClient(api_key=config.TAVILY_API_KEY)
    return tavily_client.search(query, include_raw_content=include_raw_content, max_results=max_results)

def update_state_with_search_results(search_results, idx_paragraph, state):
    """Update the state with search results."""
    for search_result in search_results["results"]:
        search = Search(url=search_result["url"], content=search_result["raw_content"])
        state.paragraphs[idx_paragraph].research.search_history.append(search)

# Agent classes
## Report Structure Agent
output_schema_report_structure = {
    "type": "array",
    "items": {
        "type": "object",
        "properties": {
            "title": {"type": "string"},
            "content": {"type": "string"}
        }
    }
}

SYSTEM_PROMPT_REPORT_STRUCTURE = f"""
You are a Deep Research assistant. Given a query, plan a structure for a report and the paragraphs to be included.
Make sure that the ordering of paragraphs makes sense.
Once the outline is created, you will be given tools to search the web and reflect for each of the section separately.
Format the output in json with the following json schema definition:

<OUTPUT JSON SCHEMA>
{json.dumps(output_schema_report_structure, indent=2)}
</OUTPUT JSON SCHEMA>

Title and content properties will be used for deeper research.
Make sure that the output is a json object with an output json schema defined above.
Only return the json object, no explanation or additional text.
"""

class ReportStructureAgent:
    def __init__(self, query: str):
        self.groq_client = Groq(api_key=config.GROQ_API_KEY)
        self.query = query

    def run(self) -> str:
        response = self.groq_client.chat.completions.create(
            model=config.LLM_REASONING,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT_REPORT_STRUCTURE},
                {"role": "user", "content": self.query}
            ]
        )
        return response.choices[0].message.content

    def mutate_state(self, state: State) -> State:
        report_structure = self.run()
        report_structure = remove_reasoning_from_output(report_structure)
        report_structure = clean_json_tags(report_structure)
        report_structure = json.loads(report_structure)
        for paragraph in report_structure:
            state.paragraphs.append(Paragraph(title=paragraph["title"], content=paragraph["content"]))
        return state

## First Search Agent
input_schema_first_search = {
    "type": "object",
    "properties": {
        "title": {"type": "string"},
        "content": {"type": "string"}
    }
}

output_schema_first_search = {
    "type": "object",
    "properties": {
        "search_query": {"type": "string"},
        "reasoning": {"type": "string"}
    }
}

SYSTEM_PROMPT_FIRST_SEARCH = f"""
You are a Deep Research assistant. You will be given a paragraph in a report, its title and expected content in the following json schema definition:

<INPUT JSON SCHEMA>
{json.dumps(input_schema_first_search, indent=2)}
</INPUT JSON SCHEMA>

You can use a web search tool that takes a 'search_query' as parameter.
Your job is to reflect on the topic and provide the most optimal web search query to enrich your current knowledge.
Format the output in json with the following json schema definition:

<OUTPUT JSON SCHEMA>
{json.dumps(output_schema_first_search, indent=2)}
</OUTPUT JSON SCHEMA>

Make sure that the output is a json object with an output json schema defined above.
Only return the json object, no explanation or additional text.
"""

class FirstSearchAgent:
    def __init__(self):
        self.groq_client = Groq(api_key=config.GROQ_API_KEY)

    def run(self, message) -> str:
        response = self.groq_client.chat.completions.create(
            model=config.LLM_REGULAR,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT_FIRST_SEARCH},
                {"role": "user", "content": message}
            ]
        )
        response = remove_reasoning_from_output(response.choices[0].message.content)
        response = clean_json_tags(response)
        response = json.loads(response)
        return response

## First Summary Agent
input_schema_first_summary = {
    "type": "object",
    "properties": {
        "title": {"type": "string"},
        "content": {"type": "string"},
        "search_query": {"type": "string"},
        "search_results": {
            "type": "array",
            "items": {"type": "string"}
        }
    }
}

output_schema_first_summary = {
    "type": "object",
    "properties": {
        "paragraph_latest_state": {"type": "string"}
    }
}

SYSTEM_PROMPT_FIRST_SUMMARY = f"""
You are a Deep Research assistant. You will be given a search query, search results and the paragraph a report that you are researching following json schema definition:

<INPUT JSON SCHEMA>
{json.dumps(input_schema_first_summary, indent=2)}
</INPUT JSON SCHEMA>

Your job is to write the paragraph as a researcher using the search results to align with the paragraph topic and structure it properly to be included in the report.
Format the output in json with the following json schema definition:

<OUTPUT JSON SCHEMA>
{json.dumps(output_schema_first_summary, indent=2)}
</OUTPUT JSON SCHEMA>

Make sure that the output is a json object with an output json schema defined above.
Only return the json object, no explanation or additional text.
"""

class FirstSummaryAgent:
    def __init__(self):
        self.groq_client = Groq(api_key=config.GROQ_API_KEY)

    def run(self, message) -> str:
        response = self.groq_client.chat.completions.create(
            model=config.LLM_REGULAR,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT_FIRST_SUMMARY},
                {"role": "user", "content": message}
            ]
        )
        return response.choices[0].message.content

    def mutate_state(self, message: str, idx_paragraph: int, state: State) -> State:
        summary = self.run(message)
        summary = remove_reasoning_from_output(summary)
        summary = clean_json_tags(summary)
        try:
            summary = json.loads(summary)
        except JSONDecodeError:
            summary = {"paragraph_latest_state": summary}
        state.paragraphs[idx_paragraph].research.latest_summary = summary["paragraph_latest_state"]
        return state

## Reflection Agent
input_schema_reflection = {
    "type": "object",
    "properties": {
        "title": {"type": "string"},
        "content": {"type": "string"},
        "paragraph_latest_state": {"type": "string"}
    }
}

output_schema_reflection = {
    "type": "object",
    "properties": {
        "search_query": {"type": "string"},
        "reasoning": {"type": "string"}
    }
}

SYSTEM_PROMPT_REFLECTION = f"""
You are a Deep Research assistant. You are responsible for constructing comprehensive paragraphs for a research report. You will be provided paragraph title and planned content summary, also the latest state of the paragraph that you have already created all in the following json schema definition:

<INPUT JSON SCHEMA>
{json.dumps(input_schema_reflection, indent=2)}
</INPUT JSON SCHEMA>

You can use a web search tool that takes a 'search_query' as parameter.
Your job is to reflect on the current state of the paragraph text and think if you haven’t missed some critical aspect of the topic and provide the most optimal web search query to enrich the latest state.
Format the output in json with the following json schema definition:

<OUTPUT JSON SCHEMA>
{json.dumps(output_schema_reflection, indent=2)}
</OUTPUT JSON SCHEMA>

Make sure that the output is a json object with an output json schema defined above.
Only return the json object, no explanation or additional text.
"""

class ReflectionAgent:
    def __init__(self):
        self.groq_client = Groq(api_key=config.GROQ_API_KEY)

    def run(self, message) -> str:
        response = self.groq_client.chat.completions.create(
            model=config.LLM_REGULAR,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT_REFLECTION},
                {"role": "user", "content": message}
            ]
        )
        response = remove_reasoning_from_output(response.choices[0].message.content)
        response = clean_json_tags(response)
        response = json.loads(response)
        return response

## Reflection Summary Agent
input_schema_reflection_summary = {
    "type": "object",
    "properties": {
        "title": {"type": "string"},
        "content": {"type": "string"},
        "search_query": {"type": "string"},
        "search_results": {
            "type": "array",
            "items": {"type": "string"}
        },
        "paragraph_latest_state": {"type": "string"}
    }
}

output_schema_reflection_summary = {
    "type": "object",
    "properties": {
        "updated_paragraph_latest_state": {"type": "string"}
    }
}

SYSTEM_PROMPT_REFLECTION_SUMMARY = f"""
You are a Deep Research assistant.
You will be given a search query, search results, paragraph title and expected content for the paragraph in a report that you are researching.
You are iterating on the paragraph and the latest state of the paragraph is also provided.
The data will be in the following json schema definition:

<INPUT JSON SCHEMA>
{json.dumps(input_schema_reflection_summary, indent=2)}
</INPUT JSON SCHEMA>

Your job is to enrich the current latest state of the paragraph with the search results considering expected content.
Do not remove key information from the latest state and try to enrich it, only add information that is missing.
Structure the paragraph properly to be included in the report.
Format the output in json with the following json schema definition:

<OUTPUT JSON SCHEMA>
{json.dumps(output_schema_reflection_summary, indent=2)}
</OUTPUT JSON SCHEMA>

Make sure that the output is a json object with an output json schema defined above.
Only return the json object, no explanation or additional text.
"""

class ReflectionSummaryAgent:
    def __init__(self):
        self.groq_client = Groq(api_key=config.GROQ_API_KEY)

    def run(self, message) -> str:
        response = self.groq_client.chat.completions.create(
            model=config.LLM_REGULAR,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT_REFLECTION_SUMMARY},
                {"role": "user", "content": message}
            ]
        )
        return response.choices[0].message.content

    def mutate_state(self, message: str, idx_paragraph: int, state: State) -> State:
        summary = self.run(message)
        summary = remove_reasoning_from_output(summary)
        summary = clean_json_tags(summary)
        try:
            summary = json.loads(summary)
        except JSONDecodeError:
            summary = {"updated_paragraph_latest_state": summary}
        state.paragraphs[idx_paragraph].research.latest_summary = summary["updated_paragraph_latest_state"]
        return state

## Report Formatting Agent
input_schema_report_formatting = {
    "type": "array",
    "items": {
        "type": "object",
        "properties": {
            "title": {"type": "string"},
            "paragraph_latest_state": {"type": "string"}
        }
    }
}

SYSTEM_PROMPT_REPORT_FORMATTING = f"""
You are a Deep Research assistant. You have already performed the research and constructed final versions of all paragraphs in the report.
You will get the data in the following json format:

<INPUT JSON SCHEMA>
{json.dumps(input_schema_report_formatting, indent=2)}
</INPUT JSON SCHEMA>

Your job is to format the Report nicely and return it in Markdown.
If Conclusion paragraph is not present, add it to the end of the report from the latest state of the other paragraphs.
Use titles of the paragraphs to create a title for the report.
"""

class ReportFormattingAgent:
    def __init__(self):
        self.groq_client = Groq(api_key=config.GROQ_API_KEY)

    def run(self, message) -> str:
        response = self.groq_client.chat.completions.create(
            model=config.LLM_REASONING,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT_REPORT_FORMATTING},
                {"role": "user", "content": message}
            ]
        )
        summary = response.choices[0].message.content
        summary = remove_reasoning_from_output(summary)
        summary = clean_markdown_tags(summary)
        return summary

# Streamlit Web Interface
st.title("Research Agent with Groq")

query = st.text_input("Enter your research query:", "")

if st.button("Generate Report"):
    if query:
        with st.spinner("Generating report..."):
            try:
                # Initialize state
                STATE = State()
                STATE.report_title = query  # Use query as initial report title

                # Define constants
                NUM_REFLECTIONS = 2
                NUM_RESULTS_PER_SEARCH = 3
                CAP_SEARCH_LENGTH = 20000

                # Initialize agents
                report_structure_agent = ReportStructureAgent(query)
                first_search_agent = FirstSearchAgent()
                first_summary_agent = FirstSummaryAgent()
                reflection_agent = ReflectionAgent()
                reflection_summary_agent = ReflectionSummaryAgent()
                report_formatting_agent = ReportFormattingAgent()

                # Generate report structure
                STATE = report_structure_agent.mutate_state(STATE)

                # Process each paragraph
                for j in range(len(STATE.paragraphs)):
                    # First Search
                    message = json.dumps({
                        "title": STATE.paragraphs[j].title,
                        "content": STATE.paragraphs[j].content
                    })
                    output = first_search_agent.run(message)
                    search_results = tavily_search(output["search_query"], max_results=NUM_RESULTS_PER_SEARCH)
                    update_state_with_search_results(search_results, j, STATE)

                    # First Summary
                    message = {
                        "title": STATE.paragraphs[j].title,
                        "content": STATE.paragraphs[j].content,
                        "search_query": search_results["query"],
                        "search_results": [result["raw_content"][:CAP_SEARCH_LENGTH] for result in search_results["results"] if result["raw_content"]]
                    }
                    STATE = first_summary_agent.mutate_state(json.dumps(message), j, STATE)

                    # Reflection iterations
                    for i in range(NUM_REFLECTIONS):
                        # Reflection
                        message = {
                            "paragraph_latest_state": STATE.paragraphs[j].research.latest_summary,
                            "title": STATE.paragraphs[j].title,
                            "content": STATE.paragraphs[j].content
                        }
                        output = reflection_agent.run(json.dumps(message))

                        # Reflection Search
                        search_results = tavily_search(output["search_query"])
                        update_state_with_search_results(search_results, j, STATE)

                        # Reflection Summary
                        message = {
                            "title": STATE.paragraphs[j].title,
                            "content": STATE.paragraphs[j].content,
                            "search_query": search_results["query"],
                            "search_results": [result["raw_content"][:CAP_SEARCH_LENGTH] for result in search_results["results"] if result["raw_content"]],
                            "paragraph_latest_state": STATE.paragraphs[j].research.latest_summary
                        }
                        STATE = reflection_summary_agent.mutate_state(json.dumps(message), j, STATE)

                # Generate final report
                report_data = [{"title": paragraph.title, "paragraph_latest_state": paragraph.research.latest_summary} for paragraph in STATE.paragraphs]
                final_report = report_formatting_agent.run(json.dumps(report_data))

                # Display the report
                st.markdown(final_report)
            except Exception as e:
                st.error(f"An error occurred: {str(e)}")
    else:
        st.warning("Please enter a query.")