from exa_py import Exa
from groq import Groq
import os

# Declare the exa search API
exa = Exa(api_key=os.getenv("EXA_API_KEY"))

# Define your API Model and key
client = Groq(api_key=os.getenv("GROQ_API_KEY"))
utilized_model = "llama3-70b-8192"

# Options for highlights from Exa search
highlights_options = {
    "num_sentences": 7,  # Length of highlights
    "highlights_per_url": 1,  # Get the best highlight for each URL
}


def get_report_fields(report_type):
    """
    Return a list of fields (tuples of field_key, field_label) based on the report type.
    """
    common_fields = [
        ("research_topic", "Research Topic"),
        ("research_question", "Research Question"),
        ("objectives", "Research Objectives (SMART)"),
        ("methodology", "Research Methodology"),
        ("timeline", "Research Timeline"),
        ("email", "Contact Email"),
    ]

    if report_type == "PhD Proposal":
        return common_fields + [("justification", "Justification for Methodology"),
                                ("literature_gap", "Literature Gap")]

    elif report_type == "Research Paper":
        return common_fields + [("introduction", "Introduction"), 
                                ("conclusion", "Conclusion")]

    elif report_type == "Thesis":
        return common_fields + [("abstract", "Abstract"),
                                ("limitations", "Limitations")]

    return common_fields

def call_dynamic_function_0(function_name, data):
    """
    Call the appropriate function to generate a report section.
    """
    function = globals().get(function_name)
    if function:
        return function(data)
    else:
        return "Error: Function not found."

  
def get_report_sections(report_type):
    """
    Return a dictionary mapping section names to corresponding functions based on report type.
    """
    report_sections = {
        "PhD Proposal": [
            ("Executive Summary", "generate_executive_summary"),
            ("Research Objectives", "generate_research_objectives"),
            ("Research Methodology", "generate_methodology_section"),
            ("Research Timeline", "generate_research_timeline"),
            ("All Sections", "generate_all_sections")  # New entry for all sections
        ],
        "Research Paper": [
            ("Introduction", "generate_proposal_introduction"),
            ("Literature Review", "generate_literature_review_outline"),
            ("Research Methodology", "generate_methodology_section"),
            ("Conclusion", "generate_contribution_statement"),
            ("All Sections", "generate_all_sections")  # New entry for all sections
        ],
        "Thesis": [
            ("Abstract", "generate_proposal_introduction"),
            ("Research Objectives", "generate_research_objectives"),
            ("Methodology", "generate_methodology_section"),
            ("Limitations", "generate_limitations_section"),
            ("All Sections", "generate_all_sections")  # New entry for all sections
        ]
    }
    return report_sections.get(report_type, [])

def generate_all_sections_0(data):
    """
    Generate all sections for the specified report type.
    """
    sections_content = {}

    # Get the report type from data
    report_type = data.get("report_type")

    # Retrieve sections based on report type
    report_sections = get_report_sections(report_type)

    # Loop through each section, skipping "All Sections" itself
    for section_name, function_name in report_sections:
        if function_name == "generate_all_sections":
            continue
        
        # Call the respective function dynamically
        section_content = call_dynamic_function(function_name, data)
        
        # Store the generated content in a dictionary
        sections_content[section_name] = section_content
    
    return sections_content  # Return the content dictionary
    
def generate_all_sections(data):
    sections_content = {}
    report_type = data.get("report_type")
    report_sections = get_report_sections(report_type)
    
    for section_name, function_name in report_sections:
        if function_name == "generate_all_sections":
            continue  # Skip "All Sections" itself
        
        update_session_tracker(section_name)
        
        # Debugging step: Ensure each section is being processed
        st.write(f"Generating content for: {section_name}")

        # Call the dynamic function and ensure data is being passed correctly
        section_function_name = dict(report_sections)[section_name]
        section_content = call_dynamic_function(section_function_name, data)
        
        if not section_content:
            st.error(f"No content generated for {section_name}")
        else:
            st.subheader(section_name)
            st.write(section_content)

        sections_content[section_name] = section_content

    return sections_content



# Define the actual generation functions, like:
def generate_executive_summary(data):
    return f"Executive Summary for {data['research_topic']}"
    
# Add other report generation functions similar to this

def call_llm(prompt):
    """
    Call the Exa and Groq APIs to generate content using the provided prompt.
    """
    # Perform search with Exa
    search_response = exa.search_and_contents(query=prompt, highlights=highlights_options, num_results=3, use_autoprompt=True)
    info = [sr.highlights[0] for sr in search_response.results]
    
    # System prompt for Groq LLM
    system_prompt = "You are an academic PhD proposal generator. Read the provided contexts and generate a well-structured research proposal based on them."
    user_prompt = f"Sources: {info}\nResearch Prompt: {prompt}"
    
    # Call the Groq model
    completion = client.chat.completions.create(
        model=utilized_model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    )
    
    # Return the generated content
    return completion.choices[0].message.content


# Functions to generate different sections of the research proposal

def generate_executive_summary(data):
    """
    Generate an executive summary based on the user's input data.
    """
    prompt = f"""
    Generate a concise executive summary for a PhD research proposal with the following details:
    
    Research Topic: {data.get("research_topic")}
    Research Question: {data.get("research_question")}
    Objectives: {data.get("objectives")}
    Methodology: {data.get("methodology")}
    Contribution: {data.get("contribution")}
    Literature Gap: {data.get("literature_gap")}
    """
    return call_llm(prompt)

      

def generate_literature_review_outline(data):
    """
    Generate a literature review outline based on the user's input data.
    """
    prompt = f"""
    Generate a structured outline for the literature review of a PhD thesis on the following topic:
    
    Research Topic: {data.get("research_topic")}
    Key Authors: {data.get("key_authors")}
    Recent Developments: {data.get("recent_developments")}
    Gaps in Literature: {data.get("literature_gap")}
    """
    return call_llm(prompt)


def generate_methodology_section(data):
    """
    Generate a methodology section for a PhD proposal.
    """
    prompt = f"""
    Write a detailed research methodology section for a PhD proposal based on the following:
    
    Research Topic: {data.get("research_topic")}
    Data Collection Methods: {data.get("data_collection")}
    Data Analysis Methods: {data.get("data_analysis")}
    Justification: {data.get("justification")}
    """
    return call_llm(prompt)


def generate_research_objectives(data):
    """
    Generate research objectives using the SMART framework.
    """
    prompt = f"""
    Generate a detailed list of short-term and long-term research objectives for the following PhD thesis topic:
    
    Research Topic: {data.get("research_topic")}
    Objectives: {data.get("objectives")}
    
    The objectives should follow the SMART criteria (Specific, Measurable, Achievable, Relevant, Time-bound).
    """
    return call_llm(prompt)


def generate_hypotheses(data):
    """
    Generate hypotheses for the research.
    """
    prompt = f"""
    Generate research hypotheses based on the following topic:
    
    Research Topic: {data.get("research_topic")}
    Research Question: {data.get("research_question")}
    """
    return call_llm(prompt)


def generate_contribution_statement(data):
    """
    Generate a contribution statement for a PhD proposal.
    """
    prompt = f"""
    Generate a statement of contribution for the following PhD research proposal:
    
    Research Topic: {data.get("research_topic")}
    Contribution to the Field: {data.get("contribution")}
    
    The statement should highlight how the research will address existing gaps and advance knowledge in the field.
    """
    return call_llm(prompt)


def generate_research_timeline(data):
    """
    Generate a detailed research timeline for a PhD thesis.
    """
    prompt = f"""
    Generate a detailed research timeline for completing a PhD thesis on the following topic:
    
    Research Topic: {data.get("research_topic")}
    Total Timeframe: {data.get("total_timeframe")}
    
    The timeline should break down tasks into manageable phases (e.g., literature review, data collection, analysis) with deadlines.
    """
    return call_llm(prompt)


def generate_proposal_introduction(data):
    """
    Generate a proposal introduction based on the user's input data.
    """
    prompt = f"""
    Write an engaging introduction for a PhD proposal on the following research topic:
    
    Research Topic: {data.get("research_topic")}
    Research Problem: {data.get("research_problem")}
    
    The introduction should provide background, introduce the problem, and explain the significance of the research.
    """
    return call_llm(prompt)


def generate_limitations_section(data):
    """
    Generate a section on research limitations.
    """
    prompt = f"""
    Generate a section describing the potential limitations and challenges of the following research:
    
    Research Topic: {data.get("research_topic")}
    Methodology: {data.get("methodology")}
    
    The limitations should address possible obstacles and suggest ways to mitigate them.
    """
    return call_llm(prompt)


def generate_future_work_section(data):
    """
    Generate a future work section for a research proposal.
    """
    prompt = f"""
    Generate a section on future work based on the following research:
    
    Research Topic: {data.get("research_topic")}
    Contribution: {data.get("contribution")}
    
    The future work section should suggest further areas for research that could build upon the findings.
    """
    return call_llm(prompt)