import pytest
import httpx

# The base URL for the local server
BASE_URL = "http://127.0.0.1:8000"
TEST_PROMPT = "Explain the theory of relativity in one sentence."

async def test_root_endpoint():
    """
    Tests if the root endpoint is alive and returns the correct status message.
    """
    print("\n--- Running test_root_endpoint ---")
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{BASE_URL}/")
        assert response.status_code == 200
        assert response.json() == {"status": "AI Model Hub is running!"}
    print("✅ Root endpoint test passed.")

async def test_chat_endpoint_deepseek():
    """
    Tests the /chat endpoint using the 'deepseek' model in the request body.
    """
    print("\n--- Running test_chat_endpoint_deepseek ---")
    # FIX: URL no longer has query parameters
    url = f"{BASE_URL}/chat"
    # FIX: 'model' is now part of the JSON payload
    payload = {"prompt": TEST_PROMPT, "model": "deepseek"}

    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(url, json=payload)

    assert response.status_code == 200, f"Expected 200, got {response.status_code}. Response: {response.text}"
    data = response.json()
    assert "answer" in data
    assert data["model_used"] == "deepseek"
    print(f"✅ DeepSeek chat test passed. Response snippet: {data['answer'][:80]}...")

async def test_chat_endpoint_gemini():
    """
    Tests the /chat endpoint using the 'gemini' model in the request body.
    """
    print("\n--- Running test_chat_endpoint_gemini ---")
    # FIX: URL no longer has query parameters
    url = f"{BASE_URL}/chat"
    # FIX: 'model' is now part of the JSON payload
    payload = {"prompt": TEST_PROMPT, "model": "gemini"}

    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(url, json=payload)

    assert response.status_code == 200, f"Expected 200, got {response.status_code}. Response: {response.text}"
    data = response.json()
    assert "answer" in data
    assert data["model_used"] == "gemini"
    print(f"✅ Gemini chat test passed. Response snippet: {data['answer'][:80]}...")

async def test_chat_with_empty_prompt():
    """
    Tests error handling for an empty prompt. Expects a 422 error.
    """
    print("\n--- Running test_chat_with_empty_prompt ---")
    url = f"{BASE_URL}/chat"
    # FIX: Payload needs a 'model' to correctly test the 'prompt' validation
    payload = {"prompt": "", "model": "deepseek"}

    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(url, json=payload)

    assert response.status_code == 422
    assert "string_too_short" in response.json()["detail"][0]["type"]
    print("✅ Empty prompt test passed.")

async def test_unsupported_model():
    """
    Tests error handling for an invalid model name. Expects a 422 error.
    """
    print("\n--- Running test_unsupported_model ---")
    url = f"{BASE_URL}/chat"
    # FIX: Send the unsupported model in the payload to trigger the correct validation
    payload = {"prompt": TEST_PROMPT, "model": "unsupported_model_123"}

    async with httpx.AsyncClient() as client:
        response = await client.post(url, json=payload)
    
    assert response.status_code == 422
    # This assertion will now pass because the correct validation error is triggered
    assert "Input should be 'deepseek' or 'gemini'" in response.json()["detail"][0]["msg"]
    print("✅ Unsupported model test passed.")

async def test_add_document_success():
    """
    Tests the /document endpoint for successful document ingestion.
    """
    print("\n--- Running test_add_document_success ---")
    url = f"{BASE_URL}/document"
    doc_data = {
        "title": "Test Integration Document",
        "text": "This document is for testing the integration endpoint.",
        "source_url": "http://example.com/integration_test"
    }
    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(url, json=doc_data)

    assert response.status_code == 200
    assert "Document 'Test Integration Document' added successfully" in response.json()["message"]
    print("✅ Add document success test passed.")

async def test_add_document_invalid_data():
    """
    Tests the /document endpoint's error handling for missing required fields.
    """
    print("\n--- Running test_add_document_invalid_data ---")
    url = f"{BASE_URL}/document"
    doc_data = {
        "text": "This document is missing a title.",
        "source_url": "http://example.com/invalid_data"
    }
    async with httpx.AsyncClient() as client:
        response = await client.post(url, json=doc_data)

    assert response.status_code == 422
    assert "field required" in response.json()["detail"][0]["msg"].lower()
    print("✅ Add document with invalid data test passed.")