#!/usr/bin/env python3
"""
Migration script to transfer data from Akumuli to VictoriaMetrics.

Usage:
    # Auto-discover series from Akumuli database:
    python migrate_from_akumuli.py --akumuli-http http://localhost:8181 \
                                    --akumuli-db /path/to/db.akumuli \
                                    --victoria-tcp localhost:4242 \
                                    --victoria-http http://localhost:8428 \
                                    [--start-time TIMESTAMP] [--end-time TIMESTAMP]

    # Or specify series manually:
    python migrate_from_akumuli.py --akumuli-http http://localhost:8181 \
                                    --victoria-tcp localhost:4242 \
                                    --victoria-http http://localhost:8428 \
                                    --series series_name \
                                    [--series "another_series{tag=value}"] \
                                    [--start-time TIMESTAMP] [--end-time TIMESTAMP]
"""
from __future__ import annotations

import anyio
import argparse
import sqlite3
import sys
import traceback

import httpx

# Import from the old asyncakumuli package for reading
try:
    import asyncakumuli
    HAS_AKUMULI = True
except ImportError:
    HAS_AKUMULI = False
    print(
        "Warning: asyncakumuli package not found. Install it to migrate from Akumuli.",
        file=sys.stderr,
    )

# Import from the new asyncvictoria package for writing
from asyncvictoria import DS as VDS
from asyncvictoria import Entry, sanitize_entry
from asyncvictoria import connect as victoria_connect


def get_all_series_from_db(db_path: str):
    """Get all series IDs from Akumuli's SQLite database.

    Args:
        db_path: Path to Akumuli's db.akumuli SQLite database

    Returns:
        List of series IDs (strings)
    """
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        cursor.execute("SELECT DISTINCT series_id FROM akumuli_series")
        series_list = [row[0] for row in cursor.fetchall()]
        conn.close()
        return series_list
    except sqlite3.Error as e:
        print(f"Error reading Akumuli database: {e}", file=sys.stderr)
        raise


async def migrate_series(
    akumuli_session,
    victoria_client,
    series: str,
    tags: dict,
    akumuli_url: str,
    start_time: float | None = None,
    end_time: float | None = None,
    batch_size: int = 1000,
):
    """Migrate a single series from Akumuli to VictoriaMetrics."""

    print(f"Migrating series: {series} with tags: {tags}")

    count = 0

    # Read data from Akumuli
    try:
        data_gen = asyncakumuli.get_data(
            akumuli_session,
            series,
            tags=tags,
            t_start=start_time,
            t_end=end_time,
            url=f"{akumuli_url}/api/query"
        )

        async for akumuli_entry in data_gen:
            # Create VictoriaMetrics Entry from Akumuli entry
            victoria_entry = Entry(
                value=akumuli_entry.value,
                series=akumuli_entry.series,
                time=akumuli_entry.time,
                tags=akumuli_entry.tags,
                mode=VDS.gauge,
            )

            # Sanitize metric and label names for VictoriaMetrics/Prometheus compatibility
            sanitized_entry = sanitize_entry(victoria_entry)

            # Write to VictoriaMetrics
            await victoria_client.put(sanitized_entry)
            count += 1

            if count % batch_size == 0:
                await victoria_client.flush()
                print(f"  Migrated {count} points...", end="\r")

        # Final flush
        await victoria_client.flush()
        print(f"  Migrated {count} points total for {series}")

    except Exception as e:
        print(f"Error migrating series {series}: {e}", file=sys.stderr)
        raise

    return count


async def migrate_all(
    akumuli_http: str,
    victoria_tcp: str,
    victoria_http: str,
    series_list: list | None = None,
    start_time: float | None = None,
    end_time: float | None = None,
    akumuli_db: str | None = None,
):
    """Main migration function."""

    if not HAS_AKUMULI:
        print("Error: asyncakumuli package is required for migration.", file=sys.stderr)
        return 1

    # Parse victoria_tcp into host and port
    if ":" in victoria_tcp:
        vic_host, vic_port = victoria_tcp.rsplit(":", 1)
        vic_port = int(vic_port)
    else:
        vic_host = victoria_tcp
        vic_port = 4242

    total_count = 0

    async with (
        httpx.AsyncClient(timeout=600) as akumuli_session,
        victoria_connect(host=vic_host, port=vic_port) as victoria_client,
    ):
        # If no series list provided, try to discover from database
        if not series_list:
            if akumuli_db:
                print(f"Auto-discovering series from Akumuli database: {akumuli_db}")
                series_list = get_all_series_from_db(akumuli_db)
                if not series_list:
                    print("No series found in database.", file=sys.stderr)
                    return 1
                print(f"Found {len(series_list)} series to migrate:")
                for s in series_list:
                    print(f"  - {s}")
            else:
                print(
                    "Error: Either --series or --akumuli-db must be specified.",
                    file=sys.stderr,
                )
                print(
                    "Use --series to specify series manually, "
                    "or --akumuli-db to auto-discover."
                )
                return 1

        for series_spec in series_list:
            # Parse series spec: "series_name" or "series_name{tag1=val1,tag2=val2}"
            if "{" in series_spec:
                series_name, tags_str = series_spec.split("{", 1)
                tags_str = tags_str.rstrip("}")
                tags = {}
                for tag_pair in tags_str.split(","):
                    k, v = tag_pair.split("=", 1)
                    tags[k.strip()] = v.strip()
            else:
                series_name = series_spec
                tags = {}

            count = await migrate_series(
                akumuli_session,
                victoria_client,
                series_name,
                tags,
                akumuli_http,
                start_time,
                end_time,
            )
            total_count += count

        # Force VictoriaMetrics to flush all data to disk
        print("\n=== Forcing VictoriaMetrics to flush data to disk ===")
        try:
            async with httpx.AsyncClient() as flush_client:
                response = await flush_client.get(f"{victoria_http}/internal/force_flush")
                response.raise_for_status()
                print("✓ Data flushed to disk successfully")
        except Exception as e:
            print(f"Warning: Failed to force flush: {e}", file=sys.stderr)

    print(f"\nMigration complete. Total points migrated: {total_count}")
    return 0


def main() -> int:
    """Main entry point for migration script."""
    parser = argparse.ArgumentParser(
        description="Migrate data from Akumuli to VictoriaMetrics"
    )
    parser.add_argument(
        "--akumuli-http",
        default="http://localhost:8181",
        help="Akumuli HTTP endpoint (default: http://localhost:8181)"
    )
    parser.add_argument(
        "--akumuli-db",
        help="Path to Akumuli SQLite database (db.akumuli) for auto-discovering series"
    )
    parser.add_argument(
        "--victoria-tcp",
        default="localhost:4242",
        help="VictoriaMetrics OpenTSDB TCP endpoint (default: localhost:4242)"
    )
    parser.add_argument(
        "--victoria-http",
        default="http://localhost:8428",
        help="VictoriaMetrics HTTP endpoint (default: http://localhost:8428)"
    )
    parser.add_argument(
        "--series",
        action="append",
        help='Series to migrate (can be specified multiple times). '
             'Format: "series_name" or "series_name{tag1=val1,tag2=val2}"'
    )
    parser.add_argument(
        "--start-time",
        type=float,
        help="Start timestamp (Unix timestamp in seconds)"
    )
    parser.add_argument(
        "--end-time",
        type=float,
        help="End timestamp (Unix timestamp in seconds)"
    )

    args = parser.parse_args()

    if not args.series and not args.akumuli_db:
        print("Error: Either --series or --akumuli-db must be specified", file=sys.stderr)
        return 1

    try:
        return anyio.run(
            migrate_all,
            args.akumuli_http,
            args.victoria_tcp,
            args.victoria_http,
            args.series,
            args.start_time,
            args.end_time,
            args.akumuli_db
        )
    except KeyboardInterrupt:
        print("\nMigration interrupted by user")
        return 1
    except Exception as e:
        print(f"Migration failed: {e}", file=sys.stderr)
        traceback.print_exc()
        return 1


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