"""Isolated Chrome WebGPU service and OpenAI bridge; no cloud/model fallback."""
import asyncio
import argparse
import importlib.util
import json
import mimetypes
from pathlib import Path
import shutil
import threading
import time
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer

from openai_server import normalize_messages, parse_tools
from playwright.async_api import async_playwright

ROOT = Path(r'D:\webbrain-training\compass-tiny-v3-spark-x25-20260926\onnx-v3')
SOURCE = Path(__file__).parent
MODEL_ID = 'webbrain-compass-tiny-v3-onnx-fp16'
PORT = 8561
VENDOR = Path(r'F:\webbrain-data\compass-tiny-v2-webgpu-recovery-v24\runtime\vendor')


async def main():
    global ROOT, PORT
    parser = argparse.ArgumentParser()
    parser.add_argument('--run-tag', default='browser-v5')
    parser.add_argument('--artifact-root', type=Path, default=ROOT)
    parser.add_argument('--package-mode', action='store_true', help='Root is the published package, not the export work directory')
    parser.add_argument('--port', type=int, default=PORT)
    args = parser.parse_args()
    ROOT, PORT = args.artifact_root.resolve(), args.port
    model_dir = ROOT if args.package_mode else ROOT / 'bundle'
    run = ROOT / args.run_tag
    run.mkdir(exist_ok=False)
    runtime = ROOT / 'runtime'
    runtime.mkdir(exist_ok=True)
    if not (runtime / 'vendor').exists():
        if not VENDOR.is_dir():
            raise RuntimeError('Missing pinned runtime/vendor files; download the complete package')
        shutil.copytree(VENDOR, runtime / 'vendor')
    dxgi_source = SOURCE / 'compass_dxgi_5090.py'
    if not dxgi_source.is_file():
        dxgi_source = Path(r'F:\webbrain-data\webbrain-data\scripts\hf\compass_dxgi_5090.py')
    module_spec = importlib.util.spec_from_file_location('dxgi', dxgi_source)
    dxgi = importlib.util.module_from_spec(module_spec)
    module_spec.loader.exec_module(dxgi)
    adapter = dxgi.discover()
    loop = asyncio.get_running_loop()
    lock = asyncio.Lock()
    state = {'ready': False, 'model': MODEL_ID, 'precision': 'FP16 ONNX; FP32 residual/norm/RoPE/softmax',
        'runtime': 'ONNX Runtime WebGPU', 'gpu': 'RTX 5090', 'cloudFallback': False,
        'maxContextTokens': 4096, 'temperature': 'greedy-only benchmark endpoint'}
    page = None

    async def rpc(method, payload):
        async with lock:
            return await page.evaluate('async ({method,payload})=>window[method](payload)', {'method': method, 'payload': payload})

    class Handler(BaseHTTPRequestHandler):
        def log_message(self, fmt, *args):
            if self.path.startswith(('/v1/', '/infer', '/template', '/health')):
                print(json.dumps({'http': fmt % args, 'time': time.time()}), flush=True)

        def json(self, status, value):
            data = json.dumps(value, ensure_ascii=False).encode()
            self.send_response(status)
            self.send_header('Content-Type', 'application/json')
            self.send_header('Content-Length', str(len(data)))
            self.end_headers()
            self.wfile.write(data)

        def do_GET(self):
            if self.path == '/health':
                return self.json(200, state)
            if self.path == '/v1/models':
                return self.json(200, {'object': 'list', 'data': [{'id': MODEL_ID, 'object': 'model', 'owned_by': 'webbrain-one', **state}]})
            if self.path == '/abi':
                return self.json(200, json.loads((ROOT / 'graph-abi.json').read_text()))
            if self.path == '/':
                data = b'<html><script type="module" src="/browser_runtime.js"></script></html>'
                self.send_response(200)
                self.send_header('Content-Type', 'text/html')
                self.send_header('Content-Length', str(len(data)))
                self.end_headers()
                return self.wfile.write(data)
            if self.path == '/browser_runtime.js':
                target = SOURCE / 'browser_runtime.js'
            elif self.path.startswith('/model/'):
                target = model_dir / self.path[len('/model/'):]
            elif self.path.startswith('/runtime/'):
                target = ROOT / self.path.lstrip('/')
            else:
                return self.json(404, {'error': 'Not found'})
            target = target.resolve()
            allowed = [ROOT.resolve(), SOURCE.resolve()]
            if not any(target.is_relative_to(p) for p in allowed) or not target.is_file():
                return self.json(404, {'error': 'Not found'})
            self.send_response(200)
            suffix = target.suffix
            self.send_header('Content-Type', {'mjs': 'text/javascript', 'js': 'text/javascript', 'wasm': 'application/wasm'}.get(suffix[1:], mimetypes.guess_type(str(target))[0] or 'application/octet-stream'))
            self.send_header('Content-Length', str(target.stat().st_size))
            self.end_headers()
            with target.open('rb') as f:
                shutil.copyfileobj(f, self.wfile, length=4 * 1024 * 1024)

        def do_POST(self):
            try:
                length = int(self.headers.get('Content-Length', '0'))
                if not 0 < length <= 4 * 1024 * 1024:
                    raise ValueError('Invalid body size')
                body = json.loads(self.rfile.read(length))
                if not state['ready']:
                    return self.json(503, {'error': 'WebGPU model not ready'})
                if self.path == '/infer':
                    method, payload = 'sparkNumerics', body
                elif self.path == '/template':
                    method, payload = 'sparkTemplate', body
                elif self.path == '/v1/chat/completions':
                    if body.get('model') != MODEL_ID or body.get('stream', False) or float(body.get('temperature', 0)) != 0:
                        raise ValueError('Use exact ONNX model id, stream=false, temperature=0')
                    maximum = int(body.get('max_completion_tokens', body.get('max_tokens', 512)))
                    if not 1 <= maximum <= 2048:
                        raise ValueError('max_tokens must be 1..2048')
                    method, payload = 'sparkGenerate', {'messages': normalize_messages(body['messages']),
                        'tools': body.get('tools'), 'maxTokens': maximum}
                else:
                    return self.json(404, {'error': 'Not found'})
                result = asyncio.run_coroutine_threadsafe(rpc(method, payload), loop).result(timeout=600)
                if self.path == '/v1/chat/completions':
                    raw = result['content']
                    parser_error = None
                    try:
                        calls, content = parse_tools(raw, body.get('tools'))
                    except Exception as exc:
                        calls, content, parser_error = [], raw, str(exc)
                    message = {'role': 'assistant', 'content': content or None}
                    finish = result['finishReason']
                    if calls:
                        message['tool_calls'] = calls
                        finish = 'tool_calls'
                    result = {'id': 'chatcmpl-' + uuid.uuid4().hex, 'object': 'chat.completion', 'created': int(time.time()),
                        'model': MODEL_ID, 'choices': [{'index': 0, 'message': message, 'finish_reason': finish}],
                        'usage': {'prompt_tokens': result['promptTokens'], 'completion_tokens': result['completionTokens'],
                                  'total_tokens': result['promptTokens'] + result['completionTokens']}, 'x_runtime': state}
                    if parser_error:
                        result['x_parser_error'] = parser_error
                self.json(200, result)
            except (ValueError, KeyError, TypeError) as exc:
                self.json(400, {'error': {'message': str(exc), 'type': 'invalid_request_error'}})
            except Exception as exc:
                self.json(500, {'error': {'message': str(exc), 'type': 'runtime_error'}})

    server = ThreadingHTTPServer(('127.0.0.1', PORT), Handler)
    threading.Thread(target=server.serve_forever, daemon=True).start()
    async with async_playwright() as p:
        browser = await p.chromium.launch(executable_path=r'C:\Program Files\Google\Chrome\Application\chrome.exe', headless=True,
            args=['--enable-unsafe-webgpu', '--ignore-gpu-blocklist', '--use-angle=d3d11',
                  f'--use-adapter-luid={adapter["high"]},{adapter["low"]}', '--force_high_performance_gpu',
                  '--enable-features=EGLDualGPURendering', '--enable-dawn-features=allow_unsafe_apis',
                  '--js-flags=--max-old-space-size=16384'])
        try:
            cdp = await browser.new_browser_cdp_session()
            info = await cdp.send('SystemInfo.getInfo')
            renderer = info['gpu']['auxAttributes']['glRenderer']
            if 'RTX 5090' not in renderer:
                raise RuntimeError('Chrome selected the wrong GPU')
            (run / 'gpu-identity.json').write_text(json.dumps({'dxgi': adapter, 'renderer': renderer}), encoding='utf-8')
            page = await browser.new_page()
            page.on('requestfailed', lambda request: print(json.dumps({'requestFailed': request.url.split('?')[0], 'reason': request.failure}), flush=True))
            page.on('console', lambda message: print(json.dumps({'console': message.type, 'text': message.text[:1500]}), flush=True))
            page.on('pageerror', lambda error: print(json.dumps({'pageError': str(error)}), flush=True))
            await page.goto(f'http://127.0.0.1:{PORT}/')
            await page.wait_for_function('typeof window.initializeSpark === "function"')
            ready = await page.evaluate('window.initializeSpark()')
            state.update(ready=True, browser=browser.version, **ready)
            (run / 'ready.json').write_text(json.dumps(state), encoding='utf-8')
            print(json.dumps({'ready': True, 'baseUrl': f'http://127.0.0.1:{PORT}/v1', **state}), flush=True)
            await asyncio.Event().wait()
        finally:
            server.shutdown()
            await browser.close()


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