#!/usr/bin/env python3
import argparse
import datetime as dt
import re
from decimal import Decimal
from typing import Any, Dict, List, Tuple

import pymysql

TABLES = ["categories", "teams", "engineers", "shift", "t_week_id", "t_day_id"]


def parse_args():
    p = argparse.ArgumentParser(description="Migrate old olympus dump into oncall-v2 schema")
    p.add_argument("--dump", required=True)
    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)
    p.add_argument("--apply", action="store_true")
    p.add_argument("--truncate", action="store_true")
    return p.parse_args()


def sanitize_binary_literals(s: str) -> str:
    out = []
    i = 0
    n = len(s)
    key = "_binary"
    while i < n:
        if s[i:i+len(key)].lower() == key:
            j = i + len(key)
            while j < n and s[j].isspace():
                j += 1
            if j < n and s[j] == "'":
                j += 1
                esc = False
                while j < n:
                    ch = s[j]
                    if esc:
                        esc = False
                    elif ch == '\\':
                        esc = True
                    elif ch == "'":
                        j += 1
                        break
                    j += 1
                out.append("NULL")
                i = j
                continue
        out.append(s[i])
        i += 1
    return ''.join(out)


def extract_insert_values(sql: str, table: str) -> str:
    pat = re.compile(rf"INSERT INTO `{table}` VALUES ", re.IGNORECASE)
    m = pat.search(sql)
    if not m:
        return ""
    i = m.end()
    # find ending semicolon not inside string
    in_str = False
    esc = False
    j = i
    while j < len(sql):
        ch = sql[j]
        if in_str:
            if esc:
                esc = False
            elif ch == '\\':
                esc = True
            elif ch == "'":
                in_str = False
        else:
            if ch == "'":
                in_str = True
            elif ch == ';':
                return sql[i:j]
        j += 1
    return ""


def split_tuples(values_blob: str) -> List[str]:
    out = []
    i = 0
    n = len(values_blob)
    while i < n:
        if values_blob[i] == '(':
            start = i + 1
            i += 1
            in_str = False
            esc = False
            depth = 1
            while i < n:
                ch = values_blob[i]
                if in_str:
                    if esc:
                        esc = False
                    elif ch == '\\':
                        esc = True
                    elif ch == "'":
                        in_str = False
                else:
                    if ch == "'":
                        in_str = True
                    elif ch == '(':
                        depth += 1
                    elif ch == ')':
                        depth -= 1
                        if depth == 0:
                            out.append(values_blob[start:i])
                            i += 1
                            break
                i += 1
        else:
            i += 1
    return out


def split_fields(tuple_body: str) -> List[str]:
    out = []
    cur = []
    in_str = False
    esc = False
    for ch in tuple_body:
        if in_str:
            cur.append(ch)
            if esc:
                esc = False
            elif ch == '\\':
                esc = True
            elif ch == "'":
                in_str = False
        else:
            if ch == "'":
                in_str = True
                cur.append(ch)
            elif ch == ',':
                out.append(''.join(cur).strip())
                cur = []
            else:
                cur.append(ch)
    out.append(''.join(cur).strip())
    return out


def parse_token(tok: str) -> Any:
    t = tok.strip()
    if t.upper() == 'NULL':
        return None
    if len(t) >= 2 and t[0] == "'" and t[-1] == "'":
        body = t[1:-1]
        body = body.replace("\\'", "'").replace('\\\\', '\\').replace('\\n', '\n').replace('\\r', '\r').replace('\\t', '\t')
        return body
    if re.fullmatch(r"-?\d+", t):
        return int(t)
    if re.fullmatch(r"-?\d+\.\d+", t):
        return Decimal(t)
    return t


def parse_table_rows(sql: str, table: str) -> List[List[Any]]:
    blob = extract_insert_values(sql, table)
    if not blob:
        return []
    if table == "engineers":
        blob = sanitize_binary_literals(blob)
    rows = []
    for tp in split_tuples(blob):
        fields = [parse_token(x) for x in split_fields(tp)]
        rows.append(fields)
    return rows


def normalize_schedule_type(v: Any) -> str:
    s = str(v or "").lower()
    if 'day' in s:
        return 'day'
    return 'week'


def to_unix(v: Any) -> int:
    if v is None:
        return 0
    s = str(v)
    try:
        d = dt.datetime.strptime(s, "%Y-%m-%d %H:%M:%S")
        return int(d.replace(tzinfo=dt.timezone.utc).timestamp())
    except Exception:
        return 0


def transform(data: Dict[str, List[List[Any]]]) -> Dict[str, List[Dict[str, Any]]]:
    categories = data["categories"]
    teams = data["teams"]
    engineers = data["engineers"]
    shifts = data["shift"]
    weeks = data["t_week_id"]
    days = data["t_day_id"]

    valid_cats: Dict[int, Dict[str, Any]] = {}
    for r in categories:
        if len(r) < 14:
            continue
        cid = int(r[0])
        level = int(r[2] or 0)
        enable = int(r[4] or 0)
        visibility = int(r[9] or 0)
        if not (level == 2 and enable == 1 and visibility == 1):
            continue
        cat_name = str(r[1] or "").strip()
        components = str(r[10] or "").strip() if len(r) > 10 else ""
        valid_cats[cid] = {
            "id": cid,
            "name": cat_name,
            "description": components or cat_name,
            "schedule_type": normalize_schedule_type(r[11]),
            "period": 1,
            "manager_emails": "",
            "group_name": "",
            "from_hour_utc": int(r[12] or 0),
            "to_hour_utc": int(r[13] or 23),
            "is_disabled": 0,
            "current_robin_index": 0,
        }

    team_to_category: Dict[int, int] = {}
    for r in teams:
        if len(r) < 3:
            continue
        tid = int(r[0])
        cid = int(r[2])
        if cid in valid_cats:
            team_to_category[tid] = cid

    out_engineers: Dict[int, Dict[str, Any]] = {}
    out_ec: Dict[Tuple[int, int], Dict[str, Any]] = {}
    for r in engineers:
        if len(r) < 10:
            continue
        eid = int(r[0])
        team_id = int(r[2] or 0)
        cid = team_to_category.get(team_id)
        if cid is None:
            continue
        name = str(r[1] or "").strip()
        en_name = str(r[8] or "").strip()
        email = str(r[9] or "").strip().lower()
        if not email:
            email = f"legacy-{eid}@example.com"
        out_engineers[eid] = {
            "id": eid,
            "name": name,
            "en_name": en_name,
            "email": email,
            "phone": "",
        }
        out_ec[(eid, cid)] = {
            "engineer_id": eid,
            "category_id": cid,
            "is_manager": 0,
        }

    # Ensure email uniqueness so all engineer IDs can be inserted (target has unique index on email).
    used_emails = set()
    for eid in sorted(out_engineers.keys()):
        email = out_engineers[eid]["email"]
        if email not in used_emails:
            used_emails.add(email)
            continue

        if "@" in email:
            local, domain = email.split("@", 1)
            new_email = f"{local}+legacy-{eid}@{domain}"
        else:
            new_email = f"legacy-{eid}@example.com"

        # extra guard in case generated email also collides
        while new_email in used_emails:
            new_email = f"legacy-{eid}-{len(used_emails)}@example.com"

        out_engineers[eid]["email"] = new_email
        used_emails.add(new_email)

    wk_meta: Dict[Decimal, Dict[str, int]] = {}
    for r in weeks:
        if len(r) < 3:
            continue
        wk = Decimal(str(r[0]))
        wk_meta[wk] = {"start_time": to_unix(r[1]), "end_time": to_unix(r[2])}

    day_meta: Dict[Decimal, Dict[str, int]] = {}
    for r in days:
        if len(r) < 3:
            continue
        dk = Decimal(str(r[0]))
        day_meta[dk] = {"start_time": to_unix(r[1]), "end_time": to_unix(r[2])}

    valid_shifts = []
    keys_by_type = {"week": set(), "day": set()}
    for r in shifts:
        if len(r) < 7:
            continue
        cid = int(r[1])
        eid = int(r[2])
        period_key = Decimal(str(r[3]))
        manual = int(r[5] or 0)
        enable = int(r[6] or 0)
        if enable != 1 or cid not in valid_cats or eid not in out_engineers:
            continue
        ptype = valid_cats[cid]["schedule_type"]
        valid_shifts.append((cid, eid, period_key, ptype, 1 if manual else 0))
        keys_by_type[ptype].add(period_key)

    ordered_weeks = sorted(keys_by_type["week"])
    ordered_days = sorted(keys_by_type["day"])
    time_periods = []
    period_to_tp: Dict[Tuple[str, Decimal], int] = {}
    next_id = 1

    for wk in ordered_weeks:
        meta = wk_meta.get(wk, {"start_time": 0, "end_time": 0})
        period_to_tp[("week", wk)] = next_id
        time_periods.append({
            "id": next_id,
            "period_type": "week",
            "name": f"Week {wk}",
            "start_time": meta["start_time"],
            "end_time": meta["end_time"],
            "sequence": next_id,
            "manual_only": 0,
            "is_special": 0,
            "merge_anchor_sequence": 0,
        })
        next_id += 1

    for dk in ordered_days:
        meta = day_meta.get(dk, {"start_time": 0, "end_time": 0})
        period_to_tp[("day", dk)] = next_id
        time_periods.append({
            "id": next_id,
            "period_type": "day",
            "name": f"Day {dk}",
            "start_time": meta["start_time"],
            "end_time": meta["end_time"],
            "sequence": next_id,
            "manual_only": 0,
            "is_special": 0,
            "merge_anchor_sequence": 0,
        })
        next_id += 1

    out_shifts = []
    for cid, eid, pkey, ptype, manual in valid_shifts:
        tp_id = period_to_tp.get((ptype, pkey))
        if tp_id is None:
            continue
        out_shifts.append({
            "category_id": cid,
            "engineer_id": eid,
            "time_period_id": tp_id,
            "manual_flag": manual,
        })

    return {
        "categories": list(valid_cats.values()),
        "engineers": list(out_engineers.values()),
        "engineer_categories": list(out_ec.values()),
        "time_periods": time_periods,
        "shifts": out_shifts,
    }


def connect(args):
    return 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": {}},
    )


def apply_data(args, data):
    conn = connect(args)
    try:
        with conn.cursor() as cur:
            if args.truncate:
                cur.execute("SET FOREIGN_KEY_CHECKS=0")
                cur.execute("TRUNCATE TABLE shifts")
                cur.execute("TRUNCATE TABLE category_robins")
                cur.execute("TRUNCATE TABLE engineer_categories")
                cur.execute("TRUNCATE TABLE time_periods")
                cur.execute("TRUNCATE TABLE engineers")
                cur.execute("TRUNCATE TABLE categories")
                cur.execute("SET FOREIGN_KEY_CHECKS=1")
                conn.commit()

            categories_inserted = 0
            for c in data["categories"]:
                cur.execute(
                    """
                    INSERT INTO categories (id,name,description,schedule_type,period,manager_emails,group_name,from_hour_utc,to_hour_utc,is_disabled,current_robin_index)
                    VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
                    ON DUPLICATE KEY UPDATE
                      name=VALUES(name),description=VALUES(description),schedule_type=VALUES(schedule_type),period=VALUES(period),manager_emails=VALUES(manager_emails),
                      group_name=VALUES(group_name),from_hour_utc=VALUES(from_hour_utc),to_hour_utc=VALUES(to_hour_utc),is_disabled=VALUES(is_disabled)
                    """,
                    (c["id"], c["name"], c["description"], c["schedule_type"], c["period"], c["manager_emails"], c["group_name"], c["from_hour_utc"], c["to_hour_utc"], c["is_disabled"], c["current_robin_index"]),
                )
                categories_inserted += cur.rowcount
            conn.commit()
            print(f"categories_done rows={categories_inserted}", flush=True)

            engineers_inserted = 0
            for e in data["engineers"]:
                cur.execute(
                    """
                    INSERT INTO engineers (id,name,en_name,email,phone)
                    VALUES (%s,%s,%s,%s,%s)
                    ON DUPLICATE KEY UPDATE name=VALUES(name),en_name=VALUES(en_name),phone=VALUES(phone)
                    """,
                    (e["id"], e["name"], e["en_name"], e["email"], e["phone"]),
                )
                engineers_inserted += cur.rowcount
            conn.commit()
            print(f"engineers_done rows={engineers_inserted}", flush=True)

            ec_inserted = 0
            for ec in data["engineer_categories"]:
                cur.execute(
                    """
                    INSERT INTO engineer_categories (engineer_id,category_id,is_manager)
                    SELECT %s,%s,%s
                    FROM DUAL
                    WHERE EXISTS (SELECT 1 FROM engineers e WHERE e.id = %s)
                      AND EXISTS (SELECT 1 FROM categories c WHERE c.id = %s)
                    """,
                    (ec["engineer_id"], ec["category_id"], ec["is_manager"], ec["engineer_id"], ec["category_id"]),
                )
                ec_inserted += cur.rowcount
            conn.commit()
            print(f"engineer_categories_done rows={ec_inserted}", flush=True)

            tp_inserted = 0
            for tp in data["time_periods"]:
                cur.execute(
                    """
                    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)
                    ON DUPLICATE KEY UPDATE
                      period_type=VALUES(period_type),name=VALUES(name),start_time=VALUES(start_time),end_time=VALUES(end_time),
                      sequence=VALUES(sequence),manual_only=VALUES(manual_only),is_special=VALUES(is_special),merge_anchor_sequence=VALUES(merge_anchor_sequence)
                    """,
                    (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"]),
                )
                tp_inserted += cur.rowcount
            conn.commit()
            print(f"time_periods_done rows={tp_inserted}", flush=True)

            shift_inserted = 0
            for s in data["shifts"]:
                cur.execute(
                    """
                    INSERT INTO shifts (category_id,engineer_id,time_period_id,manual_flag)
                    SELECT %s,%s,%s,%s
                    FROM DUAL
                    WHERE EXISTS (SELECT 1 FROM categories c WHERE c.id = %s)
                      AND EXISTS (SELECT 1 FROM engineers e WHERE e.id = %s)
                      AND EXISTS (SELECT 1 FROM time_periods tp WHERE tp.id = %s)
                    """,
                    (s["category_id"], s["engineer_id"], s["time_period_id"], s["manual_flag"], s["category_id"], s["engineer_id"], s["time_period_id"]),
                )
                shift_inserted += cur.rowcount
            conn.commit()
            print(f"shifts_done rows={shift_inserted}", flush=True)
    finally:
        conn.close()


def main():
    args = parse_args()
    with open(args.dump, "r", encoding="utf-8", errors="surrogateescape") as f:
        sql = f.read()

    parsed = {t: parse_table_rows(sql, t) for t in TABLES}
    data = transform(parsed)

    print("Preview:")
    print("  categories:", len(data["categories"]))
    print("  engineers:", len(data["engineers"]))
    print("  engineer_categories:", len(data["engineer_categories"]))
    print("  time_periods:", len(data["time_periods"]))
    print("  shifts:", len(data["shifts"]))

    if not args.apply:
        print("Dry run only. Add --apply to execute.")
        return

    apply_data(args, data)
    print("Applied.")


if __name__ == "__main__":
    main()