"""Common utility functions shared between OpenAI and DeepSeek implementations."""

import json
from typing import Dict, Tuple


def save_mapping(mapping: Dict[str, Tuple[int, str]], output_file: str):
    """Save token mapping to a JSON file."""
    serializable_mapping = {
        orig: {"id": tid, "token": tok} for orig, (tid, tok) in mapping.items()
    }

    with open(output_file, "w") as f:
        json.dump(serializable_mapping, f, indent=2)


def load_mapping(input_file: str) -> Dict[str, Tuple[int, str]]:
    """Load token mapping from a JSON file."""
    with open(input_file) as f:
        data = json.load(f)

    return {orig: (info["id"], info["token"]) for orig, info in data.items()}


def print_token_mapping(
    mapping: Dict[str, Tuple[int, str]], title: str = "Token Mapping"
):
    """Pretty print the token mapping."""
    print(f"\n{title}:")
    print("Original Token\tNew Token ID\tNew Token (showing · for space)")
    print("-" * 60)
    for orig_token, (new_id, new_token) in sorted(mapping.items()):
        orig_visible = orig_token.replace(" ", "·")
        new_visible = new_token.replace(" ", "·")
        print(f"{orig_visible}\t{new_id}\t{new_visible}")