#!/usr/bin/env python3
"""Run inside a separate LiteLLM container exec: fake OAuth failures, real API backup."""
import asyncio
from copy import deepcopy
from functools import wraps
import json
import os
from pathlib import Path
from unittest.mock import patch

import litellm
import yaml

litellm.suppress_debug_info = True


async def main():
    config = yaml.safe_load(Path('/tmp/litellm.yaml').read_text())
    primary, backup = 'codex/gpt-luna-latest', 'openai/gpt-luna-latest'
    routes = [deepcopy(r) for r in config['model_list'] if r['model_name'] in (primary, backup)]
    assert len(routes) == 2
    assert {primary: [backup]} in config['router_settings']['fallbacks']
    for route in routes:
        if route['model_name'] == backup:
            route['litellm_params']['api_key'] = os.environ['OPENAI_API_KEY']
    actual_chat, actual_responses = litellm.acompletion, litellm.aresponses
    for protocol in ('chat', 'responses'):
        for error_type, expected in ((litellm.RateLimitError, 429),
                                     (litellm.AuthenticationError, 401),
                                     (litellm.ServiceUnavailableError, 503)):
            calls = []

            async def invoke(real, *args, **kwargs):
                model = kwargs.get('model', '')
                if model.startswith('chatgpt/'):
                    calls.append('subscription-rejected')
                    raise error_type(message='Synthetic OAuth failure; isolated test only',
                                     llm_provider='chatgpt', model=model)
                assert model.startswith('openai/')
                calls.append('official-api')
                return await real(*args, **kwargs)

            @wraps(actual_chat)
            async def chat(*args, **kwargs):
                return await invoke(actual_chat, *args, **kwargs)

            @wraps(actual_responses)
            async def responses(*args, **kwargs):
                return await invoke(actual_responses, *args, **kwargs)

            with patch.object(litellm, 'acompletion', chat), patch.object(litellm, 'aresponses', responses):
                router = litellm.Router(model_list=deepcopy(routes), fallbacks=[{primary: [backup]}],
                                        num_retries=0, timeout=60)
                if protocol == 'chat':
                    result = await router.acompletion(model=primary, messages=[
                        {'role': 'user', 'content': 'Reply only OK. Synthetic test.'}], max_tokens=32)
                    assert result.choices and result.choices[0].message.content
                else:
                    result = await router.aresponses(model=primary, input=[
                        {'role': 'user', 'content': 'Reply only OK. Synthetic test.'}],
                        max_output_tokens=32, store=False)
                    assert result.status == 'completed'
                assert calls[0] == 'subscription-rejected' and calls[-1] == 'official-api', calls
                print(json.dumps({'protocol': protocol, 'injected_error': expected,
                                  'calls': calls, 'backup': backup, 'passed': True}), flush=True)


if __name__ == '__main__':
    asyncio.run(main())