from typing import TypedDict, Optional, Any
from langchain_openai import ChatOpenAI
from pyspark.sql import SparkSession

# --- Configuration ---
# Initialize Spark (The Muscle)
spark = SparkSession.builder \
    .appName("AgenticETLDemo") \
    .getOrCreate()

# Initialize the LLM (The Brain)
# Note: Make sure OPENAI_API_KEY is in your environment variables
llm = ChatOpenAI(model="gpt-5.2-pro", temperature=0)

# --- The State Definition ---
class ETLState(TypedDict):
    """The shared memory for our agentic workflow."""
    goal: str               # What are we trying to do? (e.g., "Load bad_dates.csv")
    schema_hint: str        # Information about the data structure
    pyspark_code: Optional[str] # The code generated by the LLM
    execution_result: Any   # The output data or success message
    error_log: Optional[str]# Track failures to learn from them
    attempts: int           # To prevent infinite loops

def code_generator_node(state: ETLState):
    """Node 1: The Brain writes the code."""
    print(f"\n--- [Brain] Generating Code (Attempt {state['attempts'] + 1}) ---")

    # Base prompt
    prompt = f"""
    You are an expert PySpark data engineer.
    Your Goal: {state['goal']}
    Schema Hint: {state['schema_hint']}

    Write a complete, runnable block of PySpark code that accomplishes this goal.
    The code must result in a Spark DataFrame named 'final_df'.
    ONLY return the executable Python code. No markdown formatting.
    """

    # If we failed before, add the error to the prompt so the LLM can fix it.
    if state['error_log']:
        prompt += f"\n\nIMPORTANT: The previous code failed with this error:\n{state['error_log']}\nFix the code to handle this error."

    # Call the LLM
    response = llm.invoke(prompt)
    cleaned_code = response.content.replace("python", "").replace("", "").strip()

    return {
        "pyspark_code": cleaned_code,
        "attempts": state['attempts'] + 1
    }

import pandas as pd
from io import StringIO

def execution_node(state: ETLState):
    """Node 2: The Muscle runs the code."""
    print("\n--- [Muscle] Executing Spark Code ---")
    try:
        # Create dummy data for this demo to simulate a messy CSV
        csv_data = """id,date_string,value
                        1,2023-01-15,100
                        2,15/01/2023,200
                        3,Jan 15 2023,300"""

        # We define the execution context, giving the code access to 'spark' and our data
        local_vars = {"spark": spark, "csv_data": csv_data}

        # DANGER ZONE: Executing generated code.
        # The generated code should create a 'final_df' variable.
        exec(state['pyspark_code'], globals(), local_vars)

        result_df = local_vars.get("final_df")

        if result_df is None:
            raise ValueError("Code ran but did not generate 'final_df'")

        print("SUCCESS: Data transformed successfully!")
        result_df.show()

        return {"execution_result": "Success", "error_log": None}

    except Exception as e:
        print(f"FAILURE: Spark job failed with error: {e}")
        # Pass the error back to the state so the generator can fix it
        return {"error_log": str(e), "execution_result": None}

from langgraph.graph import StateGraph, END

# Initialize the graph
workflow = StateGraph(ETLState)

# Add nodes
workflow.add_node("generator", code_generator_node)
workflow.add_node("executor", execution_node)

# Define flow
workflow.set_entry_point("generator")
workflow.add_edge("generator", "executor")

# Define the conditional router
def router(state: ETLState):
    # Stop after 3 attempts to prevent infinite loops costing you API money
    if state["attempts"] >= 3:
        print("--- Max attempts reached. Stopping. ---")
        return END

    if state["error_log"]:
        # If there is an error, go back to the drawing board
        return "generator"
    else:
        # If success, we are done
        return END

workflow.add_conditional_edges("executor", router)

# Compile the graph
app = workflow.compile()

# Initial State definition
initial_state = {
    "goal": "Read the raw CSV data provided in the variable 'csv_data'. It contains a column 'date_string' with mixed date formats. Standardize this column into a proper Spark DateType column named 'clean_date' and return the dataframe as 'final_df'.",
    "schema_hint": "CSV with headers: id,date_string,value",
    "error_log": None,
    "attempts": 0,
    "pyspark_code": None,
    "execution_result": None
}

print("Starting Agentic ETL Workflow...")
# Run the graph
final_state = app.invoke(initial_state)

print("\n--- Workflow Finished ---")
if final_state["execution_result"] == "Success":
    print("Final Status: SUCCESS")
    # In a real app, you'd now write this data to your data lake (Delta Lake, Snowflake, etc.)
else:
    print("Final Status: FAILED after max attempts.")