#!/usr/bin/env python3
"""Decode binary_plan column in TiDB slow log CSV using tidb_decode_binary_plan().

Streams rows one by one to avoid loading the entire file into memory.

Usage:
    python decode_binary_plan.py --input /app/local/slowlog-xxx.csv [--host 127.0.0.1] [--port 14000]

Output:
    /app/local/slowlog-xxx.decoded.csv  (same dir, .decoded.csv suffix)
"""

import argparse
import csv
import re
import socket
import sys
import time

import pymysql

# plan texts and binary_plan values can exceed Python csv's default 128 KB limit
csv.field_size_limit(sys.maxsize)

# Matches: tidb_decode_binary_plan('...')
BINARY_PLAN_RE = re.compile(r"tidb_decode_binary_plan\('([^']+)'\)")


def wait_for_tidb(host: str, port: int, timeout: int = 60) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            s = socket.socket()
            s.settimeout(1)
            s.connect((host, port))
            s.close()
            return True
        except OSError:
            time.sleep(1)
    return False


def decode_plan(cursor, binary_plan_value: str) -> str:
    """Extract encoded arg and call tidb_decode_binary_plan via SQL."""
    m = BINARY_PLAN_RE.search(binary_plan_value)
    if not m:
        return ""
    encoded = m.group(1)
    cursor.execute("SELECT tidb_decode_binary_plan(%s)", (encoded,))
    row = cursor.fetchone()
    return row[0] if row and row[0] else ""


def derive_output_path(input_path: str) -> str:
    if input_path.endswith(".csv"):
        return input_path[:-4] + ".decoded.csv"
    return input_path + ".decoded.csv"


def main() -> None:
    parser = argparse.ArgumentParser(description="Decode binary_plan column in slow log CSV")
    parser.add_argument("--input", required=True, help="Input CSV file path")
    parser.add_argument("--host", default="127.0.0.1", help="TiDB host (default: 127.0.0.1)")
    parser.add_argument("--port", type=int, default=14000, help="TiDB port (default: 14000)")
    args = parser.parse_args()

    output_path = derive_output_path(args.input)
    print(f"Input:  {args.input}", flush=True)
    print(f"Output: {output_path}", flush=True)

    print(f"Waiting for TiDB at {args.host}:{args.port} ...", flush=True)
    if not wait_for_tidb(args.host, args.port):
        print("ERROR: TiDB not available after timeout", file=sys.stderr)
        sys.exit(1)
    print("TiDB is ready.", flush=True)

    conn = pymysql.connect(
        host=args.host,
        port=args.port,
        user="root",
        password="",
        database="",
        connect_timeout=10,
    )
    cursor = conn.cursor()

    processed = 0
    skipped = 0

    with open(args.input, newline="", encoding="utf-8") as fin, \
         open(output_path, "w", newline="", encoding="utf-8") as fout:

        reader = csv.DictReader(fin)
        if reader.fieldnames is None:
            print("ERROR: empty or invalid CSV", file=sys.stderr)
            sys.exit(1)

        fieldnames = list(reader.fieldnames)
        if "plan" not in fieldnames:
            fieldnames.append("plan")

        writer = csv.DictWriter(fout, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()

        for row in reader:
            binary_plan = row.get("binary_plan") or ""
            if binary_plan:
                try:
                    row["plan"] = decode_plan(cursor, binary_plan)
                except Exception as e:
                    print(f"  WARN row {processed + 1}: decode failed: {e}", flush=True)
                    row["plan"] = ""
                    skipped += 1
            else:
                row["plan"] = ""

            writer.writerow(row)
            processed += 1

            if processed % 500 == 0:
                print(f"  {processed} rows processed...", flush=True)

    cursor.close()
    conn.close()
    print(f"Done. {processed} rows written to {output_path} ({skipped} decode errors)", flush=True)


if __name__ == "__main__":
    main()