from __future__ import annotations

import argparse
import json
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_INPUT = ROOT / "plugins" / "1c" / "training" / "examples" / "instruction.examples.jsonl"
DEFAULT_GENERATED_INPUT = ROOT / "plugins" / "1c" / "training" / "raw" / "generated.instruction.jsonl"
DEFAULT_OUTPUT = ROOT / "plugins" / "1c" / "training" / "prepared" / "train.chat.jsonl"


def iter_jsonl(path: Path) -> list[dict]:
    records: list[dict] = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


def normalize_record(record: dict) -> dict:
    messages = []
    for message in record.get("messages") or []:
        messages.append(
            {
                "role": message["role"],
                "content": message["content"].strip(),
            }
        )
    return {
        "messages": messages,
        "metadata": {
            "id": record.get("id"),
            **(record.get("metadata") or {}),
        },
    }


def dedupe_records(records: list[dict]) -> list[dict]:
    seen: set[str] = set()
    result = []
    for record in records:
        record_id = str(record.get("id") or (record.get("metadata") or {}).get("id") or "")
        if record_id and record_id in seen:
            continue
        if record_id:
            seen.add(record_id)
        result.append(record)
    return result


def main() -> int:
    parser = argparse.ArgumentParser(description="Prepare 1C training data in chat JSONL format.")
    parser.add_argument("--input", type=Path, action="append", help="Input JSONL file. Can be passed multiple times.")
    parser.add_argument(
        "--include-generated",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Include the generated local training JSONL when it exists.",
    )
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    args = parser.parse_args()

    input_paths = args.input or [DEFAULT_INPUT]
    if args.include_generated and DEFAULT_GENERATED_INPUT.exists() and DEFAULT_GENERATED_INPUT not in input_paths:
        input_paths.append(DEFAULT_GENERATED_INPUT)

    raw_records = []
    for input_path in input_paths:
        raw_records.extend(iter_jsonl(input_path))

    records = [normalize_record(record) for record in dedupe_records(raw_records)]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")

    print(f"Wrote {len(records)} training record(s) to {args.output}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())