import os
from prompts import *
import gradio as gr


def visualize_github_repo(repo_name, repo_branch, mistral_api_key):
    os.environ['MISTRAL_API_KEY'] = mistral_api_key

    from langgraph.graph import END, START, StateGraph
    from langchain_core.tracers.context import tracing_v2_enabled
    
    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,
    )
    
    yield "Looking at the Repo!"
    documents = load_github_codebase(repo_name, repo_branch)
    yield "Repo loaded!"

    # 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]} !"

                if "generate_mermaid_code" in s:
                    yield "Exploration done! Gathering thoughts and generating a graph!",
                if "render_mermaid" in s:

                    yield s["render_mermaid"]["messages"][-1].content


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 Public 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
        ),
        gr.Textbox(
            label="Mistral API Key",  # Input component label
            # value="main",  # Input component placeholder
            placeholder="Mistral API Key from mistral.ai",  # Input component description
        ),
    ],
    outputs = gr.Textbox(label="Mermaid Graph", placeholder="Visualization of the functionalities of the Repo",), # Output component description
    title="Repo Functionality Visualizer",
)

demo.launch()