from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.parquet as pq

from .dataset_builder import DatasetBuilder

logger = logging.getLogger("pino.factory")

GENRES = ["citrus_cologne", "fougere", "floral_woody", "amber_oriental", "wildcard"]


def _record_to_parquet_row(record: dict[str, Any]) -> dict[str, Any]:
    """Flatten a record into Parquet-compatible columns."""
    metadata = record.get("metadata", {})
    return {
        "formula_id": record.get("formula_id", ""),
        "generation_strategy": metadata.get("generation_strategy", ""),
        "status": record.get("status", ""),
        "active_components_count": metadata.get("active_components_count", 0),
        "formula": json.dumps(record.get("formula", [])),
        "depletion_rates": json.dumps(record.get("depletion_rates", {})),
        "trajectory": json.dumps(record.get("trajectory", [])),
    }


def write_jsonl_to_parquet(jsonl_path: Path, parquet_path: Path) -> int:
    """Convert an existing JSON-L dataset to Parquet."""
    rows = []
    with jsonl_path.open("r") as f:
        for line in f:
            try:
                record = json.loads(line)
                rows.append(_record_to_parquet_row(record))
            except Exception:
                continue

    schema = pa.schema([
        pa.field("formula_id", pa.string()),
        pa.field("generation_strategy", pa.string()),
        pa.field("status", pa.string()),
        pa.field("active_components_count", pa.int64()),
        pa.field("formula", pa.string()),
        pa.field("depletion_rates", pa.string()),
        pa.field("trajectory", pa.string()),
    ])

    arrays = {k: [r[k] for r in rows] for k in rows[0].keys()}
    table = pa.table(arrays, schema=schema)
    pq.write_table(table, parquet_path)
    return len(rows)


def main() -> None:
    parser = argparse.ArgumentParser(description="PINO Synthetic Data Factory")
    parser.add_argument("--output", default="data/trajectories.parquet", help="Output Parquet path")
    parser.add_argument("--count", type=int, default=5000, help="Total number of records")
    parser.add_argument("--rules", default="data/genre_rules.json")
    parser.add_argument("--registry", default="src/pino/registry.db")
    parser.add_argument("--workers", type=int, default=16)
    parser.add_argument("--seed", type=int, default=2026)
    parser.add_argument("--from-jsonl", default="data/synthetic_dataset_v2.jsonl", help="Convert existing JSON-L instead of regenerating")
    parser.add_argument("--regenerate", action="store_true", help="Regenerate from scratch")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    jsonl_path = Path(args.from_jsonl)
    if args.regenerate or not jsonl_path.exists() or jsonl_path.stat().st_size == 0:
        per_genre = args.count // len(GENRES)
        builder = DatasetBuilder(
            output_path=jsonl_path,
            rules_path=args.rules,
            registry_path=args.registry,
            per_genre=per_genre,
            max_workers=args.workers,
            seed=args.seed,
        )
        builder.build()

    count = write_jsonl_to_parquet(jsonl_path, output_path)
    logger.info("Wrote %d records to %s", count, output_path)


if __name__ == "__main__":
    main()