#!/usr/bin/env python3
"""Clean multi-vintage ('ctime') files in the data dump.

Some files carry an extra trailing 'ctime' column: each re-download of the
symbol APPENDED a full copy of the history stamped with the download time
(up to ~28 vintages per date, 2019-2022). The consumer keeps the LAST row
per date, so with mixed vintages it silently used a stale (~2021) vintage.

Fix, per file: for each date keep only the rows from the newest ctime
(several different amounts at the newest ctime are a legitimate same-day
pair and are all kept; exact (date, ctime) duplicates collapse to the last
row). The ctime column is dropped, rows sorted by date. Idempotent.

Also strips UTF-8 BOMs. Applies to the data root and overrides/frozen
(the frozen copies shadow the root and carry the same artifact).

Usage:  python3 scripts/clean_multivintage.py [--data DIR] [--frozen DIR]
        [--dry-run]
"""
import argparse
from pathlib import Path

ROOT_DEFAULT = Path("~/prog/fin/stocks").expanduser()
FROZEN_DEFAULT = Path(__file__).resolve().parent.parent / "overrides" / "frozen"
SUFFIXES = ("history", "dividend", "capitalGain", "split")


def clean_lines(lines: list[str], one_per_date: bool = False) -> list[str] | None:
    """Return cleaned lines, or None if the file has no ctime column.
    one_per_date (history/split): the newest-vintage row wins per date;
    event files keep every distinct amount at the newest ctime (a fund can
    genuinely pay several distributions the same day)."""
    header = lines[0].rstrip("\n").split(",")
    if "ctime" not in header:
        return None
    ci = header.index("ctime")
    body = [l.split(",") for l in lines[1:] if l.strip()]
    by_date: dict[str, list[tuple[str, int, list[str]]]] = {}
    for i, a in enumerate(body):
        if len(a) < 2:
            continue
        date = a[0][:10]
        ctime = a[ci] if ci < len(a) else ""
        vals = a[1:ci] + a[ci + 1:] if ci < len(a) else a[1:]
        by_date.setdefault(date, []).append((ctime, i, vals))
    out = []
    for date in sorted(by_date):
        rows = [r for r in by_date[date] if r[0] == max(x[0] for x in by_date[date])]
        if one_per_date:
            out.append(date + "," + ",".join(rows[-1][2]))
            continue
        # last occurrence of each distinct row (keeps legitimate same-day
        # pairs with different amounts, collapses exact repeats)
        last_pos = {}
        for j, (_, _, vals) in enumerate(rows):
            last_pos[tuple(vals)] = j
        for j in sorted(last_pos.values()):
            out.append(date + "," + ",".join(rows[j][2]))
    return [",".join(header[:ci] + header[ci + 1:])] + out


def process(dirs: list[Path], dry: bool) -> tuple[int, int]:
    n_files = n_rows = 0
    for d in dirs:
        if not d.is_dir():
            continue
        for sfx in SUFFIXES:
            for f in sorted(d.glob(f"*-{sfx}.csv")):
                raw = f.read_bytes()
                bom = raw.startswith(b"\xef\xbb\xbf")
                text = raw.decode("utf-8", "replace")
                lines = text.splitlines()
                out = clean_lines(lines, one_per_date=sfx in ("history", "split"))
                changed = out is not None or bom
                if not changed:
                    continue
                n_files += 1
                if out is None:
                    out = lines
                before = max(len(lines) - 1, 0)
                n_rows += before - max(len(out) - 1, 0)
                if not dry:
                    f.write_bytes(("\n".join(out) + "\n").encode("utf-8"))
    return n_files, n_rows


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", type=Path, default=ROOT_DEFAULT)
    ap.add_argument("--frozen", type=Path, default=FROZEN_DEFAULT)
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()
    nf, nr = process([args.data, args.frozen], args.dry_run)
    print(f"{'would clean' if args.dry_run else 'cleaned'} {nf} files, "
          f"{'would drop' if args.dry_run else 'dropped'} {nr} stale-vintage rows")
    return 0


if __name__ == "__main__":
    main()