from langgraph.graph import END, START, StateGraph

from langchain_core.tracers.context import tracing_v2_enabled

from prompts import *
import gradio as gr


from utils import (
    load_github_codebase,
    router,
    get_plan_for_codebase,
    parse_plan,
    explore_file,
    final_mermaid_code_generation,
    extract_mermaid_and_generate_graph,
    GraphState,
)


def visualize_github_repo(repo_name, repo_branch):
    yield ("Looking at the Repo!", None)
    documents = load_github_codebase(repo_name, repo_branch)
    yield ("Repo loaded!", None)

    # Define a new graph
    workflow = StateGraph(GraphState)

    # Define the two nodes we will cycle
    workflow.add_node("planner", get_plan_for_codebase)
    workflow.add_node("parse_plan", parse_plan)

    workflow.add_node("explore_file", explore_file)
    workflow.add_node("router", router)

    workflow.add_node("generate_mermaid_code", final_mermaid_code_generation)
    workflow.add_node("render_mermaid", extract_mermaid_and_generate_graph)

    # Set the entrypoint as `agent`
    # This means that this node is the first one called
    workflow.add_edge(START, "planner")
    workflow.add_edge("planner", "parse_plan")

    # We now add a conditional edge
    workflow.add_conditional_edges(
        # First, we define the start node. We use `agent`.
        # This means these are the edges taken after the `agent` node is called.
        "parse_plan",
        # Next, we pass in the function that will determine which node is called next.
        router,
    )
    workflow.add_conditional_edges(
        "explore_file",
        router,
    )
    workflow.add_edge("generate_mermaid_code", "render_mermaid")
    workflow.add_edge("render_mermaid", END)

    # Finally, we compile it!
    # This compiles it into a LangChain Runnable,
    # meaning you can use it as you would any other runnable.
    # Note that we're (optionally) passing the memory when compiling the graph
    app = workflow.compile()

    with tracing_v2_enabled():
        for s in app.stream(
            {"messages": [], "documents": documents},
            {"recursion_limit": 100},
        ):
            if "__end__" not in s:
                print(s)
                print("----")
                if "planner" in s:
                    yield ("Planning the Exploration !", None)
                if "parse_plan" in s:
                    yield ("Planning done! Parsing the plan!", None)
                if "explore_file" in s:
                    yield (
                        f"Exploration started! Exploring file: {s['explore_file']['explored_files'][-1]} !",
                        None,
                    )
                if "generate_mermaid_code" in s:
                    yield (
                        "Exploration done! Gathering thoughts and generating a graph!",
                        None,
                    )
                if "render_mermaid" in s:

                    yield (
                        s["render_mermaid"]["messages"][-1].content,
                        s["render_mermaid"]["final_graph"],
                    )


demo = gr.Interface(
    fn=visualize_github_repo,
    inputs=[
        gr.Textbox(
            label="Repo Name",  # Input component label
            value="abhishekkrthakur/autoxgb",  # Input component placeholder
            placeholder="Name of the Repo in format author/repo",  # Input component description
        ),
        gr.Textbox(
            label="Repo Branch",  # Input component label
            value="main",  # Input component placeholder
            placeholder="Branch to explore",  # Input component description
        ),
    ],
    outputs=[
        gr.Textbox(
            label="Mermaid Graph",  # Output component label
            placeholder="Visualization of the functionalities of the Repo",  # Output component description
        ),
        gr.Image(
            label="Rendered Graph",  # Output component label
            placeholder="Rendered Graph",  # Output component description
        ),
    ],
    title="Repo Functionality Visualizer",
)

demo.launch()