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

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

CSV_FILE = "numbers.csv"
TABLE_NAME = "phone_numbers"
COLUMN_NAME = "phone_number"
RANDOM_COLUMN_NAME = "random_code"

NUM_WORKERS = 12
BATCH_SIZE = 2_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 get_file_chunks(filename, num_chunks):
    """تقسیم فایل به بخش‌های مساوی"""
    print(f"Calculating chunks for {filename}...")
    file_size = os.path.getsize(filename)
    print(f"File size: {file_size / (1024 ** 3):.2f} GB")
    chunk_size = file_size // num_chunks
    chunks = []
    with open(filename, 'rb') as f:
        f.readline()
        start = f.tell()
        for i in range(num_chunks):
            if i < num_chunks - 1:
                f.seek(chunk_size * (i + 1))
                f.readline()
                end = f.tell()
                chunks.append((start, end))
                start = end
            else:
                chunks.append((start, file_size))
    print(f"Created {len(chunks)} chunks")
    return chunks


def count_total_rows():
    """شمارش کل رکوردها"""
    print("Counting total rows...")
    count = 0
    with open(CSV_FILE, 'r') as f:
        next(f)
        for _ in f:
            count += 1
    return count


def load_chunk(args):
    """هر Worker یک فایل temp می‌سازد و با COPY لود می‌کند"""
    chunk_start, chunk_end, worker_id, rows_in_chunk = args
    print(f"[Worker {worker_id}] Starting... chunk: {chunk_start} to {chunk_end}")

    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

        # ═══════════════════════════════════════════════════════
        # محاسبه شروع عدد توالی برای این Worker
        # ═══════════════════════════════════════════════════════
        phone_start = PHONE_START + (worker_id * rows_in_chunk)
        current_phone = phone_start
        print(f"[Worker {worker_id}] Phone range: {phone_start:,} to {phone_start + rows_in_chunk - 1:,}")

        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:
            with open(CSV_FILE, 'r', buffering=8192 * 1024) as f:
                f.seek(chunk_start)

                # رد کردن خط header اگر در ابتدای chunk باشد
                if chunk_start > 0:
                    f.readline()

                while f.tell() < chunk_end:
                    line = f.readline()
                    if line.strip():
                        # ═══════════════════════════════════════
                        # استفاده از عدد توالی به جای خواندن از CSV
                        # ═══════════════════════════════════════
                        random_code = random.randint(1000, 9999)
                        temp_f.write(f"{current_phone},{random_code}\n")

                        current_phone += 1
                        loaded += 1

                        if loaded % 5_000_000 == 0:
                            print(f"[Worker {worker_id}] Wrote {loaded:,} rows to temp file")

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

        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(start_check=100_000_000, sample_points=None):
    """بررسی صحت داده‌های لود شده"""
    if sample_points is None:
        sample_points = [0, 10_000_000, 50_000_000, 100_000_000, 124_999_999]

    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")
    print("=" * 60)
    print(f"Phone number range: {PHONE_START:,} to {PHONE_START + 125_000_000 - 1:,}")
    print(f"Total records: 125,000,000\n")

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

    print(f"CSV rows: {total_rows:,}")
    print(f"Using {num_workers} workers")
    print(f"Batch size: {BATCH_SIZE:,}\n")

    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")

        # تقسیم فایل
        print("Calculating file chunks...")
        chunks = get_file_chunks(CSV_FILE, num_workers)

        # ═══════════════════════════════════════════════════════
        # محاسبه تعداد رکوردها در هر chunk
        # ═══════════════════════════════════════════════════════
        rows_per_chunk = total_rows // num_workers

        args = [
            (start, end, i, rows_per_chunk)
            for i, (start, end) in enumerate(chunks)
        ]

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

        # اجرای موازی
        with ProcessPoolExecutor(max_workers=num_workers) as executor:
            futures = [executor.submit(load_chunk, arg) for arg in args]
            total = 0

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

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

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

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

        # ═══════════════════════════════════════════════════════
        # ساخت ایندکس‌ها با تأیید
        # ═══════════════════════════════════════════════════════
        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_rows / 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()