"""Loopback OpenAI chat adapter. Native Spark tool calls are parsed, not repaired."""
import argparse
import copy
import json
import re
import threading
import time
import uuid
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

ROOT = Path(r'D:\webbrain-training\compass-tiny-v3-spark-x25-20260926')
MODEL_ID = 'webbrain-compass-tiny-v3-bf16'
PACKAGE = ROOT / 'release-compat-v3'


def parse_tools(text, tools):
    """Only lossless native-format/schema decoding; errors stay visible."""
    definitions = {t['function']['name']: t['function'] for t in tools or []}
    calls = []
    pattern = re.compile(r'<tool_call>(.*?)</tool_call>', re.S)
    spans = list(pattern.finditer(text))
    if text.count('<tool_call>') != len(spans) or text.count('</tool_call>') != len(spans):
        raise ValueError('Incomplete native tool call')
    for span in spans:
        body = span.group(1)
        name = body.split('<arg_key>', 1)[0].strip()
        if name not in definitions:
            raise ValueError(f'Unknown model-emitted function: {name}')
        tail = body[len(body.split('<arg_key>', 1)[0]):]
        arg_pattern = re.compile(r'<arg_key>(.*?)</arg_key><arg_value>(.*?)</arg_value>', re.S)
        args = {}
        cursor = 0
        schema = definitions[name].get('parameters', {})
        for arg in arg_pattern.finditer(tail):
            if tail[cursor:arg.start()].strip():
                raise ValueError('Malformed native argument sequence')
            key, raw = arg.group(1), arg.group(2)
            if key in args:
                raise ValueError('Duplicate native argument key')
            prop = schema.get('properties', {}).get(key, {})
            args[key] = raw if prop.get('type') == 'string' else json.loads(raw)
            cursor = arg.end()
        if tail[cursor:].strip() or any(k not in args for k in schema.get('required', [])):
            raise ValueError('Malformed or incomplete native arguments')
        calls.append({'id': 'call_' + uuid.uuid4().hex[:16], 'type': 'function',
                      'function': {'name': name, 'arguments': json.dumps(args, ensure_ascii=False)}})
    return calls, pattern.sub('', text).strip()


def normalize_messages(messages):
    messages = copy.deepcopy(messages)
    for message in messages:
        for tool in message.get('tool_calls', []):
            arguments = tool['function']['arguments']
            if isinstance(arguments, str):
                tool['function']['arguments'] = json.loads(arguments)
            if not isinstance(tool['function']['arguments'], dict):
                raise ValueError('Tool arguments must decode to an object')
    return messages


class Bf16Backend:
    def __init__(self):
        import os
        import sys
        source = Path(__file__).parents[1] / 'spark_x25_17b_webbrain_compass_v3'
        sys.path.insert(0, str(source))
        from train import gpu_guard
        cfg = json.loads((source / 'train-config.json').read_text())
        gpu_guard(cfg)
        import torch
        import runtime
        from transformers import AutoTokenizer
        torch.set_num_threads(8)
        self.torch = torch
        self.tokenizer = AutoTokenizer.from_pretrained(PACKAGE, local_files_only=True)
        self.model = runtime.load_native(str(PACKAGE), dtype=torch.bfloat16, device_map={'': 0}).eval()
        self.lock = threading.Lock()
        uploaded = json.loads((ROOT / 'upload-verified.json').read_text())
        self.info = {'id': MODEL_ID, 'precision': 'BF16', 'runtime': 'Transformers/PyTorch CUDA',
                     'hardware': 'NVIDIA GeForce RTX 5090', 'revision': uploaded['revision'],
                     'max_context_tokens': 16384, 'thinking_default': False,
                     'onnx': False, 'webgpu': False, 'cloud_fallback': False}

    def generate(self, request):
        if request.get('model') != MODEL_ID:
            raise ValueError(f'Use exact model id {MODEL_ID}')
        tools = request.get('tools')
        messages = normalize_messages(request['messages'])
        options = request.get('chat_template_kwargs', {})
        if options.get('enable_thinking', False):
            raise ValueError('This benchmark endpoint serves thinking-disabled Compass')
        inputs = self.tokenizer.apply_chat_template(messages, tools=tools, enable_thinking=False,
            add_generation_prompt=True, return_tensors='pt', return_dict=True).to('cuda')
        tokens = int(inputs.input_ids.shape[1])
        maximum = int(request.get('max_completion_tokens', request.get('max_tokens', 512)))
        if maximum < 1 or maximum > 4096 or tokens + maximum > self.info['max_context_tokens']:
            raise ValueError('Requested input/output exceeds the deployed 16K context limit')
        temperature = float(request.get('temperature', 0))
        if temperature < 0:
            raise ValueError('Negative temperature')
        kwargs = {'max_new_tokens': maximum, 'do_sample': temperature > 0,
                  'pad_token_id': self.tokenizer.pad_token_id, 'use_cache': True}
        if temperature > 0:
            kwargs.update(temperature=temperature, top_p=float(request.get('top_p', 1.0)))
        with self.lock, self.torch.no_grad():
            output = self.model.generate(**inputs, **kwargs)
        tail = output[0, tokens:]
        raw = self.tokenizer.decode(tail, skip_special_tokens=True)
        finish = 'length' if len(tail) == maximum else 'stop'
        stop = request.get('stop') or []
        if isinstance(stop, str):
            stop = [stop]
        for marker in stop:
            if marker in raw:
                raw = raw.split(marker, 1)[0]
                finish = 'stop'
        parser_error = None
        try:
            calls, content = parse_tools(raw, tools)
        except (ValueError, KeyError, TypeError, json.JSONDecodeError) as exc:
            calls, content, parser_error = [], raw, str(exc)
        message = {'role': 'assistant', 'content': content or None}
        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': tokens, 'completion_tokens': len(tail), 'total_tokens': tokens + len(tail)},
                  'x_runtime': self.info}
        if parser_error:
            result['x_parser_error'] = parser_error
        return result


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--port', type=int, default=8560)
    args = parser.parse_args()
    backend = Bf16Backend()

    class Handler(BaseHTTPRequestHandler):
        def log_message(self, fmt, *args):
            # No prompts, cookies or credentials are logged.
            print(json.dumps({'http': fmt % args, 'time': time.time()}), flush=True)

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

        def do_GET(self):
            if self.path == '/v1/models':
                self.send_json(200, {'object': 'list', 'data': [{**backend.info, 'object': 'model', 'owned_by': 'webbrain-one'}]})
            elif self.path == '/health':
                self.send_json(200, {'ready': True, **backend.info})
            else:
                self.send_json(404, {'error': 'Unknown route'})

        def do_POST(self):
            if self.path != '/v1/chat/completions':
                return self.send_json(404, {'error': 'Unknown route'})
            try:
                length = int(self.headers.get('Content-Length', '0'))
                if not 0 < length <= 4 * 1024 * 1024:
                    raise ValueError('Invalid body length')
                request = json.loads(self.rfile.read(length))
                if request.get('stream', False):
                    raise ValueError('Use stream=false on this benchmark endpoint')
                result = backend.generate(request)
                self.send_json(200, result)
            except (ValueError, KeyError, TypeError) as exc:
                self.send_json(400, {'error': {'message': str(exc), 'type': 'invalid_request_error'}})
            except Exception as exc:
                self.send_json(500, {'error': {'message': str(exc), 'type': 'runtime_error'}})

    print(json.dumps({'ready': True, 'baseUrl': f'http://127.0.0.1:{args.port}/v1', **backend.info}), flush=True)
    ThreadingHTTPServer(('127.0.0.1', args.port), Handler).serve_forever()


if __name__ == '__main__':
    main()