"""langchain-modal-gpu-ez 라이브 통합 테스트.

실제 API 키로 LLM, Embeddings, Chain이 동작하는지 확인한다.
실행: python test_live.py
"""

from dotenv import load_dotenv

load_dotenv()


def test_llm():
    """LLM 텍스트 생성 테스트 (Modal T4 GPU)."""
    from langchain_modal_gpu_ez import ModalGpuEzLLM

    llm = ModalGpuEzLLM(model_id="distilgpt2", gpu="T4", max_new_tokens=10)
    result = llm.invoke("Hello world")

    assert isinstance(result, str)
    assert len(result) > 0
    print(f"  result: {result[:80]!r}")


def test_embeddings():
    """Embeddings 테스트 (로컬 실행)."""
    from langchain_modal_gpu_ez import ModalGpuEzEmbeddings

    emb = ModalGpuEzEmbeddings(model_id="BAAI/bge-small-en-v1.5", gpu="Local")
    vector = emb.embed_query("test sentence")

    assert isinstance(vector, list)
    assert len(vector) > 0
    assert all(isinstance(x, float) for x in vector)
    print(f"  vector dim: {len(vector)}, first 3: {vector[:3]}")


def test_chain():
    """LangChain 체인 테스트 (prompt | llm)."""
    from langchain_core.prompts import PromptTemplate

    from langchain_modal_gpu_ez import ModalGpuEzLLM

    llm = ModalGpuEzLLM(model_id="distilgpt2", gpu="T4", max_new_tokens=10)
    prompt = PromptTemplate.from_template("Tell me about {topic}:")
    chain = prompt | llm

    result = chain.invoke({"topic": "Python"})

    assert isinstance(result, str)
    assert len(result) > 0
    print(f"  result: {result[:80]!r}")


if __name__ == "__main__":
    tests = [
        ("LLM (distilgpt2, T4)", test_llm),
        ("Embeddings (bge-small, Local)", test_embeddings),
        ("Chain (prompt | llm)", test_chain),
    ]

    passed = 0
    failed = 0

    for name, fn in tests:
        print(f"\n[TEST] {name}")
        try:
            fn()
            print(f"  PASS")
            passed += 1
        except Exception as e:
            print(f"  FAIL: {e}")
            failed += 1

    print(f"\n{'='*40}")
    print(f"Results: {passed} passed, {failed} failed")
    if failed == 0:
        print("All tests passed!")
    else:
        print("Some tests failed.")
        exit(1)
