#!/usr/bin/env python3
"""Stock vLLM API server with its supported pre-validation HF config callable."""
from __future__ import annotations

import copy
from contextlib import contextmanager
import hashlib
import json
import shutil
import tempfile
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class CheckpointRoPEOverride:
    rope_parameters: dict

    def __call__(self, config):
        # HFConfigParser first probes callables with a documented dummy config
        # to detect model_type changes. Leave that probe untouched.
        if config.model_type == 'dummy_mellum':
            return config
        # vLLM's MellumConfig inherits Qwen3MoeConfig, whose constructor adds an
        # unused outer rope_theta scalar. Restore the exact checkpoint dictionary
        # before vLLM validates it; each decoder selects its original layer type.
        if config.model_type != 'mellum' or set(config.layer_types) != set(self.rope_parameters):
            raise ValueError('RoPE override requires the Mellum checkpoint layer types')
        config.rope_parameters = copy.deepcopy(self.rope_parameters)
        return config


def checkpoint_override(model: str) -> CheckpointRoPEOverride:
    path = Path(model) / 'config.json'
    raw = path.read_bytes()
    config = json.loads(raw)
    rope = config.get('rope_parameters')
    if (config.get('model_type') != 'mellum' or not isinstance(rope, dict)
            or set(rope) != set(config['layer_types'])
            or not all(isinstance(value, dict) for value in rope.values())):
        raise ValueError('Require the local Mellum checkpoint nested RoPE parameters')
    print(json.dumps({'compatibility': 'supported callable hf_overrides',
                      'checkpoint_config_sha256': hashlib.sha256(raw).hexdigest(),
                      'rope_parameters': rope,
                      'layer_types': config['layer_types']}), flush=True)
    return CheckpointRoPEOverride(rope)


@contextmanager
def tokenizer_directory(model: str):
    # TokenizerRegistry independently reads model config without hf_overrides.
    # A tokenizer-only local path is supported, and its absent model config is
    # explicitly tolerated. Copy original bytes; never manufacture a config.
    source = Path(model)
    with tempfile.TemporaryDirectory(prefix='tokenizer-only-') as directory:
        target = Path(directory)
        for name in ('tokenizer.json', 'tokenizer_config.json', 'chat_template.jinja',
                     'special_tokens_map.json', 'added_tokens.json'):
            if (source / name).is_file():
                shutil.copy2(source / name, target / name)
        if not (target / 'tokenizer.json').is_file() or not (target / 'tokenizer_config.json').is_file():
            raise ValueError('Require checkpoint tokenizer.json and tokenizer_config.json')
        yield str(target)


def main():
    import uvloop
    from vllm import AsyncEngineArgs
    from vllm.entrypoints.launchers.api_server.entry import run_server
    from vllm.entrypoints.launchers.cli_args import make_arg_parser, validate_parsed_serve_args
    from vllm.entrypoints.serve.utils.api_utils import cli_env_setup
    from vllm.utils.argparse_utils import FlexibleArgumentParser

    cli_env_setup()
    parser = FlexibleArgumentParser(description=__doc__)
    parser.add_argument('--compat-config-only', action='store_true',
                        help='Build actual engine config without loading weights or serving')
    args = make_arg_parser(parser).parse_args()
    if getattr(args, 'model_tag', None):
        args.model = args.model_tag
    if args.hf_overrides:
        parser.error('This launcher owns the exact checkpoint RoPE override')
    if args.tokenizer is not None:
        parser.error('This launcher uses unchanged tokenizer files from --model; omit --tokenizer')
    if args.tokenizer_mode not in ('auto', 'hf'):
        parser.error('This launcher requires the verified HF tokenizer mode')
    args.hf_overrides = checkpoint_override(args.model)
    with tokenizer_directory(args.model) as directory:
        args.tokenizer, args.tokenizer_mode = directory, 'hf'
        validate_parsed_serve_args(args)
        if args.compat_config_only:
            from vllm.tokenizers.registry import cached_tokenizer_from_config
            engine = AsyncEngineArgs.from_cli_args(args).create_engine_config()
            original = json.loads((Path(args.model) / 'config.json').read_text())
            expected = original['rope_parameters']
            actual = engine.model_config.hf_config
            if actual.rope_parameters != expected or actual.layer_types != original['layer_types']:
                raise RuntimeError('Engine config changed checkpoint RoPE semantics')
            selected = [actual.rope_parameters[k] for k in actual.layer_types]
            if selected != [expected[k] for k in original['layer_types']]:
                raise RuntimeError('Per-layer RoPE selection changed')
            tokenizer = cached_tokenizer_from_config(engine.model_config)
            print(json.dumps({'status': 'passed', 'scope': 'Actual engine config and tokenizer registry',
                              'decoder_layers': len(selected), 'rope_parameters': expected,
                              'bos_token_id': tokenizer.bos_token_id,
                              'eos_token_id': tokenizer.eos_token_id}), flush=True)
            return
        uvloop.run(run_server(args))


if __name__ == '__main__':
    main()