"""ARCHON SFT v2 — Boot-time config validator

Call validate_all() once before the training loop starts.
Any misalignment raises ConfigValidationError immediately with a diagnostic message.

Three hard gates:
  Gate 1 — curriculum_max == total_steps
  Gate 2 — dpo.enabled_from_step == phase7 DPO sub-window start
  Gate 3 — vocab_extension == tokenizer specials count

Usage:
    from _out_config_validator import validate_all
    validate_all(training_yaml_path, tokenizer_path)
"""

from __future__ import annotations

import json
import pathlib


class ConfigValidationError(RuntimeError):
    """Raised when any boot-time config gate fails."""


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _load_yaml(path: str | pathlib.Path) -> dict:
    try:
        import yaml  # PyYAML
    except ImportError as e:
        raise ImportError("PyYAML required: pip install pyyaml") from e
    with open(path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def _count_tokenizer_specials(tokenizer_path: str | pathlib.Path) -> int:
    """Count special tokens in a HF tokenizer.json file.

    A token is considered special if it appears in the 'added_tokens' list with
    'special': true, OR if it appears in tokenizer.model.vocab with id < 32 and
    matches the < ... > pattern (covers both base specials and appended ChatML tokens).
    """
    with open(tokenizer_path, "r", encoding="utf-8") as f:
        tok = json.load(f)

    # Primary: 'added_tokens' list with special=True
    added_tokens = tok.get("added_tokens", [])
    specials = [t for t in added_tokens if t.get("special", False)]
    if specials:
        return len(specials)

    # Fallback: count model vocab entries with angle-bracket naming pattern
    vocab = tok.get("model", {}).get("vocab", {})
    count = sum(1 for k in vocab if k.startswith("<") and k.endswith(">"))
    if count:
        return count

    raise ConfigValidationError(
        f"Cannot determine special tokens count from {tokenizer_path}. "
        "Tokenizer must have 'added_tokens' list or model.vocab with <...> tokens."
    )


# ---------------------------------------------------------------------------
# Individual gates
# ---------------------------------------------------------------------------

def gate_curriculum_coverage(cfg: dict) -> None:
    """Gate 1: curriculum_max_step must equal total_steps.

    Prevents the v1 bug where steps 32K-64K had no defined phase.
    """
    val_section = cfg.get("validate_config_at_boot", {})
    curriculum_max = val_section.get("curriculum_max_step")
    total = val_section.get("total_steps") or cfg.get("training", {}).get("total_steps")

    if curriculum_max is None or total is None:
        raise ConfigValidationError(
            "Gate 1 FAIL: validate_config_at_boot.curriculum_max_step or total_steps missing "
            "from training YAML."
        )

    if curriculum_max != total:
        raise ConfigValidationError(
            f"Gate 1 FAIL: curriculum_max_step={curriculum_max} != total_steps={total}. "
            f"Training steps {curriculum_max}-{total} have no defined curriculum phase. "
            "Update CURRICULUM_PHASES in mtp_task_profiles_v2.py to cover 0-{total}."
        )

    # Also cross-check against the imported module at runtime
    try:
        from _out_mtp_task_profiles_v2 import CURRICULUM_MAX_STEP, TOTAL_STEPS
        if CURRICULUM_MAX_STEP != total:
            raise ConfigValidationError(
                f"Gate 1 FAIL (module check): mtp_task_profiles_v2.CURRICULUM_MAX_STEP="
                f"{CURRICULUM_MAX_STEP} != YAML total_steps={total}."
            )
        if TOTAL_STEPS != total:
            raise ConfigValidationError(
                f"Gate 1 FAIL (module check): mtp_task_profiles_v2.TOTAL_STEPS="
                f"{TOTAL_STEPS} != YAML total_steps={total}."
            )
    except ImportError:
        pass  # module not on sys.path in some test environments; YAML check alone is sufficient


def gate_dpo_alignment(cfg: dict) -> None:
    """Gate 2: dpo.enabled_from_step must equal phase7 DPO sub-window start.

    Prevents the v1 bug where DPO loss was computed on samples that the trainer
    ignored (28K-56K desync = 44% training in noise).
    """
    val_section = cfg.get("validate_config_at_boot", {})
    dpo_enabled = (
        val_section.get("dpo_enabled_from_step")
        or cfg.get("dpo", {}).get("enabled_from_step")
    )
    phase7_start = (
        val_section.get("phase7_dpo_sub_window_start")
    )

    if dpo_enabled is None:
        raise ConfigValidationError(
            "Gate 2 FAIL: dpo.enabled_from_step missing from training YAML."
        )
    if phase7_start is None:
        raise ConfigValidationError(
            "Gate 2 FAIL: validate_config_at_boot.phase7_dpo_sub_window_start missing. "
            "Add it to the YAML so the validator can cross-check."
        )

    if dpo_enabled != phase7_start:
        raise ConfigValidationError(
            f"Gate 2 FAIL: dpo.enabled_from_step={dpo_enabled} != "
            f"phase7_dpo_sub_window_start={phase7_start}. "
            "DPO loss will be computed in a different window than the phase that emits "
            "dpo_pair samples. Fix one of the two to match."
        )

    # Cross-check against module constant
    try:
        from _out_mtp_task_profiles_v2 import DPO_ACTIVE_FROM_STEP, PHASE7_DPO_SUB_WINDOW_START
        if DPO_ACTIVE_FROM_STEP != dpo_enabled:
            raise ConfigValidationError(
                f"Gate 2 FAIL (module check): mtp_task_profiles_v2.DPO_ACTIVE_FROM_STEP="
                f"{DPO_ACTIVE_FROM_STEP} != YAML dpo.enabled_from_step={dpo_enabled}."
            )
        if PHASE7_DPO_SUB_WINDOW_START != phase7_start:
            raise ConfigValidationError(
                f"Gate 2 FAIL (module check): mtp_task_profiles_v2.PHASE7_DPO_SUB_WINDOW_START="
                f"{PHASE7_DPO_SUB_WINDOW_START} != YAML phase7_dpo_sub_window_start={phase7_start}."
            )
    except ImportError:
        pass


def gate_vocab_extension(cfg: dict, tokenizer_path: str | pathlib.Path | None) -> None:
    """Gate 3: vocab_extension must match actual tokenizer specials count.

    Prevents the v1 bug where vocab_extension=6 but tokenizer v2.1 has 22 specials,
    leaving 16 tokens with random embeddings never updated during SFT.
    """
    val_section = cfg.get("validate_config_at_boot", {})
    declared = (
        val_section.get("vocab_extension")
        or cfg.get("model", {}).get("vocab_extension")
    )

    if declared is None:
        raise ConfigValidationError(
            "Gate 3 FAIL: model.vocab_extension missing from training YAML."
        )

    if tokenizer_path is None:
        # Soft path: validate YAML self-consistency only
        expected = val_section.get("expected_tokenizer_specials")
        if expected is None:
            raise ConfigValidationError(
                "Gate 3 FAIL: tokenizer_path not provided and "
                "validate_config_at_boot.expected_tokenizer_specials missing. "
                "Cannot validate vocab_extension."
            )
        if declared != expected:
            raise ConfigValidationError(
                f"Gate 3 FAIL: vocab_extension={declared} != "
                f"expected_tokenizer_specials={expected} (from YAML self-check)."
            )
        return

    actual = _count_tokenizer_specials(tokenizer_path)
    if declared != actual:
        raise ConfigValidationError(
            f"Gate 3 FAIL: model.vocab_extension={declared} but tokenizer at "
            f"'{tokenizer_path}' has {actual} special tokens. "
            f"Delta = {actual - declared} orphaned specials with uninitialized embeddings. "
            "Update model.vocab_extension in the training YAML."
        )


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------

VALIDATION_GATES = [
    gate_curriculum_coverage,
    gate_dpo_alignment,
    gate_vocab_extension,
]
VALIDATION_GATES_COUNT: int = len(VALIDATION_GATES)  # 3


def validate_all(
    training_yaml_path: str | pathlib.Path,
    tokenizer_path: str | pathlib.Path | None = None,
) -> None:
    """Run all boot-time validation gates. Raises ConfigValidationError on first failure.

    Args:
        training_yaml_path: Path to _out_training_v2.yaml (or equivalent).
        tokenizer_path:     Path to tokenizer.json. If None, gate 3 falls back to
                            YAML self-consistency check only.
    """
    cfg = _load_yaml(training_yaml_path)

    failures: list[str] = []

    for gate_fn in [gate_curriculum_coverage, gate_dpo_alignment]:
        try:
            gate_fn(cfg)
        except ConfigValidationError as e:
            failures.append(str(e))

    try:
        gate_vocab_extension(cfg, tokenizer_path)
    except ConfigValidationError as e:
        failures.append(str(e))

    if failures:
        joined = "\n\n".join(failures)
        raise ConfigValidationError(
            f"{len(failures)}/{VALIDATION_GATES_COUNT} validation gate(s) FAILED:\n\n{joined}"
        )

    print(
        f"[config_validator] All {VALIDATION_GATES_COUNT}/{VALIDATION_GATES_COUNT} "
        "gates PASSED. Training config is consistent."
    )


# ---------------------------------------------------------------------------
# CLI: python _out_config_validator.py <yaml_path> [tokenizer_path]
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print(
            f"Usage: python {pathlib.Path(__file__).name} <training_yaml> [tokenizer_path]",
            file=sys.stderr,
        )
        sys.exit(1)

    yaml_path = sys.argv[1]
    tok_path = sys.argv[2] if len(sys.argv) > 2 else None

    try:
        validate_all(yaml_path, tok_path)
        sys.exit(0)
    except ConfigValidationError as exc:
        print(f"\n[config_validator] FATAL: {exc}", file=sys.stderr)
        sys.exit(2)