#!/usr/bin/env python3
# Issue #158 Phase 7: Migrate ORDER BY from close_time_ms to first_agg_trade_id
"""Migrate ClickHouse open_deviation_bars ORDER BY for native dedup.

Changes ORDER BY from:
    (symbol, threshold_decimal_bps, ouroboros_mode, close_time_ms)
to:
    (symbol, threshold_decimal_bps, ouroboros_mode, first_agg_trade_id)

This enables ReplacingMergeTree to natively deduplicate bars with the same
first_agg_trade_id (from different writers producing different close_time_ms
for the same trades).

Legacy rows with first_agg_trade_id=0 are excluded from migration — they
would collapse to 1 row per (symbol, threshold, ouroboros) under the new
ORDER BY and need re-population anyway.

Usage:
    python scripts/migrate_order_by_first_agg_trade_id.py              # Dry run
    python scripts/migrate_order_by_first_agg_trade_id.py --apply      # Execute
"""

from __future__ import annotations

import argparse
import logging
import sys

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger("migrate-order-by")

TABLE = "opendeviationbar_cache.open_deviation_bars"
TABLE_V2 = "opendeviationbar_cache.open_deviation_bars_v2"
TABLE_BACKUP = "opendeviationbar_cache.open_deviation_bars_backup"


def get_column_list(client: object) -> list[str]:
    """Get all columns from existing table for the INSERT ... SELECT."""
    result = client.query(
        "SELECT name FROM system.columns "
        "WHERE database = 'opendeviationbar_cache' "
        "AND table = 'open_deviation_bars' "
        "ORDER BY position"
    )
    return [row[0] for row in result.result_rows]


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Migrate ORDER BY from close_time_ms to first_agg_trade_id.",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="Execute migration (default is dry-run).",
    )
    args = parser.parse_args()

    try:
        import clickhouse_connect

        client = clickhouse_connect.get_client(host="localhost")
    except Exception:
        logger.exception("Cannot connect to ClickHouse")
        return 2

    # Check current ORDER BY
    result = client.query(
        "SELECT sorting_key FROM system.tables "
        "WHERE database = 'opendeviationbar_cache' "
        "AND name = 'open_deviation_bars'"
    )
    if not result.result_rows:
        logger.error("Table %s not found", TABLE)
        return 2

    current_order = result.result_rows[0][0]
    logger.info("Current ORDER BY: %s", current_order)

    if "first_agg_trade_id" in current_order:
        logger.info("Already migrated — nothing to do.")
        return 0

    # Count total rows and rows with valid trade IDs
    counts = client.query(
        f"SELECT count(), countIf(first_agg_trade_id > 0) FROM {TABLE}"
    )
    total, valid = counts.result_rows[0]
    legacy = total - valid
    logger.info(
        "Total rows: %d | Valid (first_agg_trade_id > 0): %d | Legacy (=0): %d",
        total, valid, legacy,
    )

    if not args.apply:
        logger.info("Dry run — use --apply to execute migration.")
        logger.info(
            "Migration will: create v2 table → copy %d valid rows → rename → add skip index",
            valid,
        )
        if legacy > 0:
            logger.warning(
                "%d legacy rows (first_agg_trade_id=0) will NOT be migrated. "
                "These need re-population with current code.",
                legacy,
            )
        return 0

    # Get CREATE TABLE statement and modify it
    columns = get_column_list(client)
    columns_str = ", ".join(columns)

    # Step 1: Get original CREATE TABLE and recreate with new ORDER BY
    create_result = client.query(f"SHOW CREATE TABLE {TABLE}")
    create_sql = create_result.result_rows[0][0]

    # Replace ORDER BY clause
    create_v2 = create_sql.replace(
        "ORDER BY (symbol, threshold_decimal_bps, ouroboros_mode, close_time_ms)",
        "ORDER BY (symbol, threshold_decimal_bps, ouroboros_mode, first_agg_trade_id)",
    )
    # Point to new table name
    create_v2 = create_v2.replace(
        f"CREATE TABLE {TABLE}",
        f"CREATE TABLE {TABLE_V2}",
    )
    # Also handle the simpler form without schema prefix
    create_v2 = create_v2.replace(
        "CREATE TABLE open_deviation_bars",
        f"CREATE TABLE {TABLE_V2}",
    )

    logger.info("Step 1: Creating %s with new ORDER BY...", TABLE_V2)
    client.command(f"DROP TABLE IF EXISTS {TABLE_V2}")
    client.command(create_v2)
    logger.info("Created %s", TABLE_V2)

    # Step 2: Copy valid rows
    logger.info("Step 2: Copying %d valid rows (first_agg_trade_id > 0)...", valid)
    client.command(
        f"INSERT INTO {TABLE_V2} ({columns_str}) "
        f"SELECT {columns_str} FROM {TABLE} WHERE first_agg_trade_id > 0",
        settings={"max_execution_time": 3600, "max_partitions_per_insert_block": 0},
    )

    # Verify count
    v2_count = client.query(f"SELECT count() FROM {TABLE_V2}")
    migrated = v2_count.result_rows[0][0]
    logger.info("Migrated %d rows (expected %d)", migrated, valid)

    if migrated != valid:
        logger.error(
            "Row count mismatch! Expected %d, got %d. Aborting.",
            valid, migrated,
        )
        client.command(f"DROP TABLE {TABLE_V2}")
        return 1

    # Step 3: Atomic rename
    logger.info("Step 3: Renaming tables (atomic swap)...")
    client.command(f"DROP TABLE IF EXISTS {TABLE_BACKUP}")
    client.command(
        f"RENAME TABLE {TABLE} TO {TABLE_BACKUP}, "
        f"{TABLE_V2} TO {TABLE}"
    )
    logger.info("Rename complete: %s → backup, %s → %s", TABLE, TABLE_V2, TABLE)

    # Step 4: Add skip index for close_time_ms range queries
    logger.info("Step 4: Adding skip index for close_time_ms...")
    client.command(
        f"ALTER TABLE {TABLE} "
        "ADD INDEX IF NOT EXISTS idx_close_time close_time_ms "
        "TYPE minmax GRANULARITY 8192"
    )
    logger.info("Skip index added")

    # Verify new ORDER BY
    result = client.query(
        "SELECT sorting_key FROM system.tables "
        "WHERE database = 'opendeviationbar_cache' "
        "AND name = 'open_deviation_bars'"
    )
    new_order = result.result_rows[0][0]
    logger.info("New ORDER BY: %s", new_order)

    logger.info(
        "Migration complete. Backup table: %s (drop when satisfied).", TABLE_BACKUP,
    )
    if legacy > 0:
        logger.warning(
            "%d legacy rows (first_agg_trade_id=0) are in the backup only. "
            "Re-populate these with current code to get valid trade IDs.",
            legacy,
        )

    return 0


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