#!/usr/bin/env python3
"""Migrate ClickHouse column names from _ms to _us suffix.

Renames:
  - open_deviation_bars: close_time_ms -> close_time_us
  - open_deviation_bars: open_time_ms  -> open_time_us
  - data_freshness:      last_close_time_ms -> last_close_time_us

Safety: ALTER TABLE RENAME COLUMN is metadata-only in ClickHouse -- instant,
no data rewrite. Safe for 42M+ rows. close_time_us is NOT in ORDER BY
(changed to first_agg_trade_id in Issue #158), so RENAME works directly.

Usage:
  python scripts/migrate_column_rename_us.py           # Dry-run (default)
  python scripts/migrate_column_rename_us.py --apply   # Execute migration
"""

from __future__ import annotations

import argparse
import os
import sys


def _get_client():  # noqa: ANN202
    """Connect to ClickHouse using environment configuration."""
    import clickhouse_connect

    hosts = os.environ.get("OPENDEVIATIONBAR_CH_HOSTS", "localhost")
    host = hosts.split(",")[0].strip()
    port = int(os.environ.get("OPENDEVIATIONBAR_CH_PORT", "8123"))
    return clickhouse_connect.get_client(host=host, port=port)


def _get_columns(client, database: str, table: str) -> set[str]:
    """Get column names for a table."""
    result = client.query(
        "SELECT name FROM system.columns "
        f"WHERE database = '{database}' AND table = '{table}'"
    )
    return {row[0] for row in result.result_rows}


def _get_row_count(client, database: str, table: str) -> int:
    """Get approximate row count for a table."""
    result = client.query(f"SELECT count() FROM {database}.{table}")
    return result.result_rows[0][0]


# Rename operations: (database, table, old_name, new_name)
RENAMES = [
    (
        "opendeviationbar_cache",
        "open_deviation_bars",
        "close_time_ms",  # old name on production CH
        "close_time_us",
    ),
    (
        "opendeviationbar_cache",
        "open_deviation_bars",
        "open_time_ms",  # old name on production CH
        "open_time_us",
    ),
    (
        "opendeviationbar_cache",
        "data_freshness",
        "last_close_time_ms",  # old name on production CH
        "last_close_time_us",
    ),
]


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Rename _ms columns to _us in ClickHouse"
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="Execute the migration (default: dry-run)",
    )
    args = parser.parse_args()

    client = _get_client()
    print(f"Connected to ClickHouse at {client.server_host_name}")

    # Check current state for each rename
    already_migrated = []
    pending = []

    for db, table, old_col, new_col in RENAMES:
        columns = _get_columns(client, db, table)

        if new_col in columns:
            already_migrated.append((db, table, old_col, new_col))
            print(f"  [DONE] {db}.{table}: {new_col} already exists")
        elif old_col in columns:
            pending.append((db, table, old_col, new_col))
            print(f"  [PENDING] {db}.{table}: {old_col} -> {new_col}")
        else:
            print(
                f"  [ERROR] {db}.{table}: neither {old_col} nor {new_col} found!",
                file=sys.stderr,
            )
            return 1

    if not pending:
        print("\nAll columns already migrated. Nothing to do.")
        return 0

    # Show row counts for verification
    print("\nRow counts (pre-migration):")
    tables_seen: set[str] = set()
    pre_counts: dict[str, int] = {}
    for db, table, _, _ in pending:
        key = f"{db}.{table}"
        if key not in tables_seen:
            tables_seen.add(key)
            count = _get_row_count(client, db, table)
            pre_counts[key] = count
            print(f"  {key}: {count:,} rows")

    # Show ALTER statements
    print("\nALTER statements:")
    for db, table, old_col, new_col in pending:
        stmt = (
            f"ALTER TABLE {db}.{table} RENAME COLUMN {old_col} TO {new_col}"
        )
        print(f"  {stmt}")

    if not args.apply:
        print("\n[DRY RUN] No changes made. Pass --apply to execute.")
        return 0

    # Execute renames
    print("\nExecuting migration...")
    for db, table, old_col, new_col in pending:
        stmt = (
            f"ALTER TABLE {db}.{table} RENAME COLUMN {old_col} TO {new_col}"
        )
        print(f"  Running: {stmt}")
        client.command(stmt)
        print(f"  OK: {old_col} -> {new_col}")

    # Verify row counts unchanged
    print("\nRow counts (post-migration):")
    for key, pre_count in pre_counts.items():
        db, table = key.split(".")
        post_count = _get_row_count(client, db, table)
        status = "OK" if post_count == pre_count else "MISMATCH"
        print(f"  {key}: {post_count:,} rows [{status}]")
        if post_count != pre_count:
            print(
                f"  [ERROR] Row count changed: {pre_count:,} -> {post_count:,}",
                file=sys.stderr,
            )
            return 1

    # Verify columns renamed
    print("\nColumn verification (post-migration):")
    for db, table, old_col, new_col in pending:
        columns = _get_columns(client, db, table)
        has_new = new_col in columns
        has_old = old_col in columns
        if has_new and not has_old:
            print(f"  {db}.{table}: {new_col} present, {old_col} gone [OK]")
        else:
            print(
                f"  [ERROR] {db}.{table}: new={has_new} old={has_old}",
                file=sys.stderr,
            )
            return 1

    print("\nMigration complete.")
    return 0


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