from __future__ import annotations

import argparse
from pathlib import Path

from engine.refinement.introspection import (
    FEW_SHOT_PATH,
    RULES_PATH,
    train_introspection_rules,
)


def _count_nonempty_blocks(path: Path) -> int:
    """Compta aproximadament quants blocs d'exemple hi ha en un fitxer.

    Per a ``few_shot_examples.txt`` comptem línies que comencen per ``# sha1sum=``.
    Per a ``rules.txt`` comptem línies no buides.
    """

    if not path.exists():
        return 0
    try:
        text = path.read_text(encoding="utf-8")
    except Exception:
        return 0

    if path.name == "few_shot_examples.txt":
        return sum(1 for line in text.splitlines() if line.lstrip().startswith("# sha1sum="))
    return sum(1 for line in text.splitlines() if line.strip())


def main() -> None:
    parser = argparse.ArgumentParser(
        description=(
            "Entrena les regles d'introspecció a partir de les correccions HITL "
            "emmagatzemades a demo/temp/audiodescriptions.db."
        )
    )
    parser.add_argument(
        "--max-examples",
        type=int,
        default=None,
        help=(
            "Nombre màxim de parelles (MoE/Salamandra, HITL) a processar. "
            "Per defecte es processen totes."
        ),
    )
    args = parser.parse_args()

    train_introspection_rules(max_examples=args.max_examples)

    n_examples = _count_nonempty_blocks(FEW_SHOT_PATH)
    n_rules = _count_nonempty_blocks(RULES_PATH)

    print(
        f"✅ Entrenament d'introspection completat. "
        f"Few-shot examples: {n_examples}, regles: {n_rules}."
    )


if __name__ == "__main__":
    main()