import psycopg2
from psycopg2 import sql
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, as_completed
import time
import random
import tempfile
import os

# ═══════════════════════════════════════════════════════════════
# تنظیمات
# ═══════════════════════════════════════════════════════════════
DB_CONFIG = {
    "host": "localhost",
    "port": 5432,
    "database": "numbers_db",
    "user": "myuser",
    "password": "mypassword"
}

TABLE_NAME = "phone_numbers"
COLUMN_NAME = "phone_number"
RANDOM_COLUMN_NAME = "random_code"

NUM_WORKERS = 12
BATCH_SIZE = 2_000_000
TOTAL_RECORDS = 125_000_000  # تعداد کل رکوردها - مستقیماً مشخص می‌شود

# ═══════════════════════════════════════════════════════════════
# شروع توالی عددی
# ═══════════════════════════════════════════════════════════════
PHONE_START = 100_000_000  # کمترین عدد ۹ رقمی


def optimize_connection(conn):
    """تنظیمات بهینه برای سرعت حداکثری"""
    with conn.cursor() as cur:
        cur.execute("SET work_mem = '512MB'")
        cur.execute("SET synchronous_commit = OFF")
        cur.execute("SET maintenance_work_mem = '2GB'")
        cur.execute("SET max_parallel_workers_per_gather = 4")
        cur.execute("SET random_page_cost = 1.1")
        conn.commit()


def load_chunk_sequential(worker_id, start_phone, rows_count):
    """
    هر Worker داده‌های توالی خود را تولید کرده و با COPY لود می‌کند
    بدون نیاز به هیچ فایل CSV
    """
    print(f"[Worker {worker_id}] Starting... Phone range: {start_phone:,} to {start_phone + rows_count - 1:,}")

    conn = psycopg2.connect(**DB_CONFIG)
    temp_file = None

    try:
        optimize_connection(conn)
        print(f"[Worker {worker_id}] Connected to DB")

        start_time = time.time()
        loaded = 0

        # ایجاد فایل موقت
        temp_fd, temp_file = tempfile.mkstemp(suffix='.csv', prefix=f'worker_{worker_id}_')
        os.close(temp_fd)

        # تولید مستقیم داده‌ها در فایل موقت
        with open(temp_file, 'w', buffering=8192 * 1024) as temp_f:
            current_phone = start_phone

            for i in range(rows_count):
                random_code = random.randint(1000, 9999)
                temp_f.write(f"{current_phone},{random_code}\n")
                current_phone += 1
                loaded += 1

                # نمایش پیشرفت هر worker
                if loaded % 5_000_000 == 0:
                    print(f"[Worker {worker_id}] Generated {loaded:,} rows so far...")

        print(f"[Worker {worker_id}] Temp file created: {loaded:,} rows")

        # لود با COPY
        with open(temp_file, 'r') as f:
            with conn.cursor() as cur:
                cur.copy_expert(
                    f"COPY {TABLE_NAME} ({COLUMN_NAME}, {RANDOM_COLUMN_NAME}) FROM STDIN WITH CSV",
                    f
                )
                conn.commit()

        elapsed = time.time() - start_time
        rate = loaded / elapsed if elapsed > 0 else 0
        print(f"[Worker {worker_id}] Done! {loaded:,} rows in {elapsed:.1f}s ({rate:,.0f} rows/sec)")

        return loaded

    except Exception as e:
        print(f"[Worker {worker_id}] ERROR: {e}")
        raise
    finally:
        conn.close()
        if temp_file and os.path.exists(temp_file):
            os.remove(temp_file)


def create_table(conn):
    """ساخت جدول Unlogged برای سرعت بیشتر"""
    with conn.cursor() as cur:
        cur.execute(f"""
            DROP TABLE IF EXISTS {TABLE_NAME} CASCADE;
            CREATE UNLOGGED TABLE {TABLE_NAME} (
                {COLUMN_NAME} BIGINT NOT NULL,
                {RANDOM_COLUMN_NAME} INTEGER NOT NULL
            ) WITH (fillfactor = 100);
        """)
        conn.commit()


def index_exists(conn, index_name):
    """بررسی وجود ایندکس"""
    with conn.cursor() as cur:
        cur.execute("""
            SELECT EXISTS (
                SELECT 1 FROM pg_indexes 
                WHERE indexname = %s
            )
        """, (index_name,))
        return cur.fetchone()[0]


def create_index(conn, column_name):
    """ساخت ایندکس با بررسی و تأیید"""
    index_name = f"idx_{TABLE_NAME}_{column_name}"
    print(f"\n>>> Creating index on '{column_name}'...")

    if index_exists(conn, index_name):
        print(f"    Index '{index_name}' already exists, dropping...")
        with conn.cursor() as cur:
            cur.execute(f"DROP INDEX IF EXISTS {index_name}")
            conn.commit()

    with conn.cursor() as cur:
        start = time.time()
        cur.execute(f"CREATE INDEX {index_name} ON {TABLE_NAME} ({column_name})")
        conn.commit()
        elapsed = time.time() - start
        print(f"    ✓ Index '{index_name}' created in {elapsed:.1f}s")

    if index_exists(conn, index_name):
        print(f"    ✓ Verified: Index '{index_name}' exists")
        return True
    else:
        print(f"    ✗ ERROR: Index '{index_name}' was not created!")
        return False


def verify_indexes():
    """تأیید نهایی وجود همه ایندکس‌ها"""
    conn = psycopg2.connect(**DB_CONFIG)
    try:
        print("\n" + "=" * 50)
        print("INDEX VERIFICATION")
        print("=" * 50)
        with conn.cursor() as cur:
            cur.execute("""
                SELECT indexname, indexdef 
                FROM pg_indexes 
                WHERE tablename = %s
                ORDER BY indexname
            """, (TABLE_NAME,))
            indexes = cur.fetchall()

            if not indexes:
                print("✗ NO INDEXES FOUND!")
                return False

            print(f"\nFound {len(indexes)} index(es):\n")
            for idx_name, idx_def in indexes:
                print(f"  ✓ {idx_name}")
                print(f"    Definition: {idx_def}\n")

            phone_idx_exists = any('phone_number' in idx[1] for idx in indexes)
            random_idx_exists = any('random_code' in idx[1] for idx in indexes)

            print("-" * 50)
            print(f"phone_number indexed: {'✓ YES' if phone_idx_exists else '✗ NO'}")
            print(f"random_code indexed:   {'✓ YES' if random_idx_exists else '✗ NO'}")
            print("-" * 50)

            return phone_idx_exists and random_idx_exists
    finally:
        conn.close()


def verify_data():
    """بررسی صحت داده‌های لود شده"""
    sample_points = [0, 10_000_000, 50_000_000, 100_000_000, TOTAL_RECORDS - 1]

    print("\n" + "=" * 50)
    print("DATA VERIFICATION")
    print("=" * 50)

    conn = psycopg2.connect(**DB_CONFIG)
    try:
        with conn.cursor() as cur:
            for point in sample_points:
                expected_phone = PHONE_START + point
                cur.execute(
                    f"SELECT {COLUMN_NAME}, {RANDOM_COLUMN_NAME} FROM {TABLE_NAME} LIMIT 1 OFFSET %s",
                    (point,)
                )
                row = cur.fetchone()
                if row:
                    status = "✓" if row[0] == expected_phone else "✗"
                    print(f"  {status} Offset {point:>12,}: Expected {expected_phone:,}, Got {row[0]:,}")
    finally:
        conn.close()


def main():
    print("=" * 60)
    print("PostgreSQL Bulk Loader - Sequential Phone Numbers (No CSV)")
    print("=" * 60)
    print(f"Total records: {TOTAL_RECORDS:,}")
    print(f"Phone number range: {PHONE_START:,} to {PHONE_START + TOTAL_RECORDS - 1:,}")
    print(f"Using {NUM_WORKERS} workers")
    print(f"Batch size per worker: {TOTAL_RECORDS // NUM_WORKERS:,}\n")

    num_workers = min(NUM_WORKERS, mp.cpu_count())

    # محاسبه تعداد رکورد برای هر worker
    rows_per_worker = TOTAL_RECORDS // num_workers
    remaining_rows = TOTAL_RECORDS % num_workers

    start_time = time.time()

    try:
        # اتصال اولیه و ساخت جدول
        conn = psycopg2.connect(**DB_CONFIG)
        print("Connected to DB, creating UNLOGGED table...")
        optimize_connection(conn)
        create_table(conn)
        conn.close()
        print("Table ready\n")

        # آماده سازی آرگومان‌ها برای workers
        args = []
        current_start = PHONE_START

        for i in range(num_workers):
            rows_count = rows_per_worker + (1 if i < remaining_rows else 0)
            args.append((i, current_start, rows_count))
            current_start += rows_count

        print("Starting workers...\n")

        # اجرای موازی
        total_loaded = 0
        with ProcessPoolExecutor(max_workers=num_workers) as executor:
            futures = {executor.submit(load_chunk_sequential, *arg): arg for arg in args}

            for future in as_completed(futures):
                try:
                    loaded = future.result()
                    total_loaded += loaded
                except Exception as e:
                    print(f"Worker failed: {e}")

                elapsed = time.time() - start_time
                rate = total_loaded / elapsed if elapsed > 0 else 0
                remaining = (TOTAL_RECORDS - total_loaded) / rate if rate > 0 else 0
                percent = (total_loaded * 100) // TOTAL_RECORDS if TOTAL_RECORDS > 0 else 0

                print(f"\n>>> Progress: {total_loaded:,}/{TOTAL_RECORDS:,} ({percent}%) | "
                      f"Speed: {rate:,.0f} rows/sec | "
                      f"ETA: {remaining / 60:.1f} min")

        print(f"\n✓ Total rows loaded: {total_loaded:,}")

        # ساخت ایندکس‌ها
        print("\n" + "=" * 60)
        print("CREATING INDEXES")
        print("=" * 60)

        conn = psycopg2.connect(**DB_CONFIG)
        success1 = create_index(conn, COLUMN_NAME)
        conn.close()

        # تأیید نهایی
        all_ok = verify_indexes()
        verify_data()

        total_time = time.time() - start_time

        print(f"\n✓ Done! Total time: {total_time / 60:.1f} minutes")
        print(f"✓ Average speed: {TOTAL_RECORDS / total_time:,.0f} rows/sec")

        if all_ok:
            print("\n✓✓✓ ALL INDEXES CREATED SUCCESSFULLY ✓✓✓")
        else:
            print("\n✗✗✗ SOME INDEXES MISSING! ✗✗✗")

    except Exception as e:
        print(f"\n✗ Error: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()