import pytest
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from src.core.state import AgentState
from src.core.graph import (
    detect_emergency,
    role_classifier_node,
    emergency_response_node,
    recovery_loop_node,
    route_after_role,
    route_patient_llm,
    route_after_validator,
    route_after_safety,
    route_after_intent,
    medical_pipeline
)
from src.core.graph_cdm import (
    data_fetch_node,
    trend_analyzer_node,
    cdm_pipeline
)

def test_detect_emergency():
    # Healthy case
    state = {"messages": [HumanMessage(content="Hello doctor, I have a normal question about my diet.")]}
    assert not detect_emergency(state)

    # Emergency cases
    emergencies = [
        "My grandfather just passed out and is unconscious!",
        "He is having a seizure right now.",
        "I think I am having a stroke or heart attack, I have chest pain."
    ]
    for msg in emergencies:
        state = {"messages": [HumanMessage(content=msg)]}
        assert detect_emergency(state) is True

@pytest.mark.asyncio
async def test_role_classifier_node():
    # User message indicating caregiver
    state = {
        "messages": [HumanMessage(content="I am caring for my mother who has diabetes.")],
        "user_role": "unknown",
        "intent_type": "unknown"
    }
    res = await role_classifier_node(state)
    assert res["user_role"] == "caregiver"
    assert "logs" in res

@pytest.mark.asyncio
async def test_emergency_response_node():
    state = {"messages": [HumanMessage(content="Heart attack!")], "logs": []}
    res = await emergency_response_node(state)
    assert res["is_valid"] is True
    assert res["is_safe"] is True
    assert len(res["messages"]) == 1
    assert "medical emergency" in res["messages"][0].content

@pytest.mark.asyncio
async def test_recovery_loop_node():
    state = {"attempts": 1, "logs": []}
    res = await recovery_loop_node(state)
    assert res["attempts"] == 2
    assert "[RECOVERY]" in res["messages"][0].content

def test_routing_logic():
    # route_after_role
    state = {"user_role": "patient", "messages": [HumanMessage(content="Normal checkup")]}
    assert route_after_role(state) == "patient_llm"

    state_emerg = {"user_role": "patient", "messages": [HumanMessage(content="passed out")]}
    assert route_after_role(state_emerg) == "emergency_response"

    state_clin = {"user_role": "clinician"}
    assert route_after_role(state_clin) == "intent_classifier"

    # route_patient_llm
    state_tools = {"messages": [AIMessage(content="Need tools", tool_calls=[{"name": "web_search", "args": {}, "id": "tc1"}])]}
    assert route_patient_llm(state_tools) == "tools_node"

    state_validator = {"messages": [AIMessage(content="Normal response")]}
    assert route_patient_llm(state_validator) == "validator"

    # route_after_validator
    assert route_after_validator({"is_valid": True}) == "safety_check"
    assert route_after_validator({"is_valid": False}) == "recovery_loop"

    # route_after_safety
    assert route_after_safety({"is_safe": True}) == "persistence_node"
    assert route_after_safety({"is_safe": False}) == "recovery_loop"

    # route_after_intent
    assert route_after_intent({"intent_type": "diagnosis"}) == "diagnosis_assist"
    assert route_after_intent({"intent_type": "treatment"}) == "treatment_assist"
    assert route_after_intent({"intent_type": "monitoring"}) == "monitoring_assist"
    assert route_after_intent({"intent_type": "other"}) == "general_assist"

@pytest.mark.asyncio
async def test_medical_pipeline_e2e(clean_db):
    patient_id = "patient-graph-test"
    clean_db._ensure_patient_exists(patient_id, clean_db)
    
    initial_state = {
        "messages": [HumanMessage(content="Hello assistant, I need general help.")],
        "user_role": "unknown",
        "intent_type": "unknown",
        "session_id": "session-graph-1",
        "patient_id": patient_id,
        "is_valid": False,
        "is_safe": False,
        "attempts": 0,
        "clinician_outputs": [],
        "patient_response": "",
        "research_output": "",
        "sources": [],
        "logs": [],
        "metrics": [],
    }

    final_state = await medical_pipeline.ainvoke(
        initial_state,
        config={"recursion_limit": 50}
    )
    assert final_state["user_role"] == "patient"
    assert final_state["is_valid"] is True
    assert final_state["is_safe"] is True
    assert len(final_state["messages"]) > 1

@pytest.mark.asyncio
async def test_cdm_pipeline_e2e(clean_db):
    patient_id = "patient-cdm-graph"
    # Seed data
    clean_db._ensure_patient_exists(patient_id, clean_db)
    clean_db._insert("observations", {
        "id": "obs-1",
        "patient_id": patient_id,
        "resource": {
            "resourceType": "Observation",
            "code": {"coding": [{"display": "Glucose"}]},
            "valueQuantity": {"value": 150}
        },
        "effectiveDateTime": "2026-06-05T12:00:00Z"
    })
    
    initial_state = {
        "messages": [HumanMessage(content="Let's analyze trends.")],
        "user_role": "patient",
        "patient_id": patient_id,
        "session_id": "session-cdm-1",
        "is_valid": False,
        "is_safe": False,
        "attempts": 0,
        "clinician_outputs": [],
        "patient_response": "",
        "research_output": "",
        "sources": [],
        "logs": [],
        "metrics": [],
    }

    final_state = await cdm_pipeline.ainvoke(
        initial_state,
        config={"recursion_limit": 50}
    )
    assert "trend_analysis" in final_state
    assert final_state["is_valid"] is True
    assert final_state["is_safe"] is True