from __future__ import annotations

import argparse
import getpass
import json
import os
import re
from pathlib import Path

import pymysql


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_SEED_DIR = ROOT / "data" / "mysql_seed_20260624"


def split_sql(sql: str) -> list[str]:
    parts: list[str] = []
    buff: list[str] = []
    for line in sql.splitlines():
        buff.append(line)
        if line.strip().endswith(";"):
            stmt = "\n".join(buff).strip()
            if stmt:
                parts.append(stmt[:-1] if stmt.endswith(";") else stmt)
            buff = []
    tail = "\n".join(buff).strip()
    if tail:
        parts.append(tail)
    return parts


def get_password(args: argparse.Namespace) -> str:
    if args.password:
        return args.password
    if os.environ.get("DARKLINE_MYSQL_PASSWORD"):
        return os.environ["DARKLINE_MYSQL_PASSWORD"]
    return getpass.getpass("MySQL password: ")


def connect_without_db(args: argparse.Namespace, password: str):
    return pymysql.connect(
        host=args.host,
        port=args.port,
        user=args.user,
        password=password,
        charset="utf8mb4",
        autocommit=True,
    )


def connect_with_db(args: argparse.Namespace, password: str):
    return pymysql.connect(
        host=args.host,
        port=args.port,
        user=args.user,
        password=password,
        database=args.database,
        charset="utf8mb4",
        autocommit=True,
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Import darkline seed data into local MySQL.")
    parser.add_argument("--host", default=os.environ.get("DARKLINE_MYSQL_HOST", "127.0.0.1"))
    parser.add_argument("--port", type=int, default=int(os.environ.get("DARKLINE_MYSQL_PORT", "3306")))
    parser.add_argument("--user", default=os.environ.get("DARKLINE_MYSQL_USER", "root"))
    parser.add_argument("--password", default=None)
    parser.add_argument("--database", default=os.environ.get("DARKLINE_MYSQL_DATABASE", "tianxia"))
    parser.add_argument("--seed-dir", default=str(DEFAULT_SEED_DIR))
    parser.add_argument("--create-database", action="store_true")
    args = parser.parse_args()

    seed_dir = Path(args.seed_dir)
    password = get_password(args)

    if args.create_database:
        conn = connect_without_db(args, password)
        with conn.cursor() as cur:
            cur.execute(
                f"CREATE DATABASE IF NOT EXISTS `{args.database}` "
                "CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci"
            )
        conn.close()

    conn = connect_with_db(args, password)
    with conn.cursor() as cur:
        schema_sql = (seed_dir / "schema.sql").read_text(encoding="utf-8")
        for stmt in split_sql(schema_sql):
            cur.execute(stmt)

        table_order = json.loads((seed_dir / "table_order.json").read_text(encoding="utf-8"))
        imported_counts = {}
        for table in table_order:
            jsonl_path = seed_dir / "tables" / f"{table}.jsonl"
            rows = [
                json.loads(line)
                for line in jsonl_path.read_text(encoding="utf-8").splitlines()
                if line.strip()
            ]
            if not rows:
                imported_counts[table] = 0
                continue
            columns = list(rows[0].keys())
            col_sql = ", ".join(f"`{c}`" for c in columns)
            placeholders = ", ".join(["%s"] * len(columns))
            sql = f"REPLACE INTO `{table}` ({col_sql}) VALUES ({placeholders})"
            cur.executemany(sql, [[row.get(c) for c in columns] for row in rows])
            imported_counts[table] = len(rows)

        verify_counts = {}
        for table in table_order:
            cur.execute(f"SELECT COUNT(*) FROM `{table}`")
            verify_counts[table] = cur.fetchone()[0]

    conn.close()
    print(
        json.dumps(
            {
                "import_status": "PASS",
                "database": args.database,
                "seed_dir": str(seed_dir),
                "imported_counts": imported_counts,
                "verify_counts": verify_counts,
            },
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()