import pprint

from app.core.rag.graph import rag_app
from app.utils import get_logger

logger = get_logger(__name__)


def test_rag_flow():
    logger.info("--- Starting RAG Flow Test ---")

    # 1. Initial constraint
    logger.info("1. User: 'I want to study'")
    state = {
        "user_input": "I want to study",
        "profile": {},
        "missing_profile_fields": [],
        "retrieved_docs": [],
    }

    result = rag_app.invoke(state)
    logger.info(f"Bot Response: {result.get('response')}")
    logger.info(f"Profile: {pprint.pformat(result.get('profile'))}")
    logger.info(f"Missing: {pprint.pformat(result.get('missing_profile_fields'))}")

    # 2. Provide Country
    logger.info("2. User: 'Germany'")
    state = {
        "user_input": "Germany",
        "profile": result.get("profile", {}),  # Pass usage profile back
        "missing_profile_fields": [],
        "retrieved_docs": [],
        "expecting_clarification": result.get("expecting_clarification", False),
    }

    result = rag_app.invoke(state)
    logger.info(f"Bot Response: {result.get('response')}")
    logger.info(f"Profile: {pprint.pformat(result.get('profile'))}")
    logger.info(f"Missing: {pprint.pformat(result.get('missing_profile_fields'))}")

    # 3. Provide Degree
    logger.info("3. User: 'Masters'")
    state = {
        "user_input": "Masters",
        "profile": result.get("profile", {}),
        "missing_profile_fields": [],
        "retrieved_docs": [],
        "expecting_clarification": result.get("expecting_clarification", False),
    }

    result = rag_app.invoke(state)
    logger.info(f"Bot Response: {result.get('response')}")
    logger.info(f"Profile: {pprint.pformat(result.get('profile'))}")
    logger.info(f"Missing: {pprint.pformat(result.get('missing_profile_fields'))}")

    # 4. Provide Field
    logger.info("3. User: 'Hotel Management'")
    state = {
        "user_input": "Hotel Management",
        "profile": result.get("profile", {}),
        "missing_profile_fields": [],
        "retrieved_docs": [],
        "expecting_clarification": result.get("expecting_clarification", False),
    }

    result = rag_app.invoke(state)
    logger.info(f"Bot Response: {result.get('response')}")
    logger.info(f"Profile: {pprint.pformat(result.get('profile'))}")
    logger.info(f"Missing: {pprint.pformat(result.get('missing_profile_fields'))}")

    # 5. Greeting
    logger.info("5. User: 'Hello'")
    state = {
        "user_input": "Hello",
        "profile": result.get("profile", {}),
        "missing_profile_fields": [],
        "retrieved_docs": [],
        "expecting_clarification": result.get("expecting_clarification", False),
    }
    result = rag_app.invoke(state)
    logger.info(f"Bot Response: {result.get('response')}")


if __name__ == "__main__":
    test_rag_flow()