from __future__ import annotations

import argparse
from pathlib import Path
from typing import Iterable, Iterator

from .common import (
    make_user_assistant_record,
    read_jsonl,
    render_records,
    resolve_paths,
    write_jsonl,
)


def iter_records(
    paths: str | Path | Iterable[str | Path],
    *,
    include_source: bool = False,
) -> Iterator[dict[str, str]]:
    for path in resolve_paths(paths):
        for _, row in read_jsonl(path):
            record = make_user_assistant_record(
                row.get("user"),
                row.get("assistant"),
                source=str(path),
                include_source=include_source,
            )

            if record is not None:
                yield record


def load_records(
    paths: str | Path | Iterable[str | Path],
    *,
    include_source: bool = False,
) -> list[dict[str, str]]:
    return list(iter_records(paths, include_source=include_source))


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Normalize MultimodalReasoning JSONL files to user/assistant rows."
    )
    parser.add_argument("inputs", nargs="+", help="Input JSONL path(s).")
    parser.add_argument("--output", help="Optional output JSONL path.")
    parser.add_argument("--take", type=int, default=0, help="Print the first N rows.")
    parser.add_argument("--include-source", action="store_true")
    args = parser.parse_args()

    records = load_records(args.inputs, include_source=args.include_source)

    if args.output:
        write_jsonl(records, args.output)

    if args.take:
        print(render_records(records[: args.take]))


if __name__ == "__main__":
    main()