"""
Agent Q3 [Evo] — LangGraph StateGraph
Self-improvement loop: ingest → train → benchmark → feedback → repeat
"""
from langgraph.graph import StateGraph, END
from typing import TypedDict

class EvoState(TypedDict):
    cycle: int
    ingest_done: bool
    train_done: bool
    benchmark_score: float
    feedback_pushed: bool
    notes: str

def ingest_node(state: EvoState) -> EvoState:
    from arxiv_ingestor import ingest
    ingest(max_results=10)
    state["ingest_done"] = True
    state["notes"] += " | arXiv ingested"
    return state

def train_node(state: EvoState) -> EvoState:
    from training_pipeline import run
    run()
    state["train_done"] = True
    state["notes"] += " | LoRA trained"
    return state

def benchmark_node(state: EvoState) -> EvoState:
    import asyncio
    from benchmark_runner import run_all
    results = asyncio.run(run_all())
    avg = sum(r["score"] for r in results) / len(results)
    state["benchmark_score"] = avg
    state["notes"] += f" | benchmark={avg:.2%}"
    return state

def feedback_node(state: EvoState) -> EvoState:
    from feedback_collector import push_to_hf
    push_to_hf()
    state["feedback_pushed"] = True
    state["notes"] += " | feedback pushed"
    return state

def should_continue(state: EvoState) -> str:
    return "continue" if state["benchmark_score"] < 0.80 else END

def build_evo_graph() -> StateGraph:
    g = StateGraph(EvoState)
    g.add_node("ingest",    ingest_node)
    g.add_node("train",     train_node)
    g.add_node("benchmark", benchmark_node)
    g.add_node("feedback",  feedback_node)
    g.set_entry_point("ingest")
    g.add_edge("ingest",    "train")
    g.add_edge("train",     "benchmark")
    g.add_edge("benchmark", "feedback")
    g.add_conditional_edges("feedback", should_continue, {"continue":"ingest", END: END})
    return g.compile()

graph = build_evo_graph()