"""Generate 500k test messages, message embeddings, and documents using raw psycopg3.

This script creates:
- 1 workspace
- 3 peers
- 3 sessions (with peers added to sessions via session_peers table)
- 500k messages distributed randomly across peers and sessions
- 500k message embeddings (one per message with constant zero vector)
- 1 collection per peer (3 total)
- 500k documents distributed across the collections

Usage:
    python manual_load.py [schema_name]

    If schema_name is not provided, uses DB_SCHEMA environment variable.

Environment variables required:
    DB_CONNECTION_URI: PostgreSQL connection URI
    DB_SCHEMA: Database schema to use (optional if passed as argument)
"""

import os
import random
import sys
import time

import psycopg
from dotenv import load_dotenv
from faker import Faker
from nanoid import generate as generate_nanoid

load_dotenv()

# Configuration
NUM_MESSAGES = 500_000
NUM_DOCUMENTS = 500_000
BATCH_SIZE = 10_000  # Batch size for inserts (messages and documents)
EMBEDDING_BATCH_SIZE = 1_000  # Smaller batch size for embeddings (they're large)
MAX_RETRIES = 3  # Number of retries for failed operations
RETRY_DELAY = 5  # Seconds to wait before retrying
NUM_PEERS = 3
NUM_SESSIONS = 3

# Load environment variables
DB_CONNECTION_URI = os.getenv("DB_CONNECTION_URI")

# Get schema from command line argument or environment variable
if len(sys.argv) > 1:
    DB_SCHEMA = sys.argv[1]
else:
    DB_SCHEMA = os.getenv("DB_SCHEMA")

if not DB_CONNECTION_URI:
    raise ValueError("DB_CONNECTION_URI environment variable is required")
if not DB_SCHEMA:
    raise ValueError(
        "DB_SCHEMA must be provided as command line argument or environment variable"
    )

# Convert SQLAlchemy-style URI to plain psycopg URI
# SQLAlchemy uses postgresql+psycopg://, psycopg uses postgresql://
if DB_CONNECTION_URI.startswith("postgresql+psycopg://"):
    DB_CONNECTION_URI = DB_CONNECTION_URI.replace("postgresql+psycopg://", "postgresql://")
elif DB_CONNECTION_URI.startswith("postgres+psycopg://"):
    DB_CONNECTION_URI = DB_CONNECTION_URI.replace("postgres+psycopg://", "postgresql://")

# Initialize Faker
fake = Faker()

# Constant embedding - all zeros for speed (1536 dimensions for OpenAI embeddings)
CONSTANT_EMBEDDING = str([0.0] * 1536)

# Generate IDs
WORKSPACE_ID = generate_nanoid()
WORKSPACE_NAME = "load_test_workspace"

PEER_IDS = [generate_nanoid() for _ in range(NUM_PEERS)]
PEER_NAMES = [f"load_peer_{i + 1}" for i in range(NUM_PEERS)]

SESSION_IDS = [generate_nanoid() for _ in range(NUM_SESSIONS)]
SESSION_NAMES = [f"load_session_{i + 1}" for i in range(NUM_SESSIONS)]

COLLECTION_IDS = [generate_nanoid() for _ in range(NUM_PEERS)]
# Collections use observer/observed paradigm - each peer observes itself
COLLECTION_OBSERVERS = PEER_NAMES
COLLECTION_OBSERVED = PEER_NAMES


def check_existing_data(conn, schema):
    """Check what data already exists in the database."""
    print("\nChecking existing data...")

    with conn.cursor() as cur:
        # Check workspace
        cur.execute(
            f'SELECT COUNT(*) FROM "{schema}"."workspaces" WHERE name = %s',
            (WORKSPACE_NAME,),
        )
        workspace_exists = cur.fetchone()[0] > 0

        # Check peers
        cur.execute(
            f'SELECT COUNT(*) FROM "{schema}"."peers" WHERE workspace_name = %s',
            (WORKSPACE_NAME,),
        )
        existing_peers = cur.fetchone()[0]

        # Check sessions
        cur.execute(
            f'SELECT COUNT(*) FROM "{schema}"."sessions" WHERE workspace_name = %s',
            (WORKSPACE_NAME,),
        )
        existing_sessions = cur.fetchone()[0]

        # Check collections (observer/observed pairs)
        cur.execute(
            f'SELECT COUNT(*) FROM "{schema}"."collections" WHERE workspace_name = %s',
            (WORKSPACE_NAME,),
        )
        existing_collections = cur.fetchone()[0]

        # Check messages
        cur.execute(
            f'SELECT COUNT(*) FROM "{schema}"."messages" WHERE workspace_name = %s',
            (WORKSPACE_NAME,),
        )
        existing_messages = cur.fetchone()[0]

        # Check message embeddings
        cur.execute(
            f'SELECT COUNT(*) FROM "{schema}"."message_embeddings" WHERE workspace_name = %s',
            (WORKSPACE_NAME,),
        )
        existing_embeddings = cur.fetchone()[0]

        # Check documents
        cur.execute(
            f'SELECT COUNT(*) FROM "{schema}"."documents" WHERE workspace_name = %s',
            (WORKSPACE_NAME,),
        )
        existing_documents = cur.fetchone()[0]

    print(f"  Workspace exists: {workspace_exists}")
    print(f"  Existing peers: {existing_peers}/{NUM_PEERS}")
    print(f"  Existing sessions: {existing_sessions}/{NUM_SESSIONS}")
    print(f"  Existing collections: {existing_collections}/{NUM_PEERS}")
    print(f"  Existing messages: {existing_messages:,}/{NUM_MESSAGES:,}")
    print(f"  Existing message embeddings: {existing_embeddings:,}/{NUM_MESSAGES:,}")
    print(f"  Existing documents: {existing_documents:,}/{NUM_DOCUMENTS:,}")

    return {
        "workspace_exists": workspace_exists,
        "existing_peers": existing_peers,
        "existing_sessions": existing_sessions,
        "existing_collections": existing_collections,
        "existing_messages": existing_messages,
        "existing_embeddings": existing_embeddings,
        "existing_documents": existing_documents,
    }


def create_entities(conn, schema, existing_data):
    """Create workspace, peers, sessions, and collections."""
    print("\nCreating workspace, peers, sessions, and collections...")

    with conn.cursor() as cur:
        # Create workspace if not exists
        if not existing_data["workspace_exists"]:
            cur.execute(
                f'INSERT INTO "{schema}"."workspaces" (id, name) VALUES (%s, %s)',
                (WORKSPACE_ID, WORKSPACE_NAME),
            )
            print(f"  ✓ Created workspace: {WORKSPACE_NAME}")
        else:
            print(f"  → Workspace already exists: {WORKSPACE_NAME}")

        # Create peers if needed
        if existing_data["existing_peers"] < NUM_PEERS:
            # Get existing peer names
            cur.execute(
                f'SELECT name FROM "{schema}"."peers" WHERE workspace_name = %s',
                (WORKSPACE_NAME,),
            )
            existing_peer_names = {row[0] for row in cur.fetchall()}

            for peer_id, peer_name in zip(PEER_IDS, PEER_NAMES, strict=False):
                if peer_name not in existing_peer_names:
                    cur.execute(
                        f'INSERT INTO "{schema}"."peers" (id, name, workspace_name) VALUES (%s, %s, %s)',
                        (peer_id, peer_name, WORKSPACE_NAME),
                    )
            new_peers = NUM_PEERS - existing_data["existing_peers"]
            print(
                f"  ✓ Created {new_peers} new peers ({existing_data['existing_peers']} already existed)"
            )
        else:
            print(f"  → All {NUM_PEERS} peers already exist")

        # Create sessions if needed
        if existing_data["existing_sessions"] < NUM_SESSIONS:
            # Get existing session names
            cur.execute(
                f'SELECT name FROM "{schema}"."sessions" WHERE workspace_name = %s',
                (WORKSPACE_NAME,),
            )
            existing_session_names = {row[0] for row in cur.fetchall()}

            for session_id, session_name in zip(
                SESSION_IDS, SESSION_NAMES, strict=False
            ):
                if session_name not in existing_session_names:
                    cur.execute(
                        f'INSERT INTO "{schema}"."sessions" (id, name, workspace_name) VALUES (%s, %s, %s)',
                        (session_id, session_name, WORKSPACE_NAME),
                    )
            new_sessions = NUM_SESSIONS - existing_data["existing_sessions"]
            print(
                f"  ✓ Created {new_sessions} new sessions ({existing_data['existing_sessions']} already existed)"
            )
        else:
            print(f"  → All {NUM_SESSIONS} sessions already exist")

        # Add peers to sessions via session_peers table
        # Check existing relationships
        cur.execute(
            f'SELECT COUNT(*) FROM "{schema}"."session_peers" WHERE workspace_name = %s',
            (WORKSPACE_NAME,),
        )
        existing_relationships = cur.fetchone()[0]
        expected_relationships = NUM_SESSIONS * NUM_PEERS

        if existing_relationships < expected_relationships:
            for session_name in SESSION_NAMES:
                for peer_name in PEER_NAMES:
                    # Check if relationship exists
                    cur.execute(
                        f'''SELECT COUNT(*) FROM "{schema}"."session_peers"
                        WHERE workspace_name = %s AND session_name = %s AND peer_name = %s''',
                        (WORKSPACE_NAME, session_name, peer_name),
                    )
                    if cur.fetchone()[0] == 0:
                        cur.execute(
                            f'''INSERT INTO "{schema}"."session_peers"
                            (workspace_name, session_name, peer_name)
                            VALUES (%s, %s, %s)''',
                            (WORKSPACE_NAME, session_name, peer_name),
                        )
            new_relationships = expected_relationships - existing_relationships
            print(
                f"  ✓ Added {new_relationships} new peer-session relationships ({existing_relationships} already existed)"
            )
        else:
            print("  → All peer-session relationships already exist")

        # Create collections (one per peer, self-observation) if needed
        if existing_data["existing_collections"] < NUM_PEERS:
            # Get existing collection observer/observed pairs
            cur.execute(
                f'SELECT observer, observed FROM "{schema}"."collections" WHERE workspace_name = %s',
                (WORKSPACE_NAME,),
            )
            existing_collections = {(row[0], row[1]) for row in cur.fetchall()}

            for collection_id, observer, observed in zip(
                COLLECTION_IDS, COLLECTION_OBSERVERS, COLLECTION_OBSERVED, strict=False
            ):
                if (observer, observed) not in existing_collections:
                    cur.execute(
                        f'INSERT INTO "{schema}"."collections" (id, observer, observed, workspace_name) VALUES (%s, %s, %s, %s)',
                        (collection_id, observer, observed, WORKSPACE_NAME),
                    )
            new_collections = NUM_PEERS - existing_data["existing_collections"]
            print(
                f"  ✓ Created {new_collections} new collections ({existing_data['existing_collections']} already existed)"
            )
        else:
            print(f"  → All {NUM_PEERS} collections already exist")

        conn.commit()


def generate_messages(conn, schema, existing_count):
    """Generate and insert messages up to NUM_MESSAGES total."""
    messages_to_create = NUM_MESSAGES - existing_count

    if messages_to_create <= 0:
        print(f"\n→ All {NUM_MESSAGES:,} messages already exist, skipping...")
        return []

    print(
        f"\nGenerating {messages_to_create:,} messages ({existing_count:,} already exist)..."
    )

    # Pre-generate all message data
    print("  Generating message data...")
    message_data = []
    for i in range(messages_to_create):
        public_id = generate_nanoid()
        peer_name = random.choice(PEER_NAMES)
        session_name = random.choice(SESSION_NAMES)
        content = fake.text(max_nb_chars=200)

        message_data.append(
            {
                "public_id": public_id,
                "peer_name": peer_name,
                "session_name": session_name,
                "content": content,
                "workspace_name": WORKSPACE_NAME,
            }
        )

    # Insert in batches
    print(f"  Inserting messages in batches of {BATCH_SIZE:,}...")
    total_batches = (messages_to_create + BATCH_SIZE - 1) // BATCH_SIZE

    with conn.cursor() as cur:
        for batch_num in range(total_batches):
            batch_start = batch_num * BATCH_SIZE
            batch_end = min(batch_start + BATCH_SIZE, messages_to_create)
            batch = message_data[batch_start:batch_end]

            # Use execute_batch for better performance
            args = [
                (
                    msg["public_id"],
                    msg["session_name"],
                    msg["content"],
                    msg["peer_name"],
                    msg["workspace_name"],
                )
                for msg in batch
            ]

            cur.executemany(
                f'''INSERT INTO "{schema}"."messages"
                (public_id, session_name, content, peer_name, workspace_name)
                VALUES (%s, %s, %s, %s, %s)''',
                args,
            )
            conn.commit()

            messages_created = existing_count + batch_end
            progress = (messages_created / NUM_MESSAGES) * 100
            print(
                f"    Batch {batch_num + 1}/{total_batches}: {messages_created:,}/{NUM_MESSAGES:,} ({progress:.1f}%)"
            )

    print(f"  ✓ Successfully created {messages_to_create:,} new messages!")
    return message_data


def generate_message_embeddings(conn, schema, message_data, existing_count):
    """Generate and insert message embeddings for messages that don't have them."""
    embeddings_to_create = NUM_MESSAGES - existing_count

    if embeddings_to_create <= 0:
        print(f"\n→ All {NUM_MESSAGES:,} message embeddings already exist, skipping...")
        return conn

    print(
        f"\nGenerating {embeddings_to_create:,} message embeddings ({existing_count:,} already exist)..."
    )

    # If we just created messages, use that data
    if message_data:
        print("  Using newly created message data...")
    else:
        # Fetch messages that don't have embeddings yet in chunks
        # Using NOT EXISTS is much faster than LEFT JOIN for large tables
        print("  Fetching messages without embeddings in chunks...")
        message_data = []
        chunk_size = 50_000  # Fetch in chunks to avoid timeout
        total_fetched = 0

        with conn.cursor() as cur:
            while total_fetched < embeddings_to_create:
                remaining = embeddings_to_create - total_fetched
                fetch_limit = min(chunk_size, remaining)

                cur.execute(
                    f'''SELECT m.public_id, m.content, m.peer_name, m.session_name, m.workspace_name
                    FROM "{schema}"."messages" m
                    WHERE m.workspace_name = %s
                    AND NOT EXISTS (
                        SELECT 1 FROM "{schema}"."message_embeddings" me
                        WHERE me.message_id = m.public_id
                    )
                    LIMIT %s''',
                    (WORKSPACE_NAME, fetch_limit),
                )
                rows = cur.fetchall()

                if not rows:
                    break  # No more messages to fetch

                message_data.extend(
                    [
                        {
                            "public_id": row[0],
                            "content": row[1],
                            "peer_name": row[2],
                            "session_name": row[3],
                            "workspace_name": row[4],
                        }
                        for row in rows
                    ]
                )

                total_fetched += len(rows)
                print(
                    f"    Fetched {total_fetched:,}/{embeddings_to_create:,} messages..."
                )

    if not message_data:
        print("  ✓ No messages found that need embeddings")
        return conn

    # Insert in batches (smaller batch size for embeddings)
    print(f"  Inserting embeddings in batches of {EMBEDDING_BATCH_SIZE:,}...")
    total_batches = (
        len(message_data) + EMBEDDING_BATCH_SIZE - 1
    ) // EMBEDDING_BATCH_SIZE

    for batch_num in range(total_batches):
        batch_start = batch_num * EMBEDDING_BATCH_SIZE
        batch_end = min(batch_start + EMBEDDING_BATCH_SIZE, len(message_data))
        batch = message_data[batch_start:batch_end]

        args = [
            (
                msg["content"],
                CONSTANT_EMBEDDING,
                msg["public_id"],
                msg["workspace_name"],
                msg["session_name"],
                msg["peer_name"],
            )
            for msg in batch
        ]

        # Retry logic for connection failures
        for attempt in range(MAX_RETRIES):
            try:
                with conn.cursor() as cur:
                    cur.executemany(
                        f'''INSERT INTO "{schema}"."message_embeddings"
                        (content, embedding, message_id, workspace_name, session_name, peer_name)
                        VALUES (%s, CAST(%s AS vector), %s, %s, %s, %s)''',
                        args,
                    )
                    conn.commit()
                break  # Success, exit retry loop
            except (psycopg.OperationalError, psycopg.InterfaceError):
                if attempt < MAX_RETRIES - 1:
                    print(
                        f"    ⚠ Connection error on batch {batch_num + 1}, retrying in {RETRY_DELAY}s... (attempt {attempt + 1}/{MAX_RETRIES})"
                    )
                    time.sleep(RETRY_DELAY)
                    # Reconnect with keepalive settings
                    try:
                        conn.close()
                    except:
                        pass
                    conn = psycopg.connect(
                        DB_CONNECTION_URI,
                        keepalives=1,
                        keepalives_idle=30,
                        keepalives_interval=10,
                        keepalives_count=5,
                    )
                else:
                    print(f"    ✗ Failed after {MAX_RETRIES} attempts")
                    raise

        embeddings_created = existing_count + batch_end
        progress = (embeddings_created / NUM_MESSAGES) * 100
        print(
            f"    Batch {batch_num + 1}/{total_batches}: {embeddings_created:,}/{NUM_MESSAGES:,} ({progress:.1f}%)"
        )

    print(f"  ✓ Successfully created {len(message_data):,} new message embeddings!")
    return conn


def generate_documents(conn, schema, existing_count):
    """Generate and insert documents up to NUM_DOCUMENTS total."""
    documents_to_create = NUM_DOCUMENTS - existing_count

    if documents_to_create <= 0:
        print(f"\n→ All {NUM_DOCUMENTS:,} documents already exist, skipping...")
        return

    print(
        f"\nGenerating {documents_to_create:,} documents ({existing_count:,} already exist)..."
    )

    # Pre-generate all document data
    print("  Generating document data...")
    document_data = []
    for i in range(documents_to_create):
        doc_id = generate_nanoid()
        # Randomly assign to a collection/peer (observer/observed pairs)
        idx = random.randint(0, NUM_PEERS - 1)
        observer = COLLECTION_OBSERVERS[idx]
        observed = COLLECTION_OBSERVED[idx]
        # Randomly assign to a session
        session_name = random.choice(SESSION_NAMES)
        content = fake.text(max_nb_chars=300)

        document_data.append(
            {
                "id": doc_id,
                "content": content,
                "observer": observer,
                "observed": observed,
                "session_name": session_name,
                "workspace_name": WORKSPACE_NAME,
            }
        )

    # Insert in batches
    print(f"  Inserting documents in batches of {BATCH_SIZE:,}...")
    total_batches = (documents_to_create + BATCH_SIZE - 1) // BATCH_SIZE

    with conn.cursor() as cur:
        for batch_num in range(total_batches):
            batch_start = batch_num * BATCH_SIZE
            batch_end = min(batch_start + BATCH_SIZE, documents_to_create)
            batch = document_data[batch_start:batch_end]

            args = [
                (
                    doc["id"],
                    doc["content"],
                    CONSTANT_EMBEDDING,
                    doc["observer"],
                    doc["observed"],
                    doc["workspace_name"],
                    doc["session_name"],
                )
                for doc in batch
            ]

            cur.executemany(
                f'''INSERT INTO "{schema}"."documents"
                (id, content, embedding, observer, observed, workspace_name, session_name)
                VALUES (%s, %s, CAST(%s AS vector), %s, %s, %s, %s)''',
                args,
            )
            conn.commit()

            documents_created = existing_count + batch_end
            progress = (documents_created / NUM_DOCUMENTS) * 100
            print(
                f"    Batch {batch_num + 1}/{total_batches}: {documents_created:,}/{NUM_DOCUMENTS:,} ({progress:.1f}%)"
            )

    print(f"  ✓ Successfully created {documents_to_create:,} new documents!")


def main():
    """Main execution function."""
    print("=" * 60)
    print("HONCHO LOAD TEST - Raw psycopg3 Implementation (Resumable)")
    print("=" * 60)
    print("Configuration:")
    print(f"  Database URI: {DB_CONNECTION_URI}")
    print(f"  Schema: {DB_SCHEMA}")
    print(f"  Target messages: {NUM_MESSAGES:,}")
    print(f"  Target message embeddings: {NUM_MESSAGES:,}")
    print(f"  Target documents: {NUM_DOCUMENTS:,}")
    print(f"  Batch size (messages/docs): {BATCH_SIZE:,}")
    print(f"  Batch size (embeddings): {EMBEDDING_BATCH_SIZE:,}")
    print("=" * 60)

    # Connect to database with keepalive settings
    print("\nConnecting to database...")
    conn = psycopg.connect(
        DB_CONNECTION_URI,
        keepalives=1,
        keepalives_idle=30,
        keepalives_interval=10,
        keepalives_count=5,
    )

    try:
        # Check what already exists
        existing_data = check_existing_data(conn, DB_SCHEMA)

        # Create all entities
        create_entities(conn, DB_SCHEMA, existing_data)

        # Generate messages
        message_data = generate_messages(
            conn, DB_SCHEMA, existing_data["existing_messages"]
        )

        # Generate message embeddings (returns potentially reconnected connection)
        # conn = generate_message_embeddings(
        #     conn, DB_SCHEMA, message_data, existing_data["existing_embeddings"]
        # )

        # Generate documents
        generate_documents(conn, DB_SCHEMA, existing_data["existing_documents"])

        print("\n" + "=" * 60)
        print("✓ LOAD TEST COMPLETE!")
        print("=" * 60)
        print("Summary:")
        print(f"  Workspace: {WORKSPACE_NAME}")
        print(f"  Peers: {NUM_PEERS}")
        print(f"  Sessions: {NUM_SESSIONS}")
        print(f"  Collections: {NUM_PEERS}")
        print(f"  Messages: {NUM_MESSAGES:,}")
        print(f"  Message embeddings: {NUM_MESSAGES:,}")
        print(f"  Documents: {NUM_DOCUMENTS:,}")
        print("=" * 60)

    except Exception as e:
        print(f"\n✗ Error occurred: {e}")
        print(
            "\nThe script is resumable - simply run it again to continue from where it left off."
        )
        raise

    finally:
        conn.close()
        print("\nDatabase connection closed.")


if __name__ == "__main__":
    main()