from langgraph.graph import StateGraph,START,END
from typing import TypedDict
from langchain_openai import ChatOpenAI
import os

OPENAI_API_KEY=os.environ.get("OPENAI_API_KEY")

llm=ChatOpenAI(
    model="gpt-4o-mini",
    api_key=OPENAI_API_KEY
)

class State(TypedDict):
    text:str
    french:str
    spanish:str
    japanese:str
    combined_output:str

def translate_french(state:State)->dict:
    response=llm.invoke(f"Translate the following to French:\n\n{state['text']}")
    return {"french":response.content.strip()}

def translate_spanish(state:State)->dict:
    response=llm.invoke(f"Translate the following to Spanish:\n\n{state['text']}")
    return {"spanish":response.content.strip()}

def translate_japanese(state:State)->dict:
    response=llm.invoke(f"Translate the following to Japanese:\n\n{state['text']}")
    return {"japanese":response.content.strip()}

def aggregator(state:State)->dict:
    combined=f"Original Text: {state['text']}\n\n"
    combined+=f"French: {state['french']}\n\n"
    combined+=f"Spanish: {state['spanish']}\n\n"
    combined+=f"Japanese: {state['japanese']}\n"
    return {"combined_output":combined}

graph=StateGraph(State)
graph.add_node("translate_french",translate_french)
graph.add_node("translate_spanish",translate_spanish)
graph.add_node("translate_japanese",translate_japanese)
graph.add_node("aggregator",aggregator)

graph.add_edge(START,"translate_french")
graph.add_edge(START,"translate_spanish")
graph.add_edge(START,"translate_japanese")

graph.add_edge("translate_french","aggregator")
graph.add_edge("translate_spanish","aggregator")
graph.add_edge("translate_japanese","aggregator")

graph.add_edge("aggregator",END)

app=graph.compile()

def run_workflow(text):
    state:State={
        "text":text,
        "french":"",
        "spanish":"",
        "japanese":"",
        "combined_output":""
    }
    response=app.invoke(state)
    workflow_graph=app.get_graph().draw_mermaid_png()

    return response,workflow_graph