#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
迁移 customer_db/customers.json -> MySQL customer_profiles
"""

from __future__ import annotations

import argparse
import json
from datetime import datetime
from pathlib import Path

import pymysql


def get_conn(host: str, port: int, user: str, password: str, database: str):
    return pymysql.connect(
        host=host,
        port=port,
        user=user,
        password=password,
        database=database,
        charset="utf8mb4",
        autocommit=False,
        cursorclass=pymysql.cursors.DictCursor,
    )


def ensure_table(conn):
    with conn.cursor() as cur:
        cur.execute(
            """
            CREATE TABLE IF NOT EXISTS customer_profiles (
                customer_id VARCHAR(128) PRIMARY KEY,
                profile_json LONGTEXT NOT NULL,
                last_update DATETIME NOT NULL
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
            """
        )
        cur.execute("SHOW INDEX FROM customer_profiles")
        exists = {str(r.get("Key_name", "")) for r in cur.fetchall()}
        if "idx_last_update" not in exists:
            cur.execute("CREATE INDEX idx_last_update ON customer_profiles(last_update)")
    conn.commit()


def migrate(json_path: Path, host: str, port: int, user: str, password: str, database: str, truncate_target: bool):
    if not json_path.exists():
        raise FileNotFoundError(f"customers.json 不存在: {json_path}")
    customers = json.loads(json_path.read_text(encoding="utf-8") or "{}")
    if not isinstance(customers, dict):
        raise RuntimeError("customers.json 格式错误，期望对象映射")

    conn = get_conn(host, port, user, password, database)
    try:
        ensure_table(conn)
        if truncate_target:
            with conn.cursor() as cur:
                cur.execute("TRUNCATE TABLE customer_profiles")
            conn.commit()

        now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        sql = (
            "REPLACE INTO customer_profiles (customer_id, profile_json, last_update) "
            "VALUES (%s, %s, %s)"
        )
        total = 0
        with conn.cursor() as cur:
            for cid, profile in customers.items():
                cur.execute(sql, (str(cid), json.dumps(profile, ensure_ascii=False), now))
                total += 1
        conn.commit()
        return total
    finally:
        conn.close()


def main():
    parser = argparse.ArgumentParser(description="迁移 customers.json 到 MySQL")
    parser.add_argument("--json-path", default=str(Path("customer_db") / "customers.json"))
    parser.add_argument("--host", required=True)
    parser.add_argument("--port", type=int, default=3306)
    parser.add_argument("--user", required=True)
    parser.add_argument("--password", required=True)
    parser.add_argument("--database", required=True)
    parser.add_argument("--truncate-target", action="store_true")
    args = parser.parse_args()

    total = migrate(
        json_path=Path(args.json_path),
        host=args.host,
        port=args.port,
        user=args.user,
        password=args.password,
        database=args.database,
        truncate_target=bool(args.truncate_target),
    )
    print(f"[DONE] customer_profiles 写入 {total} 条")


if __name__ == "__main__":
    main()