#!/usr/bin/env python3
"""
Common utility functions.
"""

import json
import mysql.connector
from mysql.connector import Error
import boto3
from botocore.client import Config
from contextlib import contextmanager
from typing import Dict, List, Any, Optional, Tuple


@contextmanager
def mysql_connection(host: str, port: int, user: str = "root", database: Optional[str] = None, timeout: int = 60):
    """
    Context manager for MySQL database connections.

    Args:
        host: Database host
        port: Database port
        user: Database user
        database: Database name (optional)

    Yields:
        mysql.connector.MySQLConnection: Database connection

    Raises:
        RuntimeError: If connection fails
    """
    connection = None
    try:
        connection = mysql.connector.connect(
            host=host,
            port=port,
            user=user,
            database=database,
            connection_timeout=timeout,
        )
        yield connection
    except Error as e:
        raise RuntimeError(f"Database connection failed: {e}")
    finally:
        if connection and connection.is_connected():
            connection.close()


def execute_sql(connection, sql: str, params: Optional[tuple] = None) -> Optional[Any]:
    """
    Execute SQL query with proper cursor management.

    Args:
        connection: Database connection
        sql: SQL query string
        params: Query parameters (optional)

    Returns:
        Query result or None

    Raises:
        RuntimeError: If query execution fails
    """
    cursor = None
    try:
        cursor = connection.cursor()
        if params:
            cursor.execute(sql, params)
        else:
            cursor.execute(sql)

        if sql.strip().upper().startswith("SELECT"):
            return cursor.fetchall()
        else:
            connection.commit()
            return cursor.rowcount
    except Error as e:
        raise RuntimeError(f"SQL execution failed: {e}")
    finally:
        if cursor:
            cursor.close()


def create_s3_client(endpoint: str, access_key: Optional[str] = None, secret_key: Optional[str] = None):
    """
    Create S3/MinIO client with standard configuration.

    Args:
        endpoint: S3 endpoint URL (without http://)
        access_key: S3 access key
        secret_key: S3 secret key

    Returns:
        boto3.client: Configured S3 client

    Raises:
        RuntimeError: If client creation fails
    """
    try:
        session = boto3.session.Session()
        return session.client(
            "s3",
            endpoint_url=f"http://{endpoint}",
            aws_access_key_id=access_key,
            aws_secret_access_key=secret_key,
            config=Config(signature_version="s3v4"),
        )
    except Exception as e:
        raise RuntimeError(f"Failed to create S3 client: {e}")


def list_s3_files_with_prefix(s3_client, bucket: str, prefix: str) -> List[Dict[str, Any]]:
    """
    List all files in S3 bucket with given prefix.

    Args:
        s3_client: S3 client instance
        bucket: S3 bucket name
        prefix: File prefix to filter by

    Returns:
        List of file dictionaries with 'key' and 'last_modified' fields

    Raises:
        RuntimeError: If S3 operation fails
    """
    try:
        files = []
        paginator = s3_client.get_paginator("list_objects_v2")
        pages = paginator.paginate(Bucket=bucket, Prefix=prefix)

        for page in pages:
            if "Contents" in page:
                for obj in page["Contents"]:
                    files.append({"key": obj["Key"], "last_modified": obj["LastModified"]})

        return files
    except Exception as e:
        raise RuntimeError(f"Failed to list S3 files: {e}")


def get_cdc_last_file(s3_client, bucket: str, prefix: str) -> Optional[str]:
    import time

    # Normalize prefix so keys under prefix match (prefix with or without trailing slash)
    search_prefix = prefix.rstrip("/") + "/"

    while True:
        try:
            all_files = list_s3_files_with_prefix(s3_client, bucket, search_prefix)
            cdc_files = [f for f in all_files if f["key"].endswith(".json") and "/CDC" in f["key"]]

            if not cdc_files:
                print("No CDC files found, retrying...")
                time.sleep(2)
                continue

            cdc_files.sort(key=lambda x: x["last_modified"], reverse=True)
            return cdc_files[0]["key"]

        except Exception as e:
            raise RuntimeError(f"Failed to get latest CDC file: {e}")


def safe_json_parse(json_str: str) -> Dict[str, Any]:
    """
    Safely parse JSON string with error handling.

    Args:
        json_str: JSON string to parse

    Returns:
        Parsed JSON as dictionary

    Raises:
        RuntimeError: If JSON parsing fails
    """
    try:
        return json.loads(json_str)
    except json.JSONDecodeError as e:
        raise RuntimeError(f"Failed to parse JSON: {e}")


def get_checkpoint_physical_time_utc(s3_client, bucket: str, cdc_metadata_key: str) -> int:
    """
    Read checkpoint-ts from CDC metadata file and return UTC timestamp.

    Args:
        s3_client: S3 client instance
        bucket: S3 bucket name
        cdc_metadata_key: S3 key of CDC metadata file

    Returns:
        Physical UTC timestamp (seconds since Unix epoch)

    Raises:
        RuntimeError: If metadata file cannot be read or parsed
    """
    try:
        response = s3_client.get_object(Bucket=bucket, Key=cdc_metadata_key)
        metadata_content = response["Body"].read().decode("utf-8")
        metadata_json = safe_json_parse(metadata_content)
        checkpoint_ts = metadata_json.get("checkpoint-ts")

        if checkpoint_ts is None:
            raise RuntimeError(f"Missing 'checkpoint-ts' in metadata file: {cdc_metadata_key}")
        if not isinstance(checkpoint_ts, int):
            raise RuntimeError(
                f"Invalid 'checkpoint-ts' type in metadata file: {cdc_metadata_key}, got {type(checkpoint_ts).__name__}"
            )
        if checkpoint_ts < 0:
            raise RuntimeError(f"Invalid negative 'checkpoint-ts' in metadata file: {cdc_metadata_key}")

        physical_ms = checkpoint_ts >> 18
        physical_time_utc_ts = physical_ms // 1000
        return physical_time_utc_ts
    except RuntimeError:
        raise
    except Exception as e:
        raise RuntimeError(f"Failed to get checkpoint physical time from {cdc_metadata_key}: {e}")


def wait_cdc_checkpoint(
    s3_client,
    bucket: str,
    cdc_metadata_key: str,
    timeout: int = 300,
) -> int:
    """
    Wait until CDC checkpoint physical time exceeds current UTC timestamp.

    Gets current UTC timestamp (t), then polls get_checkpoint_physical_time_utc
    until the returned value is greater than t.

    Args:
        s3_client: S3 client instance
        bucket: S3 bucket name
        cdc_metadata_key: S3 key of CDC metadata file
        timeout: Maximum time to wait in seconds

    Returns:
        The checkpoint physical time (seconds since Unix epoch) when condition is met

    Raises:
        TimeoutError: If checkpoint does not exceed t within timeout
    """
    import time

    t = int(time.time())
    start_time = time.time()

    while True:
        elapsed = time.time() - start_time
        if elapsed > timeout:
            raise TimeoutError(
                f"Verification timed out after {timeout} seconds while waiting for CDC checkpoint to exceed t={t}."
            )

        checkpoint_ts = get_checkpoint_physical_time_utc(s3_client, bucket, cdc_metadata_key)
        print(f"  checkpoint_ts={checkpoint_ts}, t={t}")
        if checkpoint_ts > t:
            return checkpoint_ts

        time.sleep(1)


def get_index_and_table_id(
    host: str,
    port: int,
    db_name: str,
    table_name: str,
    index_name: str,
    user: str = "root",
) -> Tuple[Optional[int], Optional[List[int]], Optional[bool]]:
    """
    Get index_id and table_id from TiDB.

    table_id is always a list when not None: for partitioned tables, list of
    TIDB_PARTITION_ID from information_schema.partitions; for non-partitioned
    tables, list of one TIDB_TABLE_ID from information_schema.tables.

    Args:
        host: TiDB host
        port: TiDB port
        db_name: Database name
        table_name: Table name
        index_name: Index name
        user: Database user

    Returns:
        Tuple of (index_id, table_id, is_partition). table_id is None or a list
        of ints (one element for non-partitioned, multiple for partitioned).
    """
    index_id, table_id, is_partition = None, None, None
    with mysql_connection(host, port, user) as connection:
        sql = f"SELECT index_id FROM information_schema.tidb_indexes WHERE table_schema = '{db_name}' and table_name = '{table_name}' and key_name = '{index_name}';"
        result = execute_sql(connection, sql)
        index_id = result[0][0] if result else None

        # Try partitioned table first: get TIDB_PARTITION_ID from information_schema.partitions
        sql = f"SELECT TIDB_PARTITION_ID FROM information_schema.partitions WHERE TABLE_SCHEMA = '{db_name}' AND TABLE_NAME = '{table_name}';"
        result = execute_sql(connection, sql)
        if result and result[0][0] is not None:
            table_id = [row[0] for row in result] if result else None
            is_partition = True
        else:
            # Non-partitioned table: get TIDB_TABLE_ID from information_schema.tables
            sql = f"SELECT TIDB_TABLE_ID FROM information_schema.tables WHERE TABLE_SCHEMA = '{db_name}' AND TABLE_NAME = '{table_name}';"
            result = execute_sql(connection, sql)
            table_id = [row[0] for row in result] if result else None
            is_partition = False

    return index_id, table_id, is_partition

def wait_shard_progress(
    host: str,
    port: int,
    db_name: str,
    table_name: str,
    index_name: str,
    bucket: str,
    prefix: str,
    s3_endpoint: str,
    s3_access_key: str = "minioadmin",
    s3_secret_key: str = "minioadmin",
    timeout: int = 300,
    user: str = "root",
) -> None:
    """
    Wait for shard progress to catch up with the latest CDC file.

    This function:
    1. Gets index_id and table_id from TiDB
    2. Waits for CDC files to stabilize
    3. Waits for all shards to process up to the latest CDC file

    Args:
        host: TiDB host
        port: TiDB port
        db_name: Database name
        table_name: Table name
        index_name: Index name to verify
        bucket: S3 bucket name
        prefix: S3 prefix (playground tag)
        s3_endpoint: S3 endpoint (e.g., "localhost:9000")
        s3_access_key: S3 access key
        s3_secret_key: S3 secret key
        timeout: Maximum time to wait in seconds
        user: Database user

    Raises:
        TimeoutError: If verification times out
        RuntimeError: If verification fails
    """
    import time

    print(f"Waiting for shard progress: {db_name}.{table_name} index={index_name}")
    start_time = time.time()

    metadata_start = time.time()
    index_id, table_id, is_partition = get_index_and_table_id(
        host, port, db_name, table_name, index_name, user
    )
    metadata_elapsed = time.time() - metadata_start
    print(f"  index_id={index_id}, table_id={table_id}, is_partition={is_partition}")

    s3_client = create_s3_client(s3_endpoint, s3_access_key, s3_secret_key)
    cdc_metadata_key = f"{prefix}/cdc/metadata"

    # Wait for CDC files to stabilize
    print(f"  Waiting for CDC files to stabilize...")
    cdc_wait_start = time.time()
    cdc_checkpoint_start = time.time()
    wait_cdc_checkpoint(s3_client, bucket, cdc_metadata_key, timeout)
    cdc_checkpoint_elapsed = time.time() - cdc_checkpoint_start
    cdc_last_file_elapsed = 0.0
    shard_poll_elapsed = 0.0

    for tid in (table_id or []):
        cdc_prefix = f"{prefix}/cdc/{tid}"
        cdc_last_file_start = time.time()
        cdc_s3_last_file = get_cdc_last_file(s3_client, bucket, cdc_prefix)
        cdc_last_file_elapsed += time.time() - cdc_last_file_start
        print(f"  Latest CDC file: {cdc_s3_last_file} (elapsed: {time.time() - cdc_wait_start:.1f}s)")

        # Wait for shard progress to catch up
        print(f"  Waiting for shard progress to catch up...")
        shard_poll_start = time.time()
        while True:
            elapsed = time.time() - start_time
            if elapsed > timeout:
                raise TimeoutError(
                    f"Verification timed out after {timeout} seconds while waiting for shard progress update."
                )

            time.sleep(1)
            with mysql_connection(host, port, user, timeout=60) as connection:
                sql = "SELECT distinct progress FROM tici.tici_shard_meta"
                if index_id is not None and tid is not None:
                    sql += f" WHERE index_id = {index_id} AND table_id = {tid};"
                else:
                    sql += ";"
                result = execute_sql(connection, sql)

                is_valid = True
                for row in result:
                    progress = safe_json_parse(row[0])
                    last_file = progress.get("cdc_s3_last_file")
                    if cdc_s3_last_file != last_file:
                        is_valid = False
                        break

                if is_valid:
                    break
        shard_poll_elapsed += time.time() - shard_poll_start

    elapsed = time.time() - start_time
    print(f"  Shard progress synced in {elapsed:.1f}s")

    # Currently, we have no way to confirm the data synchronization progress on readers.
    # Waiting for 10 seconds is a simple workaround.
    reader_grace_start = time.time()
    time.sleep(10)
    reader_grace_elapsed = time.time() - reader_grace_start
    total_elapsed = time.time() - start_time
    print(
        "  Wait timing: "
        f"metadata={metadata_elapsed:.1f}s "
        f"cdc_checkpoint={cdc_checkpoint_elapsed:.1f}s "
        f"cdc_last_file={cdc_last_file_elapsed:.1f}s "
        f"shard_poll={shard_poll_elapsed:.1f}s "
        f"reader_grace={reader_grace_elapsed:.1f}s "
        f"total={total_elapsed:.1f}s"
    )