#!/usr/bin/env python3
import argparse
import json
import os
import ssl
import sys
from urllib.parse import unquote, urlparse


def main() -> int:
    parser = argparse.ArgumentParser(description="Run one read-only Auto-oncall database query")
    parser.add_argument("--sql", required=True, help="one SELECT, SHOW, DESCRIBE, DESC, EXPLAIN, or WITH statement")
    parser.add_argument("--max-rows", type=int, default=1000)
    args = parser.parse_args()

    if args.max_rows < 1 or args.max_rows > 5000:
        parser.error("--max-rows must be between 1 and 5000")
    statement = args.sql.strip()
    first_word = statement.split(None, 1)[0].upper() if statement else ""
    if first_word not in {"SELECT", "SHOW", "DESCRIBE", "DESC", "EXPLAIN", "WITH"}:
        parser.error("only read-only SQL statements are allowed")

    database_url = os.environ.get("AUTO_ONCALL_DATABASE_URL", "").strip()
    if not database_url:
        print("AUTO_ONCALL_DATABASE_URL is not configured", file=sys.stderr)
        return 2
    parsed = urlparse(database_url)
    if parsed.scheme not in {"mysql", "mysql2"} or not parsed.hostname or not parsed.username:
        print("AUTO_ONCALL_DATABASE_URL is invalid", file=sys.stderr)
        return 2

    import pymysql
    import pymysql.cursors

    connection = pymysql.connect(
        host=parsed.hostname,
        port=parsed.port or 4000,
        user=unquote(parsed.username),
        password=unquote(parsed.password or ""),
        database=unquote(parsed.path.lstrip("/")),
        ssl=ssl.create_default_context(),
        cursorclass=pymysql.cursors.DictCursor,
        read_timeout=30,
        write_timeout=30,
        connect_timeout=15,
        autocommit=True,
    )
    try:
        with connection.cursor() as cursor:
            cursor.execute(statement)
            rows = cursor.fetchmany(args.max_rows + 1)
        truncated = len(rows) > args.max_rows
        if truncated:
            rows = rows[: args.max_rows]
        print(json.dumps({"rows": rows, "count": len(rows), "truncated": truncated}, default=str, ensure_ascii=False))
    finally:
        connection.close()
    return 0


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