from __future__ import annotations

import json
from pathlib import Path

import pyarrow as pa
from datasets import Dataset, DatasetDict, Features, Value


def convert_mixed_jsonl_to_records(repo_dir: str, output_dir: str) -> dict[str, int]:
    """
    Convert each JSONL file in repo_dir to a uniform schema where every row is
    {"record": <original_json_string>}. This avoids Hub viewer CastError when
    files have incompatible schemas.
    """
    repo_dir = Path(repo_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    features = Features({"record": Value("string")})
    counts: dict[str, int] = {}

    for path in sorted(repo_dir.glob("*.jsonl")):
        records = []
        with path.open("r", encoding="utf-8") as f:
            for line in f:
                record = json.loads(line)
                records.append({"record": json.dumps(record, ensure_ascii=False)})
        ds = Dataset.from_list(records, features=features)
        out_path = output_dir / path.name
        ds.to_json(out_path)
        counts[path.stem] = len(records)

    return counts


if __name__ == "__main__":
    counts = convert_mixed_jsonl_to_records(
        repo_dir="/home/hermes/fragrance-research/extracted",
        output_dir="/home/hermes/pino/data/literature_flat",
    )
    print(counts)