#!/usr/bin/env python
from __future__ import annotations

import argparse
from collections import Counter
from dataclasses import asdict, dataclass, field
import glob
import hashlib
import json
import os
from pathlib import Path
import random
import re
import subprocess
import sys
from typing import Iterable, Mapping, Sequence

import pyarrow as pa
import pyarrow.parquet as pq


PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


DEFAULT_JIAYI_SCORED_INPUT = (
    "/mnt/beegfs/cjy/cjy-training/laoyao_model/data/scored/"
    "edu_fineweb_math_knowledge_cross_near_h1_v1_20260620/rank_*/*.jsonl"
)
DEFAULT_CODE_INPUT = (
    "/mnt/beegfs/cjy/cjy-training/laoyao_model/data/dedup/"
    "pretrain_python_clean_code_multilang_cross_near_h1_min256_v1_20260623/*.parquet"
)
DEFAULT_OUTPUT_DIR = "data/train/pretrain_rebalanced_web40_edu20_chinese10_science10_logic10_math5_code5_v1_20260630"


DEFAULT_MIX = {
    "english_web": 0.40,
    "english_edu": 0.20,
    "chinese_clean": 0.10,
    "science": 0.10,
    "logic": 0.10,
    "math": 0.05,
    "code": 0.05,
}


SHORTFALL_SOURCE = "science"
SHORTFALL_TARGET = "english_edu"


BYTES_PER_TOKEN = {
    "english_web": 4.0,
    "english_edu": 4.0,
    "chinese_clean": 2.0,
    "science": 3.8,
    "logic": 4.0,
    "math": 3.8,
    "code": 3.4,
}


TEXT_CANDIDATE_FIELDS = (
    "text",
    "content",
    "document",
    "prompt",
    "question",
    "query",
    "context",
    "passage",
    "choices",
    "options",
    "answer",
    "answerKey",
    "answer_key",
    "opa",
    "opb",
    "opc",
    "opd",
    "cop",
    "label",
    "solution",
    "explanation",
    "exp",
)


OUTPUT_SCHEMA = pa.schema(
    [
        pa.field("text", pa.string()),
        pa.field("mixture_category", pa.string()),
        pa.field("mixture_family", pa.string()),
        pa.field("estimated_tokens", pa.int64()),
        pa.field("source", pa.string()),
        pa.field("subset", pa.string()),
        pa.field("language", pa.string()),
        pa.field("quality_bucket", pa.string()),
        pa.field("edu_quality_bucket", pa.string()),
        pa.field("math_quality_bucket", pa.string()),
        pa.field("corpus_source", pa.string()),
        pa.field("source_path", pa.string()),
        pa.field("source_row", pa.int64()),
        pa.field("input_format", pa.string()),
        pa.field("selection_seed", pa.int64()),
        pa.field("metadata_json", pa.string()),
    ]
)


MATH_MARKERS = (
    "math",
    "mathematics",
    "algebra",
    "geometry",
    "calculus",
    "theorem",
    "proof",
    "equation",
    "integral",
    "derivative",
    "probability",
    "combinatorics",
    "number theory",
    "finemath",
    "proof-pile",
    "proofpile",
)


LOGIC_MARKERS = (
    "logic",
    "reasoning",
    "brain teaser",
    "brainteaser",
    "puzzle",
    "riddle",
    "deduce",
    "deduction",
    "premise",
    "conclusion",
    "if and only if",
    "necessary and sufficient",
    "counterexample",
    "paradox",
)


@dataclass(frozen=True)
class SourceSpec:
    family: str
    paths: tuple[Path, ...]


@dataclass
class BuildStats:
    files_seen: int = 0
    docs_seen: int = 0
    docs_written: int = 0
    docs_dropped_missing_text: int = 0
    docs_dropped_quota_full: int = 0
    malformed_json: int = 0
    estimated_tokens_seen: int = 0
    estimated_tokens_written: int = 0
    shards_written: int = 0
    category_docs: Counter[str] = field(default_factory=Counter)
    category_tokens: Counter[str] = field(default_factory=Counter)
    family_docs: Counter[str] = field(default_factory=Counter)
    family_tokens: Counter[str] = field(default_factory=Counter)
    language_docs: Counter[str] = field(default_factory=Counter)
    source_docs: Counter[str] = field(default_factory=Counter)

    def to_dict(self) -> dict[str, object]:
        payload = asdict(self)
        for key in (
            "category_docs",
            "category_tokens",
            "family_docs",
            "family_tokens",
            "language_docs",
            "source_docs",
        ):
            counter = getattr(self, key)
            payload[key] = dict(counter.most_common())
        return payload


class ParquetShardWriter:
    def __init__(self, output_dir: Path, *, shard_rows: int, compression: str) -> None:
        self.output_dir = output_dir
        self.shard_rows = shard_rows
        self.compression = compression
        self.buffer: list[dict[str, object]] = []
        self.shards_written = 0

    def write(self, row: dict[str, object]) -> None:
        self.buffer.append(row)
        if len(self.buffer) >= self.shard_rows:
            self.flush()

    def flush(self) -> None:
        if not self.buffer:
            return
        path = self.output_dir / f"part_{self.shards_written:06d}.parquet"
        table = pa.Table.from_pylist(self.buffer, schema=OUTPUT_SCHEMA)
        pq.write_table(table, path, compression=self.compression)
        self.buffer.clear()
        self.shards_written += 1

    def close(self) -> None:
        self.flush()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Build a rebalanced pretraining dataset with explicit token quotas: "
            "40% English Web, 20% English Edu, 10% Chinese clean, 10% science, "
            "10% logic, 5% math, 5% code."
        )
    )
    parser.add_argument("--english-web-input", action="append", default=[], help="FineWeb parquet/jsonl source.")
    parser.add_argument(
        "--english-edu-input",
        action="append",
        default=[],
        help="FineWeb-Edu parquet/jsonl source. Prefer downloaded ModelScope/HF parquet for full builds.",
    )
    parser.add_argument("--chinese-clean-input", action="append", default=[], help="CCI3-HQ/SkyPile-150B source.")
    parser.add_argument(
        "--science-input",
        action="append",
        default=[],
        help="Science QA source, e.g. medmcqa/proofwriter/scienceqa/sciq/qasc/openbookqa.",
    )
    parser.add_argument(
        "--logic-input",
        action="append",
        default=[],
        help="Optional explicit logic source path/dir/glob. Supports parquet, jsonl, and jsonl.zst.",
    )
    parser.add_argument("--math-input", action="append", default=[], help="Optional explicit math source path/dir/glob.")
    parser.add_argument("--code-input", action="append", default=[], help="Code source path/dir/glob.")
    parser.add_argument(
        "--jiayi-scored-input",
        action="append",
        default=[],
        help=(
            "Mixed cjy scored text pool. Rows are classified by metadata/text into math, logic, "
            "english_edu, or chinese_clean instead of being forced into one bucket."
        ),
    )
    parser.add_argument("--output-dir", default=DEFAULT_OUTPUT_DIR)
    parser.add_argument("--target-tokens", type=int, default=200_000_000_000)
    parser.add_argument("--mix", action="append", default=[], metavar="CATEGORY=FRACTION")
    parser.add_argument("--seed", type=int, default=20260630)
    parser.add_argument("--output-shard-rows", type=int, default=100_000)
    parser.add_argument("--compression", default="zstd")
    parser.add_argument("--quality-keep", default="high,mid,medium,unknown")
    parser.add_argument(
        "--allow-science-shortfall-to-english-edu",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="When science is exhausted below quota, add the shortfall to english_edu effective quota.",
    )
    parser.add_argument("--max-files-per-family", type=int, default=None, help="Smoke/debug cap per input family.")
    parser.add_argument("--max-docs-per-file", type=int, default=None, help="Smoke/debug cap per file.")
    parser.add_argument("--progress-every-files", type=int, default=25)
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--dry-run", action="store_true", help="Resolve inputs and quotas without writing data.")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    mix = parse_mix(args.mix)
    quotas = {category: int(args.target_tokens * fraction) for category, fraction in mix.items()}
    effective_quotas = dict(quotas)
    specs = build_source_specs(args)
    output_dir = Path(args.output_dir)
    if args.dry_run:
        report = build_dry_run_report(specs, mix, quotas, args)
        print(json.dumps(report, ensure_ascii=False, indent=2))
        return

    prepare_output_dir(output_dir, overwrite=args.overwrite)
    writer = ParquetShardWriter(output_dir, shard_rows=args.output_shard_rows, compression=args.compression)
    stats = BuildStats()
    accepted_tokens = Counter()
    quality_keep = {item.strip().lower() for item in args.quality_keep.split(",") if item.strip()}
    if args.allow_science_shortfall_to_english_edu and not any(spec.family == SHORTFALL_SOURCE for spec in specs):
        apply_science_shortfall_to_english_edu(accepted_tokens, quotas, effective_quotas)

    try:
        total_files = sum(len(spec.paths) for spec in specs)
        completed_files = 0
        for spec in specs:
            paths = list(spec.paths)
            random.Random(args.seed + stable_int(spec.family)).shuffle(paths)
            if args.max_files_per_family is not None:
                paths = paths[: args.max_files_per_family]
            for path in paths:
                completed_files += 1
                stats.files_seen += 1
                for row_index, row in enumerate(iter_rows(path)):
                    if args.max_docs_per_file is not None and row_index >= args.max_docs_per_file:
                        break
                    stats.docs_seen += 1
                    text = row_text(row)
                    if not text:
                        stats.docs_dropped_missing_text += 1
                        continue
                    category = classify_row(row, default_family=spec.family)
                    if category not in quotas:
                        continue
                    if not keep_quality(row, category=category, quality_keep=quality_keep):
                        continue
                    estimated_tokens = estimate_tokens(row, text, category=category)
                    stats.estimated_tokens_seen += estimated_tokens
                    if accepted_tokens[category] >= effective_quotas[category]:
                        stats.docs_dropped_quota_full += 1
                        continue
                    output_row = build_output_row(
                        row,
                        text=text,
                        category=category,
                        family=spec.family,
                        path=path,
                        row_index=row_index,
                        estimated_tokens=estimated_tokens,
                        seed=args.seed,
                    )
                    writer.write(output_row)
                    accepted_tokens[category] += estimated_tokens
                    stats.docs_written += 1
                    stats.estimated_tokens_written += estimated_tokens
                    stats.category_docs[category] += 1
                    stats.category_tokens[category] += estimated_tokens
                    stats.family_docs[spec.family] += 1
                    stats.family_tokens[spec.family] += estimated_tokens
                    stats.language_docs[string_value(row.get("language")) or "unknown"] += 1
                    stats.source_docs[infer_source(row)] += 1
                if should_log_progress(completed_files, total_files, args.progress_every_files):
                    print(
                        "rebalanced_pretrain_progress "
                        f"files={completed_files}/{total_files} docs={stats.docs_seen} "
                        f"written={stats.docs_written} tokens={stats.estimated_tokens_written} "
                        f"category_tokens={dict(stats.category_tokens)}",
                        flush=True,
                    )
                if args.allow_science_shortfall_to_english_edu and spec.family == SHORTFALL_SOURCE:
                    apply_science_shortfall_to_english_edu(accepted_tokens, quotas, effective_quotas)
                if all(accepted_tokens[category] >= quota for category, quota in effective_quotas.items()):
                    break
    finally:
        writer.close()
    stats.shards_written = writer.shards_written
    if args.allow_science_shortfall_to_english_edu:
        apply_science_shortfall_to_english_edu(accepted_tokens, quotas, effective_quotas)
    write_report(output_dir / "report.json", stats, specs, mix, quotas, effective_quotas, args)
    print(
        "rebalanced_pretrain_dataset "
        f"output={output_dir} shards={stats.shards_written} docs={stats.docs_written} "
        f"tokens={stats.estimated_tokens_written} category_tokens={dict(stats.category_tokens)}"
    )


def parse_mix(overrides: Sequence[str]) -> dict[str, float]:
    mix = dict(DEFAULT_MIX)
    for item in overrides:
        if "=" not in item:
            raise ValueError(f"--mix must use CATEGORY=FRACTION: {item!r}")
        key, raw = item.split("=", 1)
        key = key.strip()
        if key not in DEFAULT_MIX:
            raise ValueError(f"unknown mix category {key!r}; expected one of {sorted(DEFAULT_MIX)}")
        mix[key] = float(raw)
    total = sum(mix.values())
    if not 0.999 <= total <= 1.001:
        raise ValueError(f"mix fractions must sum to 1.0, got {total}")
    return mix


def apply_science_shortfall_to_english_edu(
    accepted_tokens: Mapping[str, int],
    quotas: Mapping[str, int],
    effective_quotas: dict[str, int],
) -> None:
    science_target = quotas.get(SHORTFALL_SOURCE, 0)
    science_actual = min(int(accepted_tokens.get(SHORTFALL_SOURCE, 0)), science_target)
    science_shortfall = max(0, science_target - science_actual)
    effective_quotas[SHORTFALL_TARGET] = quotas.get(SHORTFALL_TARGET, 0) + science_shortfall


def build_source_specs(args: argparse.Namespace) -> tuple[SourceSpec, ...]:
    raw_specs = (
        ("english_web", args.english_web_input),
        ("science", args.science_input),
        ("english_edu", args.english_edu_input),
        ("chinese_clean", args.chinese_clean_input),
        ("auto_text", args.jiayi_scored_input),
        ("logic", args.logic_input),
        ("math", args.math_input),
        ("code", args.code_input),
    )
    specs: list[SourceSpec] = []
    for family, patterns in raw_specs:
        if not patterns:
            continue
        paths = resolve_paths(patterns)
        specs.append(SourceSpec(family=family, paths=paths))
    if not specs:
        raise ValueError("no inputs resolved")
    return tuple(specs)


def resolve_paths(patterns: Sequence[str]) -> tuple[Path, ...]:
    paths: list[Path] = []
    for pattern in patterns:
        path = Path(pattern)
        if path.is_dir():
            paths.extend(path.glob("*.jsonl"))
            paths.extend(path.glob("*.jsonl.zst"))
            paths.extend(path.glob("*.parquet"))
            paths.extend(path.glob("rank_*/*.jsonl"))
            paths.extend(path.glob("rank_*/*.jsonl.zst"))
            paths.extend(path.glob("rank_*/*.parquet"))
            continue
        matches = glob.glob(pattern, recursive=True)
        if matches:
            paths.extend(Path(match) for match in matches)
            continue
        if path.exists():
            paths.append(path)
            continue
        raise FileNotFoundError(pattern)
    supported = [path.resolve() for path in paths if file_format(path) in {"jsonl", "jsonl_zst", "parquet"}]
    unique = tuple(sorted(set(supported), key=str))
    if not unique:
        raise FileNotFoundError(f"no JSONL/JSONL.ZST/Parquet files matched: {patterns!r}")
    return unique


def iter_rows(path: Path) -> Iterable[Mapping[str, object]]:
    fmt = file_format(path)
    if fmt == "jsonl":
        with path.open("r", encoding="utf-8") as handle:
            for line in handle:
                try:
                    row = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if isinstance(row, Mapping):
                    yield row
        return
    if fmt == "jsonl_zst":
        with subprocess.Popen(
            ["zstd", "-dc", str(path)],
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            text=True,
            encoding="utf-8",
            errors="replace",
        ) as proc:
            assert proc.stdout is not None
            for line in proc.stdout:
                try:
                    row = json.loads(line)
                except json.JSONDecodeError:
                    continue
                if isinstance(row, Mapping):
                    yield row
            returncode = proc.wait()
            if returncode:
                raise RuntimeError(f"zstd failed with exit code {returncode}: {path}")
        return
    parquet_file = pq.ParquetFile(path)
    columns = [
        name
        for name in (*OUTPUT_SCHEMA.names, *TEXT_CANDIDATE_FIELDS)
        if name in parquet_file.schema_arrow.names
    ]
    if not any(name in columns for name in TEXT_CANDIDATE_FIELDS):
        return
    for batch in parquet_file.iter_batches(batch_size=8192, columns=columns, use_threads=True):
        for row in batch.to_pylist():
            if isinstance(row, Mapping):
                yield row


def classify_row(row: Mapping[str, object], *, default_family: str) -> str:
    if default_family in DEFAULT_MIX:
        return default_family
    haystack = " ".join(
        string_value(row.get(key))
        for key in (
            "source",
            "subset",
            "corpus_source",
            "mixture_domain",
            "domain",
            "language",
            "quality_bucket",
            "math_quality_bucket",
            "edu_quality_bucket",
        )
    ).lower()
    text_prefix = string_value(row.get("text"))[:4096].lower()
    combined = f"{haystack} {text_prefix}"
    if any(marker in combined for marker in MATH_MARKERS):
        return "math"
    if any(marker in combined for marker in LOGIC_MARKERS):
        return "logic"
    language = normalize_language(string_value(row.get("language")), row_text(row))
    if language == "en":
        return "english_edu"
    if language in {"zh", "zh-cn", "zh-tw", "cmn"}:
        return "chinese_clean"
    if default_family == "auto_text":
        return "english_edu"
    return "science"


def normalize_language(raw_language: str, text: str) -> str:
    language = raw_language.strip().lower().replace("_", "-")
    if language:
        return language
    sample = text[:2048]
    if re.search(r"[\u4e00-\u9fff]", sample):
        return "zh"
    ascii_letters = sum(1 for ch in sample if ("a" <= ch.lower() <= "z"))
    if sample and ascii_letters / max(len(sample), 1) > 0.45:
        return "en"
    return "unknown"


def keep_quality(row: Mapping[str, object], *, category: str, quality_keep: set[str]) -> bool:
    if category in {"code", "science"}:
        return True
    if category == "math":
        quality_keys = ("math_quality_bucket", "quality_bucket", "edu_quality_bucket")
    elif category == "english_edu":
        quality_keys = ("edu_quality_bucket", "quality_bucket", "math_quality_bucket")
    else:
        quality_keys = ("quality_bucket", "edu_quality_bucket", "math_quality_bucket")
    for key in quality_keys:
        quality = string_value(row.get(key)).strip().lower()
        if quality:
            return quality in quality_keep
    return "unknown" in quality_keep


def estimate_tokens(row: Mapping[str, object], text: str, *, category: str) -> int:
    raw_token_count = row.get("token_count")
    if isinstance(raw_token_count, int) and raw_token_count > 0:
        return raw_token_count
    if isinstance(raw_token_count, float) and raw_token_count > 0:
        return int(raw_token_count)
    return max(1, int(len(text.encode("utf-8", errors="replace")) / BYTES_PER_TOKEN[category]))


def build_output_row(
    row: Mapping[str, object],
    *,
    text: str,
    category: str,
    family: str,
    path: Path,
    row_index: int,
    estimated_tokens: int,
    seed: int,
) -> dict[str, object]:
    metadata = {
        key: value
        for key, value in row.items()
        if key not in set(OUTPUT_SCHEMA.names) and key != "text"
    }
    return {
        "text": text,
        "mixture_category": category,
        "mixture_family": family,
        "estimated_tokens": int(estimated_tokens),
        "source": string_value(row.get("source")),
        "subset": string_value(row.get("subset")),
        "language": normalize_language(string_value(row.get("language")), text),
        "quality_bucket": string_value(row.get("quality_bucket")),
        "edu_quality_bucket": string_value(row.get("edu_quality_bucket")),
        "math_quality_bucket": string_value(row.get("math_quality_bucket")),
        "corpus_source": string_value(row.get("corpus_source")),
        "source_path": str(path),
        "source_row": int(row_index),
        "input_format": file_format(path),
        "selection_seed": int(seed),
        "metadata_json": json.dumps(metadata, ensure_ascii=False, sort_keys=True),
    }


def write_report(
    path: Path,
    stats: BuildStats,
    specs: Sequence[SourceSpec],
    mix: Mapping[str, float],
    quotas: Mapping[str, int],
    effective_quotas: Mapping[str, int],
    args: argparse.Namespace,
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    actual_fraction = {
        category: stats.category_tokens.get(category, 0) / max(stats.estimated_tokens_written, 1)
        for category in mix
    }
    quota_fill = {
        category: stats.category_tokens.get(category, 0) / max(quota, 1)
        for category, quota in quotas.items()
    }
    effective_quota_fill = {
        category: stats.category_tokens.get(category, 0) / max(quota, 1)
        for category, quota in effective_quotas.items()
    }
    science_target = quotas.get(SHORTFALL_SOURCE, 0)
    science_actual = min(stats.category_tokens.get(SHORTFALL_SOURCE, 0), science_target)
    science_shortfall = max(0, science_target - science_actual)
    payload = {
        "recipe": "web40_edu20_chinese10_science10_logic10_math5_code5",
        "target_mix": dict(mix),
        "target_token_quotas": dict(quotas),
        "effective_token_quotas": dict(effective_quotas),
        "actual_fraction": actual_fraction,
        "quota_fill": quota_fill,
        "effective_quota_fill": effective_quota_fill,
        "shortfall_policy": {
            "enabled": bool(args.allow_science_shortfall_to_english_edu),
            "source": SHORTFALL_SOURCE,
            "target": SHORTFALL_TARGET,
            "science_target_tokens": science_target,
            "science_actual_tokens_capped": science_actual,
            "science_shortfall_tokens": science_shortfall,
            "english_edu_base_target_tokens": quotas.get(SHORTFALL_TARGET, 0),
            "english_edu_effective_target_tokens": effective_quotas.get(SHORTFALL_TARGET, 0),
        },
        "stats": stats.to_dict(),
        "inputs": [
            {
                "family": spec.family,
                "file_count": len(spec.paths),
                "preview": [str(path) for path in spec.paths[:10]],
            }
            for spec in specs
        ],
        "args": vars(args),
    }
    tmp_path = path.with_name(f".{path.name}.tmp.{os.getpid()}")
    tmp_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    tmp_path.replace(path)


def build_dry_run_report(
    specs: Sequence[SourceSpec],
    mix: Mapping[str, float],
    quotas: Mapping[str, int],
    args: argparse.Namespace,
) -> dict[str, object]:
    return {
        "dry_run": True,
        "recipe": "web40_edu20_chinese10_science10_logic10_math5_code5",
        "target_mix": dict(mix),
        "target_token_quotas": quotas,
        "inputs": [
            {
                "family": spec.family,
                "file_count": len(spec.paths),
                "preview": [str(path) for path in spec.paths[:10]],
            }
            for spec in specs
        ],
        "args": vars(args),
    }


def prepare_output_dir(path: Path, *, overwrite: bool) -> None:
    path.mkdir(parents=True, exist_ok=True)
    existing = list(path.glob("part_*.parquet")) + [path / "report.json"]
    existing = [item for item in existing if item.exists()]
    if not existing:
        return
    if not overwrite:
        raise FileExistsError(f"{existing[0]} already exists; pass --overwrite")
    for item in existing:
        item.unlink()


def file_format(path: Path) -> str:
    if path.name.lower().endswith(".jsonl.zst"):
        return "jsonl_zst"
    suffix = path.suffix.lower()
    if suffix == ".jsonl":
        return "jsonl"
    if suffix in {".parquet", ".pq"}:
        return "parquet"
    return ""


def infer_source(row: Mapping[str, object]) -> str:
    for key in ("source", "corpus_source", "subset"):
        value = string_value(row.get(key))
        if value:
            return value
    return "unknown"


def row_text(row: Mapping[str, object]) -> str:
    text = first_string(row, ("text", "content", "document", "prompt"))
    if text:
        return text
    question = first_string(row, ("question", "query"))
    context = first_string(row, ("context", "passage"))
    choices = format_choices(row.get("choices") or row.get("options")) or format_medmcqa_choices(row)
    answer = first_string(row, ("answer", "answerKey", "answer_key", "label", "solution", "cop"))
    explanation = first_string(row, ("explanation", "rationale", "exp"))
    parts = []
    if context:
        parts.append(f"Context:\n{context}")
    if question:
        parts.append(f"Question:\n{question}")
    if choices:
        parts.append(f"Choices:\n{choices}")
    if answer:
        parts.append(f"Answer:\n{answer}")
    if explanation:
        parts.append(f"Explanation:\n{explanation}")
    return "\n\n".join(parts)


def first_string(row: Mapping[str, object], keys: Sequence[str]) -> str:
    for key in keys:
        value = row.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip()
    return ""


def format_choices(value: object) -> str:
    if isinstance(value, str):
        return value.strip()
    if isinstance(value, Mapping):
        labels = value.get("label") or value.get("labels")
        texts = value.get("text") or value.get("texts")
        if isinstance(labels, Sequence) and not isinstance(labels, str) and isinstance(texts, Sequence) and not isinstance(texts, str):
            return "\n".join(f"{label}. {text}" for label, text in zip(labels, texts))
        return "\n".join(f"{key}. {val}" for key, val in value.items())
    if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
        return "\n".join(str(item) for item in value)
    return ""


def format_medmcqa_choices(row: Mapping[str, object]) -> str:
    pairs = (("A", "opa"), ("B", "opb"), ("C", "opc"), ("D", "opd"))
    lines = []
    for label, key in pairs:
        value = row.get(key)
        if isinstance(value, str) and value.strip():
            lines.append(f"{label}. {value.strip()}")
    return "\n".join(lines)


def string_value(value: object) -> str:
    return value if isinstance(value, str) else ""


def stable_int(value: str) -> int:
    digest = hashlib.sha256(value.encode("utf-8")).digest()
    return int.from_bytes(digest[:4], "big")


def should_log_progress(completed: int, total: int, every: int) -> bool:
    return completed == total or (every > 0 and completed % every == 0)


if __name__ == "__main__":
    main()
