"""Paths and test texts shared by the Qwen export scripts."""

import csv
import re
import unicodedata
from pathlib import Path

HERE = Path(__file__).resolve().parent
SIZES = ("1.5B", "VyLinh-3B")
# Source folder under data/ for each model. VyLinh-3B is Arcee's Vietnamese Qwen2.5-3B merge (same architecture
# and tokenizer as Qwen2.5), so the same export, prompt and checks apply.
SOURCES = {
    "1.5B": "Qwen2.5-1.5B-Instruct",  # hf download Qwen/Qwen2.5-1.5B-Instruct --local-dir data/Qwen2.5-1.5B-Instruct
    "VyLinh-3B": "Arcee-VyLinh",  # hf download arcee-ai/Arcee-VyLinh --local-dir data/Arcee-VyLinh
}
FLEURS_TSV = HERE.parent / "mobile_export" / "data" / "fleurs" / "data" / "vi_vn" / "test.tsv"


def src_dir(size: str) -> Path:
    """The downloaded Hugging Face model: see SOURCES."""
    return HERE / "data" / SOURCES[size]


def out_dir(size: str) -> Path:
    """The repo that publish.py uploads: onnx/, vocab.txt, merges.txt, LICENSE."""
    return HERE / "out" / size / "repo"


def strip(text: str) -> str:
    """Lower-cases and drops punctuation, the way ASR output looks."""
    text = unicodedata.normalize("NFC", text.lower())
    return re.sub(r"\s+", " ", re.sub(r"[^\w\s]", " ", text)).strip()


def fleurs_sentences() -> list[tuple[str, str]]:
    """(stripped, original) for each unique FLEURS Vietnamese test transcription."""
    seen, rows = set(), []
    with open(FLEURS_TSV, encoding="utf-8") as f:
        for row in csv.reader(f, delimiter="\t", quoting=csv.QUOTE_NONE):
            truth = unicodedata.normalize("NFC", row[2].strip())
            if truth not in seen:
                seen.add(truth)
                rows.append((strip(truth), truth))
    return rows


def fleurs_paragraphs(size: int = 3) -> list[tuple[str, str]]:
    """Consecutive FLEURS sentences joined into paragraphs, about the length of a 30-second recording."""
    rows = fleurs_sentences()
    groups = [rows[i:i + size] for i in range(0, len(rows) - size + 1, size)]
    return [(" ".join(s for s, _ in g), " ".join(t for _, t in g)) for g in groups]