import asyncio
import json
import sys
from unittest import mock

sys.path.insert(0, r'C:\kefu\wechat_rpa')
import httpx
import model_gateway as gateway
from model_usage import UsageRecorder
from test_model_gateway import _outlet, _openai_body


async def run():
    database = mock.Mock()
    saved = []
    database.record_model_usage_batch.side_effect = lambda records: saved.extend(dict(item) for item in records)
    recorder = UsageRecorder(database, tenant_id='isolated-audit')
    first, second = _outlet(), _outlet()
    first.config['id'], second.config['id'] = 'malformed', 'slow-valid'

    async def handler(request):
        text = json.loads(request.content)['messages'][0]['content']
        if text == 'first':
            return httpx.Response(200, json={
                'choices': [{'message': {'content': '', 'tool_calls': 1}}],
                'usage': {'total_tokens': 10}})
        await asyncio.sleep(.1)
        return httpx.Response(200, json={**_openai_body('answer'), 'usage': {'total_tokens': 20}})

    error = None
    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client:
        try:
            async with recorder:
                await asyncio.gather(
                    gateway.call_outlet(client, first, [{'role': 'user', 'content': 'first'}], usage_recorder=recorder),
                    gateway.call_outlet(client, second, [{'role': 'user', 'content': 'second'}], usage_recorder=recorder))
        except Exception as exc:
            error = type(exc).__name__
        await asyncio.sleep(.2)
    result = {'raised': error, 'saved': [{'provider_id': e['provider_id'], 'total_tokens': e['total_tokens']} for e in saved],
              'unflushed': [{'provider_id': e['provider_id'], 'total_tokens': e['total_tokens']} for e in recorder.events]}
    print(json.dumps(result))


asyncio.run(run())