from pathlib import Path
import json
import sys
import tempfile
import time

sys.path.insert(0, r'C:\kefu\wechat_rpa')
from admin_backend import Database

with tempfile.TemporaryDirectory(prefix='model-usage-query-test-') as temporary:
    db = Database(Path(temporary) / 'usage.db')
    db.migrate()
    _, since, until = db._model_usage_window(7)
    with db.connect() as connection:
        connection.execute('BEGIN')
        connection.executemany('''INSERT INTO model_calls
            (tenant_id,task_id,created_at) VALUES (?,?,?)''',
            [('tenant-a', f'task-{i}', since) for i in range(20000)])
        connection.executemany('''INSERT INTO model_usage_events
            (event_id,task_id,tenant_id,purpose,role,attempt,status,latency_ms,created_at,total_tokens)
            VALUES (?,?,?,'chat','answer',1,'success',1,?,100)''',
            [(f'event-{i}', f'task-{i}', 'tenant-a', since) for i in range(10000)])
    queries = {
        'tenant_period': (
            'SELECT SUM(total_tokens) FROM model_usage_events WHERE tenant_id=? AND created_at>=? AND created_at<?',
            ['tenant-a', since, until]),
        'legacy_coverage': ('''SELECT COUNT(*) FROM model_calls AS calls
            WHERE datetime(calls.created_at)>=datetime(?) AND datetime(calls.created_at)<datetime(?)
            AND calls.tenant_id=? AND NOT EXISTS (
              SELECT 1 FROM model_usage_events AS usage
              WHERE calls.task_id != '' AND usage.task_id != '' AND usage.task_id=calls.task_id
              AND (usage.tenant_id=calls.tenant_id OR
                   (usage.tenant_id IN ('','default') AND calls.tenant_id IN ('','default'))))''',
            [since, until, 'tenant-a']),
    }
    report = {}
    with db.connect() as connection:
        for name, (query, parameters) in queries.items():
            plan = [dict(row) for row in connection.execute('EXPLAIN QUERY PLAN ' + query, parameters)]
            started = time.perf_counter()
            result = connection.execute(query, parameters).fetchone()[0]
            report[name] = {'plan': plan, 'result': result, 'elapsed_ms': round((time.perf_counter()-started)*1000, 2)}
    assert report['legacy_coverage']['result'] == 10000
    assert report['tenant_period']['result'] == 1000000
    Path(__file__).with_name('query-plans.json').write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding='utf-8')
    print(json.dumps(report, ensure_ascii=False))