import asyncio
import logging
from agents.graph import build_oncoagent_graph
from agents.state import AgentState

logging.basicConfig(level=logging.INFO)

async def main():
    graph = build_oncoagent_graph()
    state = {
        "clinical_text": "55-year-old female patient presents with postmenopausal bleeding. Ultrasound shows an endometrial thickening of 12mm. The endometrial biopsy report confirms Grade 1 endometrioid adenocarcinoma. What is the recommended treatment?",
        "session_id": "test_123"
    }
    
    # Debug: Check raw retrieval distances
    from rag_engine.retriever import OncoRAGRetriever
    retriever = OncoRAGRetriever()
    candidates, distances = retriever._bi_encoder_retrieve(state["clinical_text"], 5)
    print("\n--- Raw Retrieval Debug ---")
    for cand, dist in zip(candidates, distances):
        print(f"Dist: {dist:.4f} | Header: {cand['header']}")
    print("--- End Raw Retrieval Debug ---\n")
    
    config = {"configurable": {"thread_id": "test_123"}}
    async for event in graph.astream(state, config=config):
        for node_name, output in event.items():
            print(f"\n--- Node: {node_name} ---")
            if "rag_context" in output:
                print(f"RAG Context length: {len(output['rag_context'])}")
            if "rag_retrieval_count" in output:
                print(f"Retrieved: {output['rag_retrieval_count']} | Graded: {output.get('rag_grading_pass_count')}")
            if "clinical_recommendation" in output:
                print(f"Recommendation: {output['clinical_recommendation'][:200]}...")


if __name__ == "__main__":
    asyncio.run(main())