#!/usr/bin/env python3
"""
Test runner for MySQL test cases.

This script parses test files containing MySQL queries and expected results,
executes the queries against a MySQL database, and compares the actual results
with the expected results.

Supported directives (multi-line supported):
- mysql> <query>              Execute a MySQL query and compare with expected result
- wait> <type> ...           Wait for a condition to be met
- expect_plan_index_covering> <query>  Expect EXPLAIN plan to use a covering index without TiKV lookup
- insert_hdfs_log> <table> [count]  Insert HDFS logs into an existing TiDB table via ci/read_hdfs.py

Wait directive types:
- shard_progress table=<db.table> index=<index_name> [timeout=300] [interval=5]

Expect plan index covering:
- expect_plan_index_covering> <query>  Run EXPLAIN and assert the plan uses an index scan without TiKV table lookup

Insert HDFS logs:
- insert_hdfs_log> hdfs_10w 100000
- insert_hdfs_log> table=hdfs_10w count=100000
"""

import argparse
import json
import os
import re
import subprocess
import sys
import tempfile
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import List, Dict, Any, Optional, Union, Tuple

from schema_mapper import SchemaMapper, make_schema_name, validate_schema_name
from utils import mysql_connection, execute_sql, wait_shard_progress


@dataclass
class MySQLDirective:
    """Represents a single MySQL directive with query and expected result."""
    query: str
    expected_result: List[List[str]]
    line_no: int

    def __str__(self):
        return f"Query: {self.query}"


@dataclass
class WaitDirective:
    """Represents a wait directive for synchronization."""
    wait_type: str
    params: Dict[str, str]
    line_no: int

    def __str__(self):
        params_str = " ".join(f"{k}={v}" for k, v in self.params.items())
        return f"Wait: {self.wait_type} {params_str}"


@dataclass
class ExpectPlanIndexCoveringDirective:
    """Expect EXPLAIN plan to use a covering index without TiKV table lookup."""
    query: str
    line_no: int

    def __str__(self):
        return f"ExpectPlanIndexCovering: {self.query}"


@dataclass
class InsertHdfsLogDirective:
    """Insert HDFS log dataset into an existing table."""
    table_name: str
    line_no: int
    row_count: int = 100000

    def __str__(self):
        return f"InsertHdfsLog: table={self.table_name} count={self.row_count}"


@dataclass
class DirectiveTiming:
    """Timing record for one executed directive."""
    file_path: str
    line_no: int
    directive_type: str
    label: str
    elapsed: float
    passed: bool


# Type alias for parsed directives
Directive = Union[MySQLDirective, WaitDirective, ExpectPlanIndexCoveringDirective, InsertHdfsLogDirective]


class TestParser:
    """Parser for MySQL test files."""

    # Regex to match directive prefixes.
    DIRECTIVE_PATTERN = re.compile(
        r'^(mysql|wait|expect_plan_index_covering|insert_hdfs_log)>\s*',
        re.MULTILINE,
    )

    def __init__(self, file_path: str):
        self.file_path = file_path
        self.directives: List[Directive] = []

    def parse(self) -> List[Directive]:
        """Parse the test file and extract directives."""
        try:
            with open(self.file_path, 'r', encoding='utf-8') as f:
                content = f.read()

            # Find all directive matches with their positions
            matches = list(self.DIRECTIVE_PATTERN.finditer(content))

            for i, match in enumerate(matches):
                directive_type = match.group(1)
                start_pos = match.end()
                line_no = content.count('\n', 0, match.start()) + 1

                # Find end position (start of next directive or end of file)
                if i + 1 < len(matches):
                    end_pos = matches[i + 1].start()
                else:
                    end_pos = len(content)

                section = content[start_pos:end_pos].strip()

                if directive_type == 'mysql':
                    directive = self._parse_mysql_section(section, line_no)
                    if directive:
                        self.directives.append(directive)
                elif directive_type == 'wait':
                    directive = self._parse_wait_section(section, line_no)
                    if directive:
                        self.directives.append(directive)
                elif directive_type == 'expect_plan_index_covering':
                    directive = self._parse_expect_plan_index_covering_section(
                        section, line_no)
                    if directive:
                        self.directives.append(directive)
                elif directive_type == 'insert_hdfs_log':
                    directive = self._parse_insert_hdfs_log_section(
                        section, line_no)
                    if directive:
                        self.directives.append(directive)

            return self.directives

        except Exception as e:
            print(f"Error parsing test file {self.file_path}: {e}")
            return []

    def _parse_mysql_section(self, section: str, line_no: int) -> Optional[MySQLDirective]:
        """Parse a mysql> section. Supports multi-line queries."""
        lines = section.split('\n')
        query_lines: List[str] = []
        result_start: Optional[int] = None

        for i, line in enumerate(lines):
            stripped = line.strip()
            # Ignore shell-style comments (e.g. disabled expected blocks after #mysql>)
            if stripped.startswith('#'):
                continue
            # Result table: separator (+/-) or row (|)
            if stripped.startswith('+') or stripped.startswith('-') or stripped.startswith('|'):
                result_start = i
                break
            if stripped:
                query_lines.append(stripped)

        if not query_lines:
            return None

        query = ' '.join(query_lines)
        if query.endswith(';'):
            query = query[:-1].strip()

        expected_result = (
            self._parse_result_table(
                lines[result_start:]) if result_start is not None else []
        )
        return MySQLDirective(query, expected_result, line_no)

    def _parse_wait_section(self, section: str, line_no: int) -> Optional[WaitDirective]:
        """Parse a wait> section. Supports multi-line by joining non-empty lines."""
        line = ' '.join(ln.strip() for ln in section.split('\n') if ln.strip())
        if not line:
            return None

        parts = line.split()
        if not parts:
            return None

        wait_type = parts[0]
        params: Dict[str, str] = {}

        for part in parts[1:]:
            if '=' in part:
                key, value = part.split('=', 1)
                params[key] = value

        return WaitDirective(wait_type, params, line_no)

    def _parse_expect_plan_index_covering_section(
        self, section: str, line_no: int
    ) -> Optional[ExpectPlanIndexCoveringDirective]:
        """Parse expect_plan_index_covering> section: single EXPLAIN query."""
        lines = [ln.strip() for ln in section.split('\n') if ln.strip()]
        if not lines:
            return None

        query = ' '.join(lines)
        if query.endswith(';'):
            query = query[:-1].strip()
        if not query.upper().startswith('EXPLAIN'):
            query = 'EXPLAIN ' + query

        return ExpectPlanIndexCoveringDirective(query=query, line_no=line_no)

    def _parse_insert_hdfs_log_section(self, section: str, line_no: int) -> Optional[InsertHdfsLogDirective]:
        """Parse insert_hdfs_log> section. Supports:
        - insert_hdfs_log> hdfs_10w 100000
        - insert_hdfs_log> table=hdfs_10w count=100000
        """
        line = ' '.join(ln.strip() for ln in section.split('\n') if ln.strip())
        if not line:
            return None

        parts = line.split()
        key_values: Dict[str, str] = {}
        positional: List[str] = []
        for part in parts:
            if '=' in part:
                key, value = part.split('=', 1)
                key_values[key.strip()] = value.strip()
            else:
                positional.append(part)

        table_name = key_values.get('table')
        count_text = key_values.get('count')

        if table_name is None and positional:
            table_name = positional[0]
        if count_text is None and len(positional) >= 2:
            count_text = positional[1]

        if not table_name:
            print("ERROR: insert_hdfs_log requires a table name")
            return None

        row_count = 100000
        if count_text is not None:
            try:
                row_count = int(count_text)
            except ValueError:
                print(f"ERROR: insert_hdfs_log invalid count: {count_text}")
                return None
        if row_count <= 0:
            print(f"ERROR: insert_hdfs_log count must be positive: {row_count}")
            return None

        return InsertHdfsLogDirective(
            table_name=table_name,
            line_no=line_no,
            row_count=row_count,
        )

    def _parse_result_table(self, lines: List[str]) -> List[List[str]]:
        """Parse the expected result table from the test output."""
        result = []
        found_header = False

        for line in lines:
            line = line.strip()
            # Do not parse commented-out pseudo result tables (#| ... |)
            if line.startswith('#'):
                continue

            # Skip empty lines and separators
            if not line or line.startswith('+') or line.startswith('-'):
                continue

            # Check if this is a table row
            if '|' in line:
                # Remove leading/trailing |
                line = line.strip('|')
                # Split by | and clean up each cell
                row = [cell.strip() for cell in line.split('|')]
                if row and any(cell for cell in row):  # Skip empty rows
                    if not found_header:
                        # First row is usually the header, skip it for comparison
                        # but mark that we found it
                        found_header = True
                        continue
                    else:
                        # This is a data row
                        result.append(row)

        return result


class TestRunner:
    """Runs test cases against a MySQL database."""

    def __init__(
        self,
        host: str,
        port: int,
        user: str,
        password: str,
        database: Optional[str] = None,
        s3_endpoint: str = "localhost:9000",
        s3_bucket: str = "ticidefaultbucket",
        s3_prefix: str = "tici_default_prefix",
        s3_access_key: str = "minioadmin",
        s3_secret_key: str = "minioadmin",
        logical_schema: str = "test",
        physical_schema: Optional[str] = None,
    ):
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.s3_endpoint = s3_endpoint
        self.s3_bucket = s3_bucket
        self.s3_prefix = s3_prefix
        self.s3_access_key = s3_access_key
        self.s3_secret_key = s3_secret_key
        self.project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
        self.asset_dir = os.getenv("ASSET_DIR", self.project_root)
        self.directive_timings: List[DirectiveTiming] = []
        self.logical_schema = logical_schema
        self.physical_schema = physical_schema
        self.schema_mapper = (
            SchemaMapper(logical_schema, physical_schema)
            if physical_schema and physical_schema != logical_schema else None
        )

    def execute_query(self, query: str) -> List[List[str]]:
        """Execute a query and return the results."""
        if self.schema_mapper:
            query = self.schema_mapper.remap_sql(query)
        with mysql_connection(self.host, self.port, self.user, self.database) as connection:
            results = execute_sql(connection, query)

            if isinstance(results, int):
                return []
            # Convert results to strings for comparison
            if results:
                return [[str(cell) for cell in row] for row in results]
            return []

    def execute_query_with_columns(self, query: str) -> Tuple[List[str], List[List[str]]]:
        """Execute a query and return (column_names, rows). Used for EXPLAIN result validation."""
        if self.schema_mapper:
            query = self.schema_mapper.remap_sql(query)
        with mysql_connection(self.host, self.port, self.user, self.database) as connection:
            cursor = connection.cursor()
            try:
                cursor.execute(query)
                columns = [d[0]
                           for d in cursor.description] if cursor.description else []
                rows = cursor.fetchall()
                str_rows = [[str(cell) for cell in row] for row in rows]
                return (columns, str_rows)
            finally:
                cursor.close()

    def run_mysql_directive(self, directive: MySQLDirective) -> bool:
        """Run a single query directive and return whether it passed."""
        try:
            print(f"RUN: {directive.query}")
            actual_result = self.execute_query(directive.query)

            # Compare results
            if self._compare_results(directive.expected_result, actual_result):
                return True
            else:
                print(f"FAIL: {directive.query}")
                print(f"  Expected: {directive.expected_result}")
                print(f"  Actual: {actual_result}")
                return False

        except Exception as e:
            print(f"ERROR: {directive.query}")
            print(f"  Error: {e}")
            return False

    def run_expect_plan_index_covering(self, case: ExpectPlanIndexCoveringDirective) -> bool:
        """Expect EXPLAIN plan to use an index scan without TiKV table lookup."""
        forbidden_task = 'tikv'
        forbidden_operators = ('IndexLookUp', 'TableRowIDScan')
        try:
            print(f"RUN: {case.query}")
            columns, rows = self.execute_query_with_columns(case.query)

            id_idx = None
            task_idx = None
            access_idx = None
            for i, col in enumerate(columns):
                normalized_col = col.strip().lower() if col else ''
                if normalized_col == 'id':
                    id_idx = i
                elif normalized_col == 'task':
                    task_idx = i
                elif normalized_col == 'access object':
                    access_idx = i

            if task_idx is None:
                print(
                    f"FAIL: No 'task' column in EXPLAIN result (columns: {columns})")
                return False
            if id_idx is None:
                print(
                    f"FAIL: No 'id' column in EXPLAIN result (columns: {columns})")
                return False

            has_index_scan = False
            has_index_access = False
            for row_idx, row in enumerate(rows):
                if task_idx >= len(row):
                    continue
                task_val = row[task_idx] or ''
                if forbidden_task.lower() in task_val.lower():
                    print(
                        f"FAIL: plan should not read TiKV rows, task contains '{forbidden_task}' in row {row_idx + 1}: {task_val}")
                    return False

                operator_id = row[id_idx] if id_idx < len(row) else ''
                operator_id = operator_id or ''
                for forbidden_operator in forbidden_operators:
                    if forbidden_operator.lower() in operator_id.lower():
                        print(
                            f"FAIL: plan should be index covering, operator contains '{forbidden_operator}' in row {row_idx + 1}: {operator_id}")
                        return False

                if 'index' in operator_id.lower() and 'scan' in operator_id.lower():
                    has_index_scan = True

                if access_idx is not None and access_idx < len(row):
                    access_object = row[access_idx] or ''
                    if 'index:' in access_object.lower():
                        has_index_access = True

            if not has_index_scan or not has_index_access:
                print(
                    "FAIL: plan should use a covering index, but no index scan/access was found")
                return False

            print("PASS: plan uses a covering index without TiKV table lookup")
            return True

        except Exception as e:
            print(f"ERROR: {case.query}")
            print(f"  Error: {e}")
            return False

    def run_insert_hdfs_log(self, directive: InsertHdfsLogDirective) -> bool:
        """Insert HDFS dataset into existing target table by invoking ci/read_hdfs.py."""
        if not re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', directive.table_name):
            print(f"ERROR: Invalid table name for insert_hdfs_log: {directive.table_name}")
            return False

        dataset_file = os.path.join(self.asset_dir, "hdfs-logs-multitenants.json")
        if not os.path.exists(dataset_file):
            print(f"ERROR: Dataset file not found: {dataset_file}")
            return False

        read_hdfs_script = os.path.join(self.project_root, "ci", "read_hdfs.py")
        if not os.path.exists(read_hdfs_script):
            print(f"ERROR: read_hdfs.py not found: {read_hdfs_script}")
            return False

        cmd = [
            sys.executable,
            read_hdfs_script,
            directive.table_name,
            "--input",
            dataset_file,
            "--max_rows",
            str(directive.row_count),
            "--insert_only",
            "--tidb_host",
            self.host,
            "--tidb_port",
            str(self.port),
            "--database",
            self.physical_schema or self.database or self.logical_schema,
        ]

        print(
            f"RUN: insert_hdfs_log table={directive.table_name} count={directive.row_count} "
            f"input={dataset_file}"
        )
        try:
            subprocess.run(cmd, check=True)
            print(f"PASS: insert_hdfs_log table={directive.table_name} count={directive.row_count}")
            return True
        except subprocess.CalledProcessError as e:
            print(f"FAIL: insert_hdfs_log table={directive.table_name} count={directive.row_count}")
            print(f"  Error: read_hdfs.py exited with code {e.returncode}")
            return False

    def run_wait_directive(self, directive: WaitDirective) -> bool:
        """Execute a wait directive and return whether it succeeded."""
        try:
            if directive.wait_type == 'shard_progress':
                params = directive.params
                if self.schema_mapper:
                    params = self.schema_mapper.remap_wait_params(params)
                return self._wait_shard_progress(params)
            else:
                print(f"ERROR: Unknown wait type: {directive.wait_type}")
                return False
        except Exception as e:
            print(f"ERROR: {directive}")
            print(f"  Error: {e}")
            return False

    def _wait_shard_progress(self, params: Dict[str, str]) -> bool:
        """Execute shard_progress wait directive."""
        table = params.get('table', '')
        if '.' not in table:
            print(
                f"ERROR: Invalid table format '{table}', expected 'db.table'")
            return False

        db_name, table_name = table.split('.', 1)
        index_name = params.get('index', '')
        # Adapted for the last character of index_name end by ';'
        index_name = index_name.rstrip(';')
        if not index_name:
            print("ERROR: 'index' parameter is required for shard_progress")
            return False

        timeout = int(params.get('timeout', '300'))
        interval = int(params.get('interval', '5'))

        wait_shard_progress(
            host=self.host,
            port=self.port,
            db_name=db_name,
            table_name=table_name,
            index_name=index_name,
            bucket=self.s3_bucket,
            prefix=self.s3_prefix,
            s3_endpoint=self.s3_endpoint,
            s3_access_key=self.s3_access_key,
            s3_secret_key=self.s3_secret_key,
            timeout=timeout,
            user=self.user,
        )
        return True

    def _normalize_cell_for_compare(self, cell: str) -> str:
        """Normalize cell value for comparison. NULL (MySQL) and None (Python driver) are equivalent."""
        if cell is None:
            return "NULL"
        s = str(cell).strip()
        if s.upper() == "NULL" or s == "None":
            return "NULL"
        return s

    def _compare_results(self, expected: List[List[str]], actual: List[List[str]]) -> bool:
        """Compare expected and actual results."""
        if len(expected) != len(actual):
            return False

        for _, (exp_row, act_row) in enumerate(zip(expected, actual)):
            if len(exp_row) != len(act_row):
                return False

            for _, (exp_cell, act_cell) in enumerate(zip(exp_row, act_row)):
                exp_norm = self._normalize_cell_for_compare(exp_cell)
                act_norm = self._normalize_cell_for_compare(act_cell)
                # Handle numeric comparisons (only when both are numeric after normalization)
                if exp_norm != "NULL" and act_norm != "NULL":
                    if exp_norm.isdigit() and act_norm.isdigit():
                        if int(exp_norm) != int(act_norm):
                            return False
                    elif exp_norm != act_norm:
                        return False
                elif exp_norm != act_norm:
                    return False

        return True

    def run_directive(self, directive: Directive) -> bool:
        """Run a single directive (test case, wait, explain_check, or insert_hdfs_log)."""
        if isinstance(directive, MySQLDirective):
            return self.run_mysql_directive(directive)
        elif isinstance(directive, WaitDirective):
            return self.run_wait_directive(directive)
        elif isinstance(directive, ExpectPlanIndexCoveringDirective):
            return self.run_expect_plan_index_covering(directive)
        elif isinstance(directive, InsertHdfsLogDirective):
            return self.run_insert_hdfs_log(directive)
        else:
            print(f"ERROR: Unknown directive type: {type(directive)}")
            return False

    def _directive_type(self, directive: Directive) -> str:
        """Return the directive type name used in timing summaries."""
        if isinstance(directive, MySQLDirective):
            return "mysql"
        if isinstance(directive, WaitDirective):
            return "wait"
        if isinstance(directive, ExpectPlanIndexCoveringDirective):
            return "expect_plan_index_covering"
        if isinstance(directive, InsertHdfsLogDirective):
            return "insert_hdfs_log"
        return type(directive).__name__

    def _directive_label(self, directive: Directive) -> str:
        """Return a compact directive label for timing summaries."""
        if isinstance(directive, MySQLDirective):
            return directive.query[:160]
        if isinstance(directive, WaitDirective):
            params = " ".join(f"{k}={v}" for k, v in directive.params.items())
            return f"{directive.wait_type} {params}".strip()
        if isinstance(directive, ExpectPlanIndexCoveringDirective):
            return directive.query[:160]
        if isinstance(directive, InsertHdfsLogDirective):
            return f"table={directive.table_name} count={directive.row_count}"
        return str(directive)[:160]

    def _record_directive_timing(
        self,
        file_path: Optional[str],
        directive: Directive,
        passed: bool,
        elapsed: float,
    ) -> None:
        """Record timing for one executed directive."""
        self.directive_timings.append(DirectiveTiming(
            file_path=file_path or "<unknown>",
            line_no=directive.line_no,
            directive_type=self._directive_type(directive),
            label=self._directive_label(directive),
            elapsed=elapsed,
            passed=passed,
        ))

    def run_directives(self, directives: List[Directive], file_path: Optional[str] = None) -> Dict[str, Any]:
        """Run all directives and return results summary."""
        total = len(directives)
        executed = 0
        passed = 0
        failed = 0
        errors = 0

        print(f"\nRunning {total} directive(s)...")

        for i, directive in enumerate(directives, 1):
            print(f"[{i}/{total}] ", end="")
            executed += 1
            directive_start = time.monotonic()

            try:
                directive_passed = self.run_directive(directive)
                self._record_directive_timing(
                    file_path=file_path,
                    directive=directive,
                    passed=directive_passed,
                    elapsed=time.monotonic() - directive_start,
                )
                if directive_passed:
                    passed += 1
                else:
                    failed += 1
                    print(
                        "STOP: failure encountered, skipping remaining directives "
                        f"in {file_path} at line {directive.line_no}."
                    )
                    break
            except Exception as e:
                errors += 1
                self._record_directive_timing(
                    file_path=file_path,
                    directive=directive,
                    passed=False,
                    elapsed=time.monotonic() - directive_start,
                )
                print(f"STOP: error encountered: {e}, skipping remaining directives "
                      f"in {file_path} at line {directive.line_no}.")
                break

        results = {
            'total': executed,
            'passed': passed,
            'failed': failed,
            'errors': errors,
            'success_rate': (passed / executed * 100) if executed > 0 else 0
        }

        return results


def find_test_files(test_dir: str) -> List[str]:
    """Find all test files in the given directory."""
    test_files = []

    if os.path.isfile(test_dir):
        # If it's a single file, return it
        if test_dir.endswith('.test'):
            test_files.append(test_dir)
    elif os.path.isdir(test_dir):
        # If it's a directory, find all .test files
        for root, dirs, files in os.walk(test_dir):
            for file in files:
                if file.endswith('.test'):
                    test_files.append(os.path.join(root, file))

    return sorted(test_files)


def _quote_schema(schema_name: str) -> str:
    """Return a validated schema name quoted as a MySQL identifier."""
    return "`{}`".format(validate_schema_name(schema_name))


def create_database(host: str, port: int, user: str, database: str) -> None:
    """Create a database for an isolated e2e test file."""
    with mysql_connection(host, port, user) as connection:
        execute_sql(connection, "CREATE DATABASE IF NOT EXISTS {}".format(_quote_schema(database)))


def drop_database(host: str, port: int, user: str, database: str) -> None:
    """Drop a database created for an isolated e2e test file."""
    with mysql_connection(host, port, user) as connection:
        execute_sql(connection, "DROP DATABASE IF EXISTS {}".format(_quote_schema(database)))


def directive_timing_to_dict(timing: DirectiveTiming) -> Dict[str, Any]:
    """Convert directive timing to a JSON-serializable dictionary."""
    return {
        'file_path': timing.file_path,
        'line_no': timing.line_no,
        'directive_type': timing.directive_type,
        'label': timing.label,
        'elapsed': timing.elapsed,
        'passed': timing.passed,
    }


def directive_timing_from_dict(data: Dict[str, Any]) -> DirectiveTiming:
    """Convert a JSON dictionary back to DirectiveTiming."""
    return DirectiveTiming(
        file_path=data['file_path'],
        line_no=data['line_no'],
        directive_type=data['directive_type'],
        label=data['label'],
        elapsed=data['elapsed'],
        passed=data['passed'],
    )


def write_summary_json(
    path: Optional[str],
    results: Dict[str, Any],
    directive_timings: List[DirectiveTiming],
    runner_elapsed: float,
) -> None:
    """Write a machine-readable runner summary for parent processes."""
    if not path:
        return

    payload = {
        'results': results,
        'directive_timings': [
            directive_timing_to_dict(timing) for timing in directive_timings
        ],
        'runner_elapsed': runner_elapsed,
    }
    with open(path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, sort_keys=True)


def run_test_files(test_files: List[str], runner: TestRunner) -> Dict[str, Any]:
    """Run all test files and return combined results."""
    total_files = len(test_files)
    total_directives = 0
    passed = 0
    failed = 0
    errors = 0
    file_results = {}

    print(f"\nFound {total_files} test file(s) to run...")

    for i, test_file in enumerate(test_files, 1):
        print(f"\n[{i}/{total_files}] Running test file: {test_file}")
        print("-" * 60)
        file_start = time.monotonic()

        # Parse test file
        parser = TestParser(test_file)
        directives = parser.parse()

        if not directives:
            print(f"No directives found in {test_file}")
            file_results[test_file] = {
                'total': 0, 'passed': 0, 'failed': 0, 'errors': 0, 'success_rate': 0,
                'elapsed': time.monotonic() - file_start
            }
            continue

        print(f"Parsed {len(directives)} directive(s) from {test_file}")

        # Run directives for this file
        results = runner.run_directives(directives, file_path=test_file)
        results['elapsed'] = time.monotonic() - file_start
        file_results[test_file] = results

        # Update totals
        total_directives += results['total']
        passed += results['passed']
        failed += results['failed']
        errors += results['errors']

        # Print file summary
        print(f"\nFile summary: {results['passed']}/{results['total']} passed "
              f"({results['success_rate']:.1f}%, elapsed={results['elapsed']:.1f}s)")

    # Calculate overall results
    overall_results = {
        'total_files': total_files,
        'total_directives': total_directives,
        'passed': passed,
        'failed': failed,
        'errors': errors,
        'success_rate': (passed / total_directives * 100) if total_directives > 0 else 0,
        'file_results': file_results
    }

    return overall_results


def run_one_test_file_subprocess(
    test_file: str,
    worker_index: int,
    args: argparse.Namespace,
    summary_dir: str,
) -> Dict[str, Any]:
    """Run one test file in a child runner process and capture its summary."""
    physical_schema = make_schema_name(args.parallel_schema_prefix, worker_index, test_file)
    summary_path = os.path.join(summary_dir, "summary_{}.json".format(worker_index))
    cmd = [
        sys.executable,
        os.path.abspath(__file__),
        test_file,
        "--host", args.host,
        "--port", str(args.port),
        "--user", args.user,
        "--password", args.password,
        "--database", physical_schema,
        "--s3-endpoint", args.s3_endpoint,
        "--s3-bucket", args.s3_bucket,
        "--s3-prefix", args.s3_prefix,
        "--s3-access-key", args.s3_access_key,
        "--s3-secret-key", args.s3_secret_key,
        "--logical-schema", args.logical_schema,
        "--physical-schema", physical_schema,
        "--parallel", "1",
        "--summary-json", summary_path,
    ]
    if args.keep_schemas_on_success:
        cmd.append("--keep-schemas-on-success")

    completed = subprocess.run(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        universal_newlines=True,
    )

    summary = None
    if os.path.exists(summary_path):
        with open(summary_path, "r", encoding="utf-8") as f:
            summary = json.load(f)

    return {
        'file': test_file,
        'schema': physical_schema,
        'returncode': completed.returncode,
        'output': completed.stdout,
        'summary': summary,
    }


def run_test_files_parallel(test_files: List[str], args: argparse.Namespace) -> Dict[str, Any]:
    """Run test files concurrently in isolated child runner processes."""
    total_files = len(test_files)
    total_directives = 0
    passed = 0
    failed = 0
    errors = 0
    file_results = {}
    directive_timings: List[DirectiveTiming] = []

    print(f"\nFound {total_files} test file(s) to run with parallel={args.parallel}...")

    with tempfile.TemporaryDirectory(prefix="tici-e2e-summary-") as summary_dir:
        with ThreadPoolExecutor(max_workers=args.parallel) as executor:
            futures = {
                executor.submit(
                    run_one_test_file_subprocess,
                    test_file,
                    i + 1,
                    args,
                    summary_dir,
                ): test_file
                for i, test_file in enumerate(test_files)
            }

            for future in as_completed(futures):
                test_file = futures[future]
                try:
                    item = future.result()
                except Exception as e:
                    result = {
                        'total': 0,
                        'passed': 0,
                        'failed': 0,
                        'errors': 1,
                        'success_rate': 0,
                        'elapsed': 0.0,
                        'schema': '<unknown>',
                    }
                    print(f"\n[failed] Running test file: {test_file}")
                    print("-" * 60)
                    print(f"ERROR: parallel worker failed before producing output: {e}")
                else:
                    print(
                        f"\n[completed] Running test file: {test_file} "
                        f"schema={item['schema']}"
                    )
                    print("-" * 60)
                    print(item['output'], end="")

                    if item['summary'] is None:
                        result = {
                            'total': 0,
                            'passed': 0,
                            'failed': 0,
                            'errors': 1,
                            'success_rate': 0,
                            'elapsed': 0.0,
                            'schema': item['schema'],
                        }
                        print(
                            "ERROR: child runner did not produce summary JSON "
                            f"(exit={item['returncode']})"
                        )
                    else:
                        child_results = item['summary']['results']
                        result = child_results['file_results'].get(test_file)
                        if result is None and child_results['file_results']:
                            result = next(iter(child_results['file_results'].values()))
                        if result is None:
                            result = {
                                'total': 0,
                                'passed': 0,
                                'failed': 0,
                                'errors': 1,
                                'success_rate': 0,
                                'elapsed': 0.0,
                                'schema': item['schema'],
                            }
                        result = dict(result)
                        result['schema'] = item['schema']
                        directive_timings.extend(
                            directive_timing_from_dict(timing)
                            for timing in item['summary'].get('directive_timings', [])
                        )
                        if item['returncode'] != 0 and result['failed'] == 0 and result['errors'] == 0:
                            result['errors'] = 1
                            result['success_rate'] = 0

                file_results[test_file] = result
                total_directives += result['total']
                passed += result['passed']
                failed += result['failed']
                errors += result['errors']

                print(
                    f"\nFile summary: {result['passed']}/{result['total']} passed "
                    f"({result['success_rate']:.1f}%, elapsed={result.get('elapsed', 0.0):.1f}s, "
                    f"schema={result.get('schema', '<none>')})"
                )

    return {
        'total_files': total_files,
        'total_directives': total_directives,
        'passed': passed,
        'failed': failed,
        'errors': errors,
        'success_rate': (passed / total_directives * 100) if total_directives > 0 else 0,
        'file_results': file_results,
        'directive_timings': directive_timings,
    }


def print_timing_summary(
    results: Dict[str, Any],
    directive_timings: List[DirectiveTiming],
    runner_elapsed: float,
    top_n: int = 10,
) -> None:
    """Print timing summary for e2e diagnostics."""
    print("\n" + "="*60)
    print("E2E TEST RUNNER TIMING SUMMARY")
    print("="*60)
    print(f"Total runner elapsed: {runner_elapsed:.1f}s")

    by_type: Dict[str, Dict[str, float]] = {}
    for timing in directive_timings:
        bucket = by_type.setdefault(timing.directive_type, {
            "count": 0,
            "elapsed": 0.0,
        })
        bucket["count"] += 1
        bucket["elapsed"] += timing.elapsed

    print("By directive type:")
    for directive_type in sorted(by_type):
        bucket = by_type[directive_type]
        print(
            f"  {directive_type}: count={int(bucket['count'])} "
            f"elapsed={bucket['elapsed']:.1f}s"
        )

    print("Slowest test files:")
    file_items = sorted(
        results['file_results'].items(),
        key=lambda item: item[1].get('elapsed', 0.0),
        reverse=True,
    )
    for file_path, file_result in file_items[:top_n]:
        print(f"  {file_result.get('elapsed', 0.0):.1f}s {file_path}")

    print("Slowest directives:")
    for timing in sorted(directive_timings, key=lambda t: t.elapsed, reverse=True)[:top_n]:
        status = "PASS" if timing.passed else "FAIL"
        print(
            f"  {timing.elapsed:.1f}s {timing.directive_type} {status} "
            f"{timing.file_path}:{timing.line_no} {timing.label}"
        )


def main():
    parser = argparse.ArgumentParser(
        description="Run MySQL test cases from test files or directories",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s tests/test-cases/multi-config/sample.test --host 127.0.0.1 --port 4000 --user root
  %(prog)s tests/test-cases/ --host localhost --database mydb
  %(prog)s tests/test-cases --host 127.0.0.1 --port 4000

Supported directives in test files:
  mysql> SELECT ...                 Execute SQL and compare result
  wait> shard_progress ...          Wait for shard progress to sync
  expect_plan_index_covering> EXPLAIN ...
                                    Expect plan to use a covering index without TiKV lookup
  insert_hdfs_log> <table> [count]  Insert HDFS logs into existing table (default count=100000)

Expect plan index covering:
  expect_plan_index_covering> explain select d from test.return_null use index (idx1) where c = '2025-01-01 12:00:00.123';

Wait directive examples:
  wait> shard_progress table=test.hdfs_10w index=ft_body_index timeout=300 interval=5

Insert HDFS logs examples:
  insert_hdfs_log> hdfs_10w 100000
  insert_hdfs_log> table=hdfs_10w count=50000
        """
    )

    # Test directory/file argument
    parser.add_argument('test_dir', help='Path to the test directory or file')

    # Database connection arguments
    parser.add_argument('--host', default='127.0.0.1', help='MySQL host (default: 127.0.0.1)')
    parser.add_argument('--port', type=int, default=4000, help='MySQL port (default: 4000)')
    parser.add_argument('--user', default='root', help='MySQL username (default: root)')
    parser.add_argument('--password', default='', help='MySQL password (default: empty)')
    parser.add_argument('--database', help='Database name (optional)')

    # S3 arguments for wait directives
    parser.add_argument('--s3-endpoint', default='localhost:9000', help='S3 endpoint (default: localhost:9000)')
    parser.add_argument('--s3-bucket', default='ticidefaultbucket', help='S3 bucket name (default: ticidefaultbucket)')
    parser.add_argument('--s3-prefix', default='tici_default_prefix', help='S3 prefix (default: tici_default_prefix)')
    parser.add_argument('--s3-access-key', default='minioadmin', help='S3 access key (default: minioadmin)')
    parser.add_argument('--s3-secret-key', default='minioadmin', help='S3 secret key (default: minioadmin)')
    parser.add_argument('--parallel', type=int, default=4, help='Number of test files to run concurrently (default: 4)')
    parser.add_argument('--logical-schema', default='test', help='Logical schema name used inside .test files (default: test)')
    parser.add_argument('--physical-schema', help='Physical schema name for a single isolated test file')
    parser.add_argument('--parallel-schema-prefix', default='test_e2e', help='Schema prefix for parallel child runners')
    parser.add_argument('--summary-json', help='Write a JSON summary to this path')
    parser.add_argument(
        '--keep-schemas-on-success',
        action='store_true',
        help='Keep isolated schemas after successful execution',
    )

    args = parser.parse_args()
    runner_start = time.monotonic()

    try:
        if args.parallel < 1:
            raise ValueError("--parallel must be >= 1")

        validate_schema_name(args.logical_schema)
        validate_schema_name(args.parallel_schema_prefix.rstrip("_"))
        if args.physical_schema:
            validate_schema_name(args.physical_schema)

        # Find test files
        test_files = find_test_files(args.test_dir)

        if not test_files:
            print(f"No test files found in {args.test_dir}")
            return 1

        use_parallel = args.parallel > 1 and len(test_files) > 1
        if use_parallel and args.physical_schema:
            raise ValueError("--parallel cannot be combined with --physical-schema when running multiple files")

        # Run tests
        if use_parallel:
            results = run_test_files_parallel(test_files, args)
            directive_timings = results.pop('directive_timings')
        else:
            schema_created = False
            if args.physical_schema:
                print(f"Creating isolated schema: {args.physical_schema}")
                create_database(args.host, args.port, args.user, args.physical_schema)
                schema_created = True

            runner = TestRunner(
                host=args.host,
                port=args.port,
                user=args.user,
                password=args.password,
                database=args.database or args.physical_schema,
                s3_endpoint=args.s3_endpoint,
                s3_bucket=args.s3_bucket,
                s3_prefix=args.s3_prefix,
                s3_access_key=args.s3_access_key,
                s3_secret_key=args.s3_secret_key,
                logical_schema=args.logical_schema,
                physical_schema=args.physical_schema,
            )
            results = run_test_files(test_files, runner)
            directive_timings = runner.directive_timings
            if args.physical_schema:
                for file_result in results['file_results'].values():
                    file_result['schema'] = args.physical_schema

        # Print overall summary
        print("\n" + "="*60)
        print("OVERALL TEST SUMMARY")
        print("="*60)
        print(f"Total files: {results['total_files']}")
        print(f"Total directives: {results['total_directives']}")
        print(f"Passed: {results['passed']}")
        print(f"Failed: {results['failed']}")
        print(f"Errors: {results['errors']}")
        print(f"Success rate: {results['success_rate']:.1f}%")

        # Print per-file results if multiple files
        if results['total_files'] > 1:
            print(f"\nPer-file results:")
            for file_path, file_result in results['file_results'].items():
                filename = os.path.basename(file_path)
                schema_text = ""
                if file_result.get('schema'):
                    schema_text = f", schema={file_result['schema']}"
                print(
                    f"  {filename}: {file_result['passed']}/{file_result['total']} "
                    f"({file_result['success_rate']:.1f}%{schema_text})"
                )

        runner_elapsed = time.monotonic() - runner_start
        print_timing_summary(results, directive_timings, runner_elapsed)

        success = results['failed'] == 0 and results['errors'] == 0
        write_summary_json(args.summary_json, results, directive_timings, runner_elapsed)

        if args.physical_schema and schema_created:
            if success and not args.keep_schemas_on_success:
                print(f"Dropping isolated schema: {args.physical_schema}")
                drop_database(args.host, args.port, args.user, args.physical_schema)
            else:
                print(f"Keeping isolated schema: {args.physical_schema}")

        # Return appropriate exit code
        return 0 if success else 1

    except Exception as e:
        print(f"Error: {e}")
        return 1


if __name__ == "__main__":
    sys.exit(main())