"""Central configuration for the translation pipeline and review tool.

To target a new language, change the three TARGET_* values (and PACK_NAME).
Everything else derives from them.
"""

from pathlib import Path

ROOT = Path(__file__).parent

# --- Target language -------------------------------------------------------
TARGET_LANG_CODE = "de"  # ISO 639-1 code understood by TranslateGemma
TARGET_LANG_NAME = "German"
PACK_NAME = "German_JSON"

SOURCE_LANG_CODE = "en"
SOURCE_LANG_NAME = "English"

# --- Paths ------------------------------------------------------------------
SOURCE_DIR = ROOT / "English_JSON"
PACK_DIR = ROOT / PACK_NAME
TRANSLATIONS_DIR = ROOT / "translations" / TARGET_LANG_CODE
CACHE_PATH = TRANSLATIONS_DIR / ".cache.json"
FILE_CONTEXT_PATH = ROOT / "file_context.json"
# Public fallback used on HF Spaces: context for the demo sample files only.
FILE_CONTEXT_SAMPLE_PATH = ROOT / "file_context.sample.json"
FILE_MAPPING_PATH = ROOT / "FILE_MAPPING.md"
CHAR_DATA_DIR = ROOT / "character_data"
CHAR_WIKI_DIR = ROOT / "character_wikis"

# --- Hugging Face in-process model inference --------------------------------
# Both translation stages now run inside Python with Transformers. The model
# repos below are the canonical Hugging Face checkpoints, not local HTTP
# servers. TranslateGemma is gated, so the runtime needs an HF token whose
# account has accepted Google's Gemma terms.
TRANSLATE_MODEL_ID = "google/translategemma-12b-it"
TONE_MODEL_ID = "google/gemma-4-12B-it"

# "auto" follows the model card examples and works locally and on GPU Spaces.
# Set HF_DEVICE_MAP = None to load on CPU first and explicitly move to cuda.
HF_DEVICE_MAP = "auto"
HF_DTYPE = "auto"
HF_ATTN_IMPLEMENTATION = None

TRANSLATE_MAX_NEW_TOKENS = 1024
TONE_MAX_NEW_TOKENS = 512

# Gemma's officially recommended sampling parameters for creative rewriting.
TONE_GENERATION_KWARGS = {
    "do_sample": True,
    "temperature": 1.0,
    "top_k": 64,
    "top_p": 0.95,
}

# ZeroGPU only exposes the real GPU inside @spaces.GPU functions. The decorator
# is a no-op outside ZeroGPU, so keeping these values here makes deployment a
# config change instead of a code fork.
ZERO_GPU_DURATION_S = 600
ZERO_GPU_SIZE = "xlarge"

# Keeping both 12B checkpoints resident can exceed the default 48GB ZeroGPU
# slice. The pipeline unloads the other model before loading the requested one
# unless this is set to True on a larger GPU.
HF_KEEP_BOTH_MODELS = False
HF_PRELOAD_MODELS = ()

REQUEST_TIMEOUT_S = 600
MAX_ATTEMPTS = 3
RETRY_BACKOFF_S = 5