"""Static configuration for the LoRA training worker.

Everything load-bearing is an explicit module constant here (explicit-over-implicit): the
paths the worker is allowed to touch, the fixed sd-scripts invocation surface, the tier
table, and the SDXL-LoRA hyperparameters. Nothing about a training run is free-form — the
worker only ever builds ONE command shape (INV-T7), and every knob that shapes it is
visible in this file.

The worker runs as `llmuser` on irv-ml1. It never imports torch / sd-scripts; it only
*subprocesses* the fluxgym venv's `accelerate launch`. So this process stays tiny and the
fluxgym venv stays pristine.
"""

from __future__ import annotations

import os
from pathlib import Path

# ---- Network ---------------------------------------------------------------------------
# 0.0.0.0 (not 127.0.0.1): arbo runs CONTAINERIZED on the traefik-net bridge and reaches
# the host worker via host.docker.internal:host-gateway — host loopback is unreachable from
# that bridge. irv-ml1 is WireGuard-only + ACL'd, so 0.0.0.0 exposure is bounded to the tunnel.
HOST = os.environ.get("LORA_WORKER_HOST", "0.0.0.0")
PORT = int(os.environ.get("LORA_WORKER_PORT", "8203"))

# ---- Fluxgym / sd-scripts (the ONLY thing the worker executes) -------------------------
FLUXGYM_ROOT = Path("/opt/fluxgym")
FLUXGYM_VENV = FLUXGYM_ROOT / ".venv"
ACCELERATE_BIN = FLUXGYM_VENV / "bin" / "accelerate"
SD_SCRIPTS_DIR = FLUXGYM_ROOT / "sd-scripts"
SDXL_TRAIN_SCRIPT = SD_SCRIPTS_DIR / "sdxl_train_network.py"

# ---- Allowed path roots (INV-T7 path safety) -------------------------------------------
# dataset_dir + output_dir MUST live under the shared handoff root (the group-shared,
# setgid /worktank/arbo/train that both arbo-container-uid and llmuser can rw). base_model
# must live under a known model root. Anything else → rejected before a process is spawned.
# base_model must be an absolute HOST path (the worker runs on the host, not in arbo's
# container — a container-internal path like /comfy/... or /basedir/... won't match).
#   /storetank/arbo/models — canonical model store (SDXL checkpoints; the 2026-06-13 move
#     to the 1.8TB /storetank volume; arbo mounts it → /basedir/models inside its container).
#   /opt/fluxgym/models — fluxgym base models (flux; Phase 4).
#   /worktank/arbo — handoff root (a base staged into the handoff, if ever).
# (Dropped the pre-move /worktank/models [gone] + /worktank/comfyui [empty host path] roots.)
HANDOFF_ROOT = Path(os.environ.get("LORA_WORKER_HANDOFF_ROOT", "/worktank/arbo/train"))
ALLOWED_MODEL_ROOTS = tuple(
    Path(p)
    for p in os.environ.get(
        "LORA_WORKER_MODEL_ROOTS",
        "/storetank/arbo/models:/opt/fluxgym/models:/worktank/arbo",
    ).split(":")
    if p
)

# ---- LoRA publish step (Phase 2 — auto-registration into ComfyUI Generate) -------------
# On a train reaching `succeeded`, IN ADDITION to output/{name}.safetensors (the unchanged
# download source), the worker COPIES the LoRA into ComfyUI's loras search path under
# PUBLISH_SUBDIR/{train_id}/ and reports `published_lora_name` (the path RELATIVE to the
# loras root — the exact string a ComfyUI LoraLoader.lora_name widget takes). ComfyUI
# (verified 2026-07-07) resolves loras SUBfolders, so the nested scheme works. A publish
# failure NEVER fails the train — it just omits published_lora_name (download still works).
# Stays inside INV-T7 (a copy to a fixed computed path; no new free-form args).
LORAS_PUBLISH_ROOT = Path(os.environ.get("LORA_WORKER_LORAS_ROOT", "/storetank/arbo/models/loras"))
PUBLISH_SUBDIR = "trained"  # loras/trained/{train_id}/{name}.safetensors

# ---- Worker state + logs (survives a worker restart for boot reconciliation, §4.3) -----
STATE_DIR = Path(os.environ.get("LORA_WORKER_STATE_DIR", "/opt/lora-training-worker/state"))
LOG_DIR = Path(os.environ.get("LORA_WORKER_LOG_DIR", "/opt/lora-training-worker/logs"))
JOBS_FILE = STATE_DIR / "jobs.json"
MAX_RETAINED_JOBS = 50  # keep terminal jobs queryable for arbo's status proxy

# ---- Tier table (§4.6) -----------------------------------------------------------------
# tier -> (max_train_steps, network_dim, resolution). Quality (1024) is A6000-only; the
# device-fit guard lives in invocation.build_command (quality on the 3090 → rejected).
TIERS = {
    "fast": {"steps": 400, "dim": 16, "resolution": 768},
    "balanced": {"steps": 1500, "dim": 32, "resolution": 768},
    "quality": {"steps": 3000, "dim": 32, "resolution": 1024},
}

# ---- Device model (CUDA_DEVICE_ORDER=PCI_BUS_ID) ---------------------------------------
# Under PCI_BUS_ID (which the recipe pins), index 0 = RTX 3090, index 1 = RTX A6000.
# Quality (1024) will not fit the 3090's ~24GB LoRA envelope → allowed on the A6000 only.
DEVICE_3090 = 0
DEVICE_A6000 = 1
QUALITY_ONLY_DEVICES = (DEVICE_A6000,)

# ---- SDXL-LoRA hyperparameters (agent-discretion defaults; confirm vs Sindra) ----------
# The contract (§4.6) pins the flag SET + tier->(steps,dim,res) but NOT lr/scheduler/batch.
# These are standard, conservative SDXL-LoRA values, surfaced here so they're one-line
# auditable + tunable. Flagged to comfy-dev to cross-check against the proven Sindra runs.
SDXL_HPARAMS = {
    "learning_rate": "1e-4",
    "lr_scheduler": "cosine",
    "lr_warmup_steps": "0",
    "train_batch_size_lean": "1",  # 3090
    "train_batch_size_full": "2",  # A6000
    # alpha = dim * ratio. 0.5 (alpha=dim/2) matches the PROVEN Sindra v1/v2 runs (comfy-dev
    # cross-check 2026-07-06) — it's the scaling that produced the validated likeness. alpha=dim
    # (1.0) is a valid stronger default but wasn't what Sindra used; operator can override.
    "network_alpha_ratio": 0.5,
    "min_snr_gamma": "5",
    "noise_offset": "0.1",
    "save_precision": "bf16",
    "mixed_precision": "bf16",
    "max_data_loader_n_workers": "2",
}

# ---- TTS-liveness heuristic (GET /gpu-status → tts_on_3090) -----------------------------
# arbo's scheduler steers a lean train off the 3090 when TTS is live there. We report the
# raw per-device VRAM + a boolean derived from (a) a compute-app on the 3090 whose cmdline
# matches a TTS marker, else (b) a used-MB floor fallback.
TTS_CMDLINE_MARKERS = ("chatterbox", "omnivoice", "tts_server", "chatterbox_fast")
TTS_3090_USED_MB_FLOOR = int(os.environ.get("LORA_WORKER_TTS_FLOOR_MB", "2000"))