from graph_builder import (
    build_graph_from_documents,
    save_graph,
    load_graph
)
from react_agent import react_agent  
from visualization import visualize_graph

GRAPH_PATH = "graph/kg.graphml"

# Create graph or load existing one
def get_or_build_graph(force_rebuild=False):

    if not force_rebuild:
        try:
            print("\n[INFO] Loading existing graph...")
            G = load_graph(GRAPH_PATH)

            print(f"[INFO] Loaded graph: {len(G.nodes())} nodes, {len(G.edges())} edges")
            return G

        except Exception as e:
            print(f"[WARN] Could not load graph: {e}")
            print("[INFO] Rebuilding graph...")

    print("\n[INFO] Building graph from documents...")
    G = build_graph_from_documents()

    print(f"[INFO] Built graph: {len(G.nodes())} nodes, {len(G.edges())} edges")

    print("\n[INFO] Saving graph...")
    save_graph(G, GRAPH_PATH)
    return G

# Testing questions (query)
def ask_question(G, question):

    print("\n" + "=" * 60)
    print(f"QUESTION: {question}")
    print("=" * 60)

    history = []

    result = react_agent(
        user_message=question,
        history=history,
        G=G
    )

    print("\n--- FINAL ANSWER ---")
    print(result.answer)

    print("\n--- EVIDENCE (GRAPH) ---")
    print(result.evidence)

    return result


# Visualisation
def show_graph(G):

    print("\n[INFO] Generating visualization...")

    visualize_graph(G)

    print("[INFO] Graph saved as kg.html")

# Run main
if __name__ == "__main__":

    # Create graph or load existing one
    G = get_or_build_graph(force_rebuild=False)

    # Run visualisation.py (optional and for testing)
    #show_graph(G)

    # Testing questions (query)
    test_questions = [
        #"What is taught in Informatics?",
        #"What is used in machine learning?",
        #"What Literary Texts are taught in Slovak Language And Slovak Literature?",
        "Čo sa učí na predmete Slovenský jazyk a literatúra?"
    ]

    for q in test_questions:
        ask_question(G, q)

    print("\n[INFO] PIPELINE COMPLETED")