"""Convert raw CSV/Parquet into qlib-compatible per-symbol CSV, then dump to .bin format."""

from __future__ import annotations

import argparse
from pathlib import Path

import pandas as pd

from config.settings import load_settings


REQUIRED_QLIB_FIELDS = ["open", "close", "high", "low", "volume", "factor"]


def normalize_columns(df: pd.DataFrame, aliases: dict[str, str]) -> pd.DataFrame:
    rename = {}
    for col in df.columns:
        key = str(col).strip()
        if key in aliases:
            rename[col] = aliases[key]
    out = df.rename(columns=rename)
    return out.loc[:, ~out.columns.duplicated()]


def parquet_to_symbol_csv(
    parquet_path: Path,
    output_dir: Path,
    aliases: dict[str, str],
    chunk_size: int = 500_000,
) -> int:
    """Split a wide parquet into per-symbol CSV files for dump_bin."""
    import pyarrow.parquet as pq

    output_dir.mkdir(parents=True, exist_ok=True)
    pf = pq.ParquetFile(parquet_path)
    writers: dict[str, list[pd.DataFrame]] = {}

    for batch in pf.iter_batches(batch_size=chunk_size):
        chunk = batch.to_pandas()
        chunk = normalize_columns(chunk, aliases)
        if "date" not in chunk.columns or "symbol" not in chunk.columns:
            raise KeyError(f"Missing date/symbol after alias map. columns={chunk.columns.tolist()}")

        chunk["date"] = pd.to_datetime(chunk["date"])
        chunk["symbol"] = chunk["symbol"].astype(str).str.lower()
        if "factor" not in chunk.columns:
            chunk["factor"] = 1.0

        for sym, grp in chunk.groupby("symbol"):
            writers.setdefault(sym, []).append(grp)

    count = 0
    for sym, parts in writers.items():
        df = pd.concat(parts, ignore_index=True)
        df = df.sort_values("date").drop_duplicates("date")
        out_path = output_dir / f"{sym}.csv"
        keep = ["date", "symbol"] + [c for c in REQUIRED_QLIB_FIELDS if c in df.columns]
        df[keep].to_csv(out_path, index=False)
        count += 1
    return count


def dump_csv_dir_to_qlib(
    csv_dir: Path,
    qlib_dir: Path,
    include_fields: str,
    freq: str = "day",
    max_workers: int = 8,
    mode: str = "dump_all",
):
    from scripts.dump_bin import DumpDataAll, DumpDataUpdate

    kwargs = dict(
        data_path=str(csv_dir),
        qlib_dir=str(qlib_dir),
        include_fields=include_fields,
        file_suffix=".csv",
        date_field_name="date",
        symbol_field_name="symbol",
        freq=freq,
        max_workers=max_workers,
    )
    if mode == "dump_update":
        DumpDataUpdate(**kwargs)()
    else:
        DumpDataAll(**kwargs)()


def run_from_parquet(parquet_path: Path | None = None, mode: str = "dump_all") -> Path:
    settings = load_settings()
    dump_cfg = settings.dump_config

    csv_dir = settings.path(dump_cfg.get("raw_data_dir", "data/raw/csv_by_symbol"))
    qlib_dir = settings.path(dump_cfg.get("qlib_dir", "data/qlib_data/cn_data"))
    aliases = dump_cfg.get("column_aliases", {})

    if parquet_path is not None:
        n = parquet_to_symbol_csv(Path(parquet_path), csv_dir, aliases)
        print(f"Converted {n} symbols to {csv_dir}")

    dump_csv_dir_to_qlib(
        csv_dir=csv_dir,
        qlib_dir=qlib_dir,
        include_fields=dump_cfg.get("include_fields", "open,close,high,low,volume,factor"),
        freq=dump_cfg.get("freq", "day"),
        max_workers=int(dump_cfg.get("max_workers", 8)),
        mode=mode,
    )
    print(f"Qlib bin data ready at: {qlib_dir}")
    return qlib_dir


if __name__ == "__main__":
    import sys

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

    parser = argparse.ArgumentParser(description="Convert parquet/csv to qlib bin format")
    parser.add_argument("--parquet", type=str, default=None, help="Source parquet path")
    parser.add_argument("--mode", choices=["dump_all", "dump_update"], default="dump_all")
    args = parser.parse_args()
    run_from_parquet(args.parquet, mode=args.mode)