from langgraph.graph import StateGraph, END

from state import AgentState

from agents.routing_agent import RouterAgent
from agents.image_agent import ImageAgent
from agents.memory import Memory
from agents.planner import PlannerAgent


# -------------------------------
# Dummy LLM
# -------------------------------

class DummyLLM:

    def invoke(self, prompt):

        class Response:
            content = "generate"

        return Response()


llm = DummyLLM()

memory = Memory()

router = RouterAgent(llm)

planner = PlannerAgent()

agent = ImageAgent(router, memory, planner)


# -------------------------------
# ROUTE
# -------------------------------

def route_node(state):

    task = router.route(state["user_input"])

    state["task"] = task

    return state


# -------------------------------
# PLAN
# -------------------------------

def plan_node(state):

    state["steps"] = planner.plan(state["task"])

    return state


# -------------------------------
# LOAD UPLOADED IMAGE
# -------------------------------

def upload_node(state):

    uploaded = state.get("uploaded_files", [])

    if uploaded:

        # Store uploaded image in memory
        memory.uploaded_files = uploaded
        memory.last_used_image = uploaded[0]

    return state


# -------------------------------
# AUTO ANALYZE
# -------------------------------

def analyze_node(state):

    image = memory.last_used_image

    if image:

        analysis = agent.auto_analyze(image)

        state["analysis"] = analysis

    return state


# -------------------------------
# EXECUTE
# -------------------------------

def execute_node(state):

    result = agent.run(state["user_input"])

    state["result"] = result

    return state


# -------------------------------
# GRAPH
# -------------------------------

workflow = StateGraph(AgentState)

workflow.add_node("upload", upload_node)

workflow.add_node("route", route_node)

workflow.add_node("plan", plan_node)

workflow.add_node("analyze", analyze_node)

workflow.add_node("execute", execute_node)


workflow.set_entry_point("upload")

workflow.add_edge("upload", "route")

workflow.add_edge("route", "plan")

workflow.add_edge("plan", "analyze")

workflow.add_edge("analyze", "execute")

workflow.add_edge("execute", END)

app = workflow.compile()


# -------------------------------
# MAIN
# -------------------------------

def run_graph(user_input, uploaded_files=None):

    state = {

        "user_input": user_input,

        "task": "",

        "steps": [],

        "uploaded_files": uploaded_files or [],

        "analysis": None,

        "result": None,

        "history": []
    }

    return app.invoke(state)