#!/usr/bin/env python3
from __future__ import annotations

import argparse
import os
import ssl
import sys

import pymysql
from pymysql.constants import CLIENT


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Execute SQL against TiDB over a TLS-capable PyMySQL connection.")
    parser.add_argument("--host", required=True)
    parser.add_argument("--port", required=True, type=int)
    parser.add_argument("--user", required=True)
    parser.add_argument("--password", default="")
    parser.add_argument("--database", required=True)
    parser.add_argument(
        "--tls",
        default=os.environ.get("TIDB_TLS", "true"),
        help="Enable TLS when connecting to TiDB (default: true).",
    )
    return parser.parse_args()


def tls_enabled(raw: str) -> bool:
    return str(raw).strip().lower() in {"1", "true", "yes", "on"}


def tidb_tls_context() -> ssl.SSLContext:
    context = ssl.create_default_context()
    context.check_hostname = False
    context.verify_mode = ssl.CERT_NONE
    return context


def main() -> int:
    args = parse_args()
    sql_text = sys.stdin.read()
    if not sql_text.strip():
        sys.stderr.write("no SQL provided on stdin\n")
        return 2

    connect_kwargs = {
        "host": args.host,
        "port": args.port,
        "user": args.user,
        "password": args.password,
        "database": args.database,
        "autocommit": True,
        "client_flag": CLIENT.MULTI_STATEMENTS,
        "connect_timeout": 10,
        "read_timeout": 30,
        "write_timeout": 30,
    }
    if tls_enabled(args.tls):
        connect_kwargs["ssl"] = tidb_tls_context()

    conn = pymysql.connect(**connect_kwargs)
    try:
        with conn.cursor() as cur:
            cur.execute(sql_text)
            while cur.nextset():
                pass
    finally:
        conn.close()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())