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

import pymysql

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

# daily support categories + related compute/na/tools set user requested
TARGET_CATEGORY_IDS = [49, 50, 52, 54, 57, 58, 59, 63, 71, 75]


def parse_args():
    p = argparse.ArgumentParser(description="Import only category_robins from old olympus dump with engineers.available=1 filter")
    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("--create-db-if-missing", action="store_true")
    p.add_argument("--only-target-categories", action="store_true", default=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()
    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 build_robin_from_shift(parsed: Dict[str, List[List[Any]]], only_targets: bool = True):
    categories = parsed["categories"]
    engineers = parsed["engineers"]
    shifts = parsed["shift"]

    # category filter:
    # - default follows migrator (L2 + enabled + visible)
    # - when only_targets=True, explicitly include target IDs even if level!=2 (e.g. 49/50)
    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 only_targets:
            if cid not in TARGET_CATEGORY_IDS:
                continue
            if not (enable == 1 and visibility == 1):
                continue
        else:
            if not (level == 2 and enable == 1 and visibility == 1):
                continue

        valid_cats[cid] = {
            "id": cid,
            "name": str(r[1] or "").strip(),
            "schedule_type": normalize_schedule_type(r[11]),
            "period": 1,
        }

    # available=1 only
    engineer_available: Dict[int, int] = {}
    for r in engineers:
        if len(r) < 8:
            continue
        eid = int(r[0])
        available = int(r[7] or 0)
        engineer_available[eid] = available

    # collect (category,start_week)->engineers
    by_cat_start: Dict[int, Dict[Decimal, List[int]]] = defaultdict(lambda: defaultdict(list))
    for r in shifts:
        if len(r) < 7:
            continue
        cid = int(r[1])
        eid = int(r[2])
        start_week = Decimal(str(r[4]))
        enable = int(r[6] or 0)
        if cid not in valid_cats:
            continue
        if enable != 1:
            continue
        if engineer_available.get(eid, 0) != 1:
            continue
        if eid not in by_cat_start[cid][start_week]:
            by_cat_start[cid][start_week].append(eid)

    # convert to step_order list by start_week ascending
    robin_rows: List[Tuple[int, int, int]] = []  # (category_id, engineer_id, step_order)
    steps_count: Dict[int, int] = {}
    for cid in sorted(valid_cats.keys()):
        starts = sorted(by_cat_start.get(cid, {}).keys())
        step_order = 0
        for sw in starts:
            eids = by_cat_start[cid][sw]
            for eid in sorted(eids):
                robin_rows.append((cid, eid, step_order))
            step_order += 1
        steps_count[cid] = step_order

    return valid_cats, robin_rows, steps_count


def ensure_db(host, port, user, password, database):
    conn = pymysql.connect(
        host=host,
        port=port,
        user=user,
        password=password,
        charset="utf8mb4",
        autocommit=True,
        cursorclass=pymysql.cursors.DictCursor,
        ssl={"ssl": {}},
    )
    try:
        with conn.cursor() as cur:
            cur.execute(f"CREATE DATABASE IF NOT EXISTS `{database}` CHARACTER SET utf8mb4 COLLATE utf8mb4_bin")
    finally:
        conn.close()


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 prepare_schema(conn):
    with conn.cursor() as cur:
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS categories (
              id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,
              name LONGTEXT,
              schedule_type LONGTEXT,
              period BIGINT,
              manager_emails LONGTEXT,
              group_name LONGTEXT,
              from_hour_utc BIGINT DEFAULT 0,
              to_hour_utc BIGINT DEFAULT 23,
              is_disabled TINYINT(1),
              current_robin_index BIGINT,
              PRIMARY KEY (id)
            )
            """
        )
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS engineers (
              id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,
              name LONGTEXT,
              en_name LONGTEXT,
              email VARCHAR(191),
              phone LONGTEXT,
              PRIMARY KEY (id),
              UNIQUE KEY idx_engineers_email (email)
            )
            """
        )
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS category_robins (
              id BIGINT UNSIGNED NOT NULL AUTO_INCREMENT,
              category_id BIGINT UNSIGNED,
              engineer_id BIGINT UNSIGNED,
              step_order BIGINT,
              PRIMARY KEY (id)
            )
            """
        )
    conn.commit()


def apply(args, valid_cats, robin_rows):
    conn = connect(args)
    try:
        prepare_schema(conn)
        with conn.cursor() as cur:
            # ensure categories exist with sensible defaults
            for cid, c in valid_cats.items():
                cur.execute(
                    """
                    INSERT INTO categories (id,name,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)
                    ON DUPLICATE KEY UPDATE
                      name=VALUES(name), schedule_type=VALUES(schedule_type), period=VALUES(period),
                      is_disabled=COALESCE(is_disabled,0)
                    """,
                    (cid, c["name"], c["schedule_type"], c["period"], "admin@example.com", "", 0, 23, 0, 0),
                )

            # ensure engineers exist
            eids = sorted({eid for _, eid, _ in robin_rows})
            for eid in eids:
                cur.execute(
                    """
                    INSERT INTO engineers (id,name,en_name,email,phone)
                    VALUES (%s,%s,%s,%s,%s)
                    ON DUPLICATE KEY UPDATE id=id
                    """,
                    (eid, f"legacy-{eid}", f"legacy-{eid}", f"legacy-{eid}@example.com", ""),
                )

            # replace robin for target categories only
            target_ids = sorted(valid_cats.keys())
            if target_ids:
                fmt = ",".join(["%s"] * len(target_ids))
                cur.execute(f"DELETE FROM category_robins WHERE category_id IN ({fmt})", target_ids)

            if robin_rows:
                cur.executemany(
                    "INSERT INTO category_robins (category_id,engineer_id,step_order) VALUES (%s,%s,%s)",
                    robin_rows,
                )

        conn.commit()
    finally:
        conn.close()


def verify(args, target_ids):
    conn = connect(args)
    try:
        with conn.cursor() as cur:
            fmt = ",".join(["%s"] * len(target_ids))
            cur.execute(
                f"""
                SELECT c.id, c.name, c.schedule_type, c.period,
                       COUNT(cr.id) AS robin_rows,
                       COUNT(DISTINCT cr.step_order) AS step_count
                FROM categories c
                LEFT JOIN category_robins cr ON cr.category_id = c.id
                WHERE c.id IN ({fmt})
                GROUP BY c.id,c.name,c.schedule_type,c.period
                ORDER BY c.id
                """,
                target_ids,
            )
            rows = cur.fetchall()
            return rows
    finally:
        conn.close()


def main():
    args = parse_args()

    if args.create_db_if_missing:
        ensure_db(args.host, args.port, args.user, args.password, args.database)

    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}
    valid_cats, robin_rows, steps_count = build_robin_from_shift(parsed, only_targets=args.only_target_categories)

    print("Preview:")
    print("  categories(target):", len(valid_cats))
    print("  robin_rows:", len(robin_rows))
    for cid in sorted(valid_cats.keys()):
        print(f"  - {cid} {valid_cats[cid]['name']} steps={steps_count.get(cid,0)}")

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

    apply(args, valid_cats, robin_rows)
    print("Applied robin import.")

    report = verify(args, sorted(valid_cats.keys()))
    print("Verify:")
    for r in report:
        print(f"  {r['id']}\t{r['name']}\t{r['schedule_type']}\tperiod={r['period']}\trobin_rows={r['robin_rows']}\tsteps={r['step_count']}")


if __name__ == "__main__":
    main()