"""HK provider routes and separately opted-in personal Vertex dictation."""
import os
from pathlib import Path

import yaml

from provider_routes import build_config


def build_hk_config(catalog, env):
    config = build_config(catalog, env, lambda path: False)
    config['general_settings'] = {'master_key': 'os.environ/LITELLM_MASTER_KEY'}
    routes = config['model_list']
    routes[:] = [route for route in routes if route['model_name'] != 'gemini-transcribe-personal']
    if env.get('ZENMUX_SUBSCRIPTION_API_KEY'):
        asr_models = {
            'asr-gpt': 'openai/gpt-transcribe',
            'asr-qwen': 'qwen/qwen-audio-3.1-asr-flash',
            'asr-doubao': 'bytedance/doubao-seed-asr-2.0',
            'asr-mimo': 'xiaomi/mimo-v2.5-asr',
        }
        for alias, model in asr_models.items():
            routes.append({'model_name': alias, 'litellm_params': {
                'model': 'openai/' + model,
                'api_key': 'os.environ/ZENMUX_SUBSCRIPTION_API_KEY',
                'api_base': 'https://zenmux.ai/api/v1',
                'max_retries': 0}})
    if env.get('GOOGLE_AI_STUDIO_API_KEY'):
        routes.append({'model_name': 'gemini-transcribe', 'litellm_params': {
            'model': 'gemini/gemini-3.5-transcribe',
            'api_key': 'os.environ/GOOGLE_AI_STUDIO_API_KEY'}})
    if env.get('VERTEXAI_ENABLED', '').lower() == 'true' and env.get('VERTEXAI_PROJECT'):
        routes.append({'model_name': 'gemini-transcribe-personal', 'litellm_params': {
            'model': 'vertex_ai/gemini-3.5-transcribe',
            'vertex_project': 'os.environ/VERTEXAI_PROJECT',
            'vertex_location': 'global',
            'personal_dictation': True}})
        # Keep old client aliases; the old 3.6 alias now follows the verified Flash map.
        legacy = {
            'gemini-3.6-flash': catalog['google_vertex']['gemini-flash-vertex'],
            'gemini-3.5-flash-lite': catalog['google_vertex']['gemini-flash-lite-vertex'],
            'google/gemini-embedding-2': 'gemini-embedding-2',
            'google/gemini-3.8-flash': catalog['google_vertex']['gemini-flash-vertex'],
        }
        for alias, model in legacy.items():
            routes.append({'model_name': alias, 'litellm_params': {
                'model': 'vertex_ai/' + model,
                'vertex_project': 'os.environ/VERTEXAI_PROJECT',
                'vertex_location': 'os.environ/VERTEXAI_LOCATION'}})
    return config


if __name__ == '__main__':
    catalog = yaml.safe_load(Path('/etc/litellm/models.yaml').read_text())
    config = build_hk_config(catalog, os.environ)
    output = Path('/tmp/litellm.yaml')
    output.write_text(yaml.safe_dump(config, allow_unicode=True, sort_keys=False))
    output.chmod(0o600)
    print('Configured routes:', ', '.join(route['model_name'] for route in config['model_list']))