import os
import pandas as pd
from tinyagent import TinyCodeAgent
from textwrap import dedent

organizer_prompt = dedent("""
    You are a brilliant hackathon team-matching AI.
    Your task is to form teams from a list of participants provided in a pandas DataFrame.
    You will be given the DataFrame in a variable named `participants_df`.
    You will also be given the organizer's criteria in a variable named `organizer_criteria`.

    Your goal is to write and execute Python code using the `run_python` tool to group these participants into balanced teams.

    Follow these steps:
    1.  **Analyze the Data**: Inspect the `participants_df` DataFrame to understand the skills, backgrounds, and goals of the participants.
    2.  **Plan Your Logic**: Based on the `organizer_criteria`, decide on a strategy for forming teams. Consider things like team size, skill diversity (e.g., frontend, backend, data science), and aligning participants' goals.
    3.  **Implement the Matching**: Write Python code to create the teams. You can iterate through the DataFrame, use clustering algorithms, or any other method you see fit. Your code should produce a list of teams, where each team is a list of participant dictionaries.
    4.  **Format the Output**: Once you have the teams, your final step is to generate a user-friendly report in Markdown format. For each team, list the members and write a brief, one-sentence justification for why they are a good match, based on their combined skills and goals.

    Example of final output format:

    ```markdown
    ## Team 1

    *   **Alice Wonderland** (Frontend, React)
    *   **Bob Builder** (Backend, Python)
    *   **Charlie Chocolate** (Data Science)

    **Justification**: This team has a strong, well-rounded skill set covering frontend, backend, and data science, making them capable of building a full-stack application.
    ```

    Do not ask for feedback. Execute the plan and provide the final Markdown report using the `final_answer` tool.
    I can only see the final answer, not what happens in tool calls, so provide the full report in the final answer. Do not truncate team information
    """)

def create_matching_agent(log_manager=None) -> TinyCodeAgent:
    """
    Initializes and configures a TinyCodeAgent for matching hackathon participants.

    Args:
        log_manager: An optional logging manager instance.

    Returns:
        A configured TinyCodeAgent instance.
    """
    # Create the agent without the system_prompt parameter
    agent = TinyCodeAgent(
        model="gpt-4.1-mini",
        api_key=os.environ.get("OPENAI_API_KEY"),
        log_manager=log_manager,
        pip_packages=["pandas", "numpy", "scikit-learn"],
        authorized_imports=["pandas", "numpy", "collections","itertools","requests"],
        local_execution=False, # Use remote Modal for security by default
    )
    
    # Set the system prompt separately
    
    
    return agent

async def run_matching(
    agent: TinyCodeAgent,
    participants_df: pd.DataFrame,
    organizer_criteria: str
) -> str:
    """
    Runs the matching process using the configured agent.

    Args:
        agent: The TinyCodeAgent instance.
        participants_df: A DataFrame with participant data.
        organizer_criteria: A string containing the organizer's preferences.

    Returns:
        The final markdown report of the matched teams.
    """
    # Set the participant data and criteria as variables for the agent's environment
    print(participants_df.head())
    agent.set_user_variables({
        "participants_df": participants_df,
        "organizer_criteria": organizer_criteria
    })    

    # The user prompt is simple, as the main instructions are in the system prompt
    task = dedent("""
                                  
    You are a brilliant hackathon team-matching AI.
    Your task is to form teams from a list of participants provided in a pandas DataFrame.
    You will be given the DataFrame in a variable named `participants_df`.

    Your goal is to write and execute Python code using the `run_python` tool to group these participants into balanced teams.""")


    task = organizer_prompt+'\n\n'

    task += ("Form the teams based on the provided data and criteria."
             "\n Please go through all of them and give me details of all groups. "
                   f"\n<Organizer Criteria>\n{organizer_criteria}\n</Organizer Criteria>")
    
    final_report = await agent.run(task, max_turns=15)
    print(agent.messages)
    return final_report