#!/usr/bin/env python3
"""Count restored MySQL 5.6 tables in an isolated Docker container.

Outputs only schema/table names and row counts; no row contents or credentials.
"""

import csv
import re
import subprocess
import sys
from pathlib import Path


CONTAINER = "dm8-mysql56-newsource-20260928"
SCHEMAS = ("zhct", "aizhct")
CHUNK_SIZE = 20


def mysql(sql: str) -> str:
    result = subprocess.run(
        ["docker", "exec", CONTAINER, "mysql", "--batch", "--raw", "--skip-column-names", "-e", sql],
        check=True,
        capture_output=True,
        text=True,
    )
    return result.stdout


def main() -> int:
    output = Path(__file__).resolve().parent / "source-table-counts.csv"
    rows: list[tuple[str, str, int]] = []
    for schema in SCHEMAS:
        names = mysql(
            "SELECT TABLE_NAME FROM information_schema.TABLES "
            f"WHERE TABLE_SCHEMA='{schema}' AND TABLE_TYPE='BASE TABLE' ORDER BY TABLE_NAME"
        ).splitlines()
        if not names or any(not re.fullmatch(r"[A-Za-z0-9_]+", name) for name in names):
            raise ValueError(f"Unexpected or missing table name in {schema}")
        for start in range(0, len(names), CHUNK_SIZE):
            chunk = names[start : start + CHUNK_SIZE]
            union = " UNION ALL ".join(
                f"SELECT '{name}' AS table_name, COUNT(*) AS row_count FROM `{schema}`.`{name}`"
                for name in chunk
            )
            for line in mysql(union).splitlines():
                table_name, row_count = line.split("\t", 1)
                rows.append((schema, table_name, int(row_count)))
        if sum(1 for row in rows if row[0] == schema) != len(names):
            raise ValueError(f"Count mismatch in {schema}")
        print(f"{schema}: tables={len(names)} rows={sum(row[2] for row in rows if row[0] == schema)}")

    with output.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(("schema", "table_name", "row_count"))
        writer.writerows(rows)
    print(f"wrote {output} ({len(rows)} tables)")
    return 0


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