# What this tests ?
## Tests /config/update + Test /chat/completions -> assert logs are sent to Langfuse

import pytest
import asyncio
import aiohttp
import os
import dotenv
from dotenv import load_dotenv
import pytest

load_dotenv()


async def config_update(session):
    url = "http://0.0.0.0:4000/config/update"
    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
    data = {
        "litellm_settings": {
            "success_callback": ["langfuse"],
        },
        "environment_variables": {
            "LANGFUSE_PUBLIC_KEY": os.environ["LANGFUSE_PUBLIC_KEY"],
            "LANGFUSE_SECRET_KEY": os.environ["LANGFUSE_SECRET_KEY"],
        },
    }

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        print(response_text)
        print()

        if status != 200:
            raise Exception(f"Request did not return a 200 status code: {status}")
        return await response.json()


async def chat_completion(session, key, model="azure-gpt-3.5", request_metadata=None):
    url = "http://0.0.0.0:4000/chat/completions"
    headers = {
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
    }
    data = {
        "model": model,
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "Hello!"},
        ],
        "metadata": request_metadata,
    }

    print("data sent in test=", data)

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        print(response_text)
        print()

        if status != 200:
            raise Exception(f"Request did not return a 200 status code: {status}")


@pytest.mark.asyncio
@pytest.mark.flaky(retries=6, delay=1)
async def test_team_logging():
    """
    1. Add Langfuse as a callback with /config/update
    2. Call /chat/completions
    3. Assert the logs are sent to Langfuse
    """
    try:
        async with aiohttp.ClientSession() as session:

            # Add Langfuse as a callback with /config/update
            await config_update(session)

            # 2. Call /chat/completions with a specific trace id
            import uuid

            _trace_id = f"trace-{uuid.uuid4()}"
            _request_metadata = {
                "trace_id": _trace_id,
            }

            await chat_completion(
                session,
                key="sk-1234",
                model="fake-openai-endpoint",
                request_metadata=_request_metadata,
            )

            # Test - if the logs were sent to the correct team on langfuse
            import langfuse

            langfuse_client = langfuse.Langfuse(
                public_key=os.getenv("LANGFUSE_PUBLIC_KEY"),
                secret_key=os.getenv("LANGFUSE_SECRET_KEY"),
            )

            await asyncio.sleep(10)

            print(f"searching for trace_id={_trace_id} on langfuse")

            generations = langfuse_client.get_generations(trace_id=_trace_id).data

            # 1 generation with this trace id
            assert len(generations) == 1

    except Exception as e:
        pytest.fail("Team 2 logging failed: " + str(e))