#!/usr/bin/env python3
import argparse
import importlib.util
from pathlib import Path
import pymysql


def load_migrator(path: str):
    spec = importlib.util.spec_from_file_location("migrator", path)
    mod = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(mod)
    return mod


def chunks(lst, n):
    for i in range(0, len(lst), n):
        yield lst[i:i+n]


def main():
    p = argparse.ArgumentParser(description="One-time fill time_periods + shifts from old olympus dump")
    p.add_argument("--dump", required=True)
    p.add_argument("--migrator", default="/Users/hu/aisupport/scripts/bootstrap/migrate_old_olympus_to_oncall_v2.py")
    p.add_argument("--host", required=True)
    p.add_argument("--port", type=int, default=4000)
    p.add_argument("--user", required=True)
    p.add_argument("--password", required=True)
    p.add_argument("--database", required=True)
    args = p.parse_args()

    mod = load_migrator(args.migrator)

    print("loading dump...", flush=True)
    sql = Path(args.dump).read_text(encoding="utf-8", errors="surrogateescape")
    parsed = {t: mod.parse_table_rows(sql, t) for t in mod.TABLES}
    data = mod.transform(parsed)
    tps = data["time_periods"]
    shifts = data["shifts"]

    print(f"parsed: time_periods={len(tps)} shifts={len(shifts)}", flush=True)

    conn = pymysql.connect(
        host=args.host,
        port=args.port,
        user=args.user,
        password=args.password,
        database=args.database,
        charset="utf8mb4",
        autocommit=False,
        cursorclass=pymysql.cursors.DictCursor,
        ssl={"ssl": {}},
    )
    try:
        with conn.cursor() as cur:
            print("cleanup existing shifts/time_periods...", flush=True)
            cur.execute("DELETE FROM shifts")
            cur.execute("DELETE FROM time_periods")
            conn.commit()

            print("insert time_periods...", flush=True)
            tp_sql = (
                "INSERT INTO time_periods "
                "(id,period_type,name,start_time,end_time,sequence,manual_only,is_special,merge_anchor_sequence) "
                "VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)"
            )
            tp_rows = [
                (
                    tp["id"], tp["period_type"], tp["name"], tp["start_time"], tp["end_time"], tp["sequence"],
                    tp["manual_only"], tp["is_special"], tp["merge_anchor_sequence"],
                )
                for tp in tps
            ]
            inserted_tp = 0
            for batch in chunks(tp_rows, 500):
                cur.executemany(tp_sql, batch)
                inserted_tp += len(batch)
                conn.commit()
            print(f"time_periods_done rows={inserted_tp}", flush=True)

            print("insert shifts...", flush=True)
            sh_sql = "INSERT INTO shifts (category_id,engineer_id,time_period_id,manual_flag) VALUES (%s,%s,%s,%s)"
            sh_rows = [
                (s["category_id"], s["engineer_id"], s["time_period_id"], s["manual_flag"])
                for s in shifts
            ]
            inserted_sh = 0
            for batch in chunks(sh_rows, 1000):
                cur.executemany(sh_sql, batch)
                inserted_sh += len(batch)
                conn.commit()
            print(f"shifts_done rows={inserted_sh}", flush=True)

        print("done", flush=True)
    finally:
        conn.close()


if __name__ == "__main__":
    main()