import os
import shutil
import gradio as gr
from autogen import ConversableAgent, AssistantAgent
from autogen.coding import LocalCommandLineCodeExecutor

# Setup
WORK_DIR = "coding"
CSV_PATH = os.path.join(WORK_DIR, "input.csv")
PLOT_PATH = os.path.join(WORK_DIR, "plot.png")
CODE_PATH = os.path.join(WORK_DIR, "generated_code.py")
os.makedirs(WORK_DIR, exist_ok=True)

# Agents
llm_config = {"model": "gpt-4-turbo"}

code_writer_agent = AssistantAgent(
    name="code_writer_agent",
    llm_config=llm_config,
    human_input_mode="NEVER",
    system_message=(
        "You are a coding assistant. Write Python code that reads 'input.csv' in the current directory, "
        "performs analysis based on the user prompt, and saves the resulting plot to 'plot.png'. "
        "If plot is not possible, output text or tabular results to 'results.txt'. "
        "Only generate the code, no explanations or comments."
    ),
)

executor = LocalCommandLineCodeExecutor(timeout=60, work_dir=WORK_DIR)

code_executor_agent = ConversableAgent(
    name="code_executor_agent",
    llm_config=False,
    code_execution_config={"executor": executor},
    human_input_mode="NEVER",
    default_auto_reply="TERMINATE",
)

# Main function
def process_csv_and_prompt(prompt, file):
    # Save uploaded file
    if file is not None:
        shutil.copyfile(file.name, CSV_PATH)

    # Clean up old outputs
    if os.path.exists(PLOT_PATH):
        os.remove(PLOT_PATH)
    if os.path.exists(CODE_PATH):
        os.remove(CODE_PATH)
    results_path = os.path.join(WORK_DIR, "results.txt")
    if os.path.exists(results_path):
        os.remove(results_path)

    # Run agents
    code_executor_agent.reset()
    code_writer_agent.reset()
    full_prompt = f"{prompt}\n(Remember: the file is 'input.csv' in current directory)"
    chat = code_executor_agent.initiate_chat(code_writer_agent, message=full_prompt)

    # Save generated code
    if chat and "content" in chat.messages[-1]:
        generated_code = chat.messages[-1]["content"]
        with open(CODE_PATH, "w") as f:
            f.write(generated_code)
    else:
        generated_code = "No code captured."

    # Decide outputs
    plot = PLOT_PATH if os.path.exists(PLOT_PATH) else None
    table_text = None
    if os.path.exists(results_path):
        with open(results_path, "r") as f:
            table_text = f.read()

    return plot, generated_code, table_text or "No tabular/text results produced."

# Gradio Interface
iface = gr.Interface(
    fn=process_csv_and_prompt,
    inputs=[
        gr.Textbox(label="Enter Analysis Prompt", placeholder="e.g. Generate 1000 Random numbers and plot histogram"),
        gr.File(label="Upload CSV", file_types=[".csv"]),
    ],
    outputs=[
        gr.Image(label="Generated Plot"),
        gr.Code(label="Generated Code"),
        gr.Textbox(label="Tabular/Text Results")
    ],
    title="Multi agent coding example",
    description="Upload CSV + analysis prompt. System generates Python code, executes it, and shows results as plot, code, and/or text.",
    allow_flagging="never",
)

if __name__ == "__main__":
    iface.launch()