import os
import psycopg2
import pandas as pd
import json
from dotenv import load_dotenv

# Import the pgvector adapter for psycopg2
from pgvector.psycopg2 import register_vector

load_dotenv()

def get_db_connection():
    conn = psycopg2.connect(
        dbname=os.getenv("POSTGRES_DB", "embeddings_db"),
        user=os.getenv("POSTGRES_USER", "db_admin"),
        password=os.getenv("POSTGRES_PASSWORD"),
        host=os.getenv("POSTGRES_HOST", "localhost"),
        port=os.getenv("POSTGRES_PORT", 5432)
    )
    # Register the vector type with the connection
    register_vector(conn)
    return conn

def create_schema_and_table(conn, drop_existing_table: bool):
    create_extension = "CREATE EXTENSION IF NOT EXISTS vector;"
    create_table = """
    CREATE TABLE IF NOT EXISTS educational_chunks (
      id SERIAL PRIMARY KEY,
      grade TEXT NOT NULL,
      subject TEXT,
      unit TEXT,
      concept TEXT,
      lesson TEXT,
      from_page INT,
      to_page INT,
      chunk_index INT,
      chunk_text TEXT NOT NULL,
      is_arabic BOOLEAN NOT NULL,
      embedding VECTOR(1536) NOT NULL
    );
    """
    create_indexes = [
        "CREATE INDEX IF NOT EXISTS idx_embedding ON educational_chunks USING hnsw (embedding vector_cosine_ops);",
        "CREATE INDEX IF NOT EXISTS idx_grade ON educational_chunks (grade);",
        "CREATE INDEX IF NOT EXISTS idx_is_arabic ON educational_chunks (is_arabic);",
        "CREATE INDEX IF NOT EXISTS idx_subject ON educational_chunks (subject);",
        "CREATE INDEX IF NOT EXISTS idx_grade_is_arabic ON educational_chunks (grade, is_arabic);"
    ]

    cur = conn.cursor()
    cur.execute(create_extension)
    print("CREATE EXTENSION vector operation fine.")
    
    if drop_existing_table:
        drop_table = "DROP TABLE IF EXISTS educational_chunks;"
        cur.execute(drop_table)
        print("DROP TABLE educational_chunks operation fine.")
        
    cur.execute(create_table)
    print("CREATE TABLE educational_chunks operation fine.")
    for idx_query in create_indexes:
        cur.execute(idx_query)
        print(f"CREATE INDEX operation fine for: {idx_query}")
    conn.commit()
    cur.close()


def insert_chunks_from_csv(csv_file: str):
    df = pd.read_csv(csv_file)

    required_cols = [
        "Grade", "Subject", "Unit", "Concept", "Lesson",
        "From page", "To page", "Chunk index", "Chunk text",
        "Is Arabic", "Embedding"
    ]

    for col in required_cols:
        if col not in df.columns:
            raise ValueError(f"Missing required column in CSV: {col}")

    conn = get_db_connection()
    cur = conn.cursor()

    insert_query = """
        INSERT INTO educational_chunks
        (grade, subject, unit, concept, lesson,
         from_page, to_page, chunk_index, chunk_text,
         is_arabic, embedding)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
    """

    batch_size = 50
    buffer = []

    for idx, row in df.iterrows():
        try:
            embedding = json.loads(row["Embedding"])
            buffer.append((
                str(row["Grade"]),
                row["Subject"],
                row.get("Unit"),
                row.get("Concept"),
                row.get("Lesson"),
                int(row["From page"]) if not pd.isna(row["From page"]) else None,
                int(row["To page"]) if not pd.isna(row["To page"]) else None,
                int(row["Chunk index"]),
                row["Chunk text"],
                bool(row["Is Arabic"]),
                embedding
            ))
        except Exception as e:
            print(f"Skipping row {idx} due to error: {e}")
            continue

        if len(buffer) >= batch_size:
            cur.executemany(insert_query, buffer)
            conn.commit()
            print(f"Inserted {len(buffer)} rows. Operation fine.")
            buffer = []

    if buffer:
        cur.executemany(insert_query, buffer)
        conn.commit()
        print(f"Inserted final {len(buffer)} rows. Operation fine.")

    cur.close()
    conn.close()
    print("All data inserted successfully.")

def setup_embeddings_database(drop_existing_tables: bool = False):
    """
    Sets up the educational chunks table and populates it with embeddings from CSV files.
    Args:
        drop_existing_tables: If True, drops the existing table before creating it.
    """
    try:
        conn = get_db_connection()
        create_schema_and_table(conn, drop_existing_tables)
        
        csv_files = ["prime4_ar_embeddings.csv", "Prime5_en_chunked_with_embeddings.csv", "prime6_ar_embeddings.csv", "Prime6_en_chunked_with_embeddings.csv"]
        for file in csv_files:
            if os.path.exists(file):
                print(f"Inserting data from {file}...")
                insert_chunks_from_csv(file)
            else:
                print(f"File not found: {file}")

    except psycopg2.OperationalError as e:
        print(f"Database connection failed: {e}")
    except Exception as e:
        print(f"An error occurred: {e}")
    finally:
        if 'conn' in locals() and conn:
            conn.close()
            print("Database connection closed.")

if __name__ == "__main__":
    # To run with a clean slate, pass True
    # setup_embeddings_database(drop_existing_tables=True)
    
    # To run without dropping the table (default)
    setup_embeddings_database()