import os
import psycopg2
import pandas as pd
from psycopg2.extras import execute_values
from dotenv import load_dotenv

load_dotenv()

def get_db_connection():
    return 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)
    )

def insert_lessons_from_csv(file_path, conn, grade, subject):
    df = pd.read_csv(file_path)

    df.rename(columns={
        "الوحدة": "Unit",
        "المفهوم": "Concept",
        "الدرس": "Lesson",
        "من صفحة": "From page",
        "إلى صفحة": "To page",
        "النص": "Lesson text"
    }, inplace=True)

    required_columns = ["Unit", "Concept", "Lesson", "From page", "To page", "Lesson text"]
    missing = [col for col in required_columns if col not in df.columns]
    if missing:
        print(f"⚠️ Missing columns in {file_path}: {', '.join(missing)}")
        return

    rows = []
    for _, row in df.iterrows():
        rows.append((
            grade,
            subject,
            row["Unit"],
            row["Concept"],
            row["Lesson"],
            row["From page"],
            row["To page"],
            row["Lesson text"]
        ))

    query = """
        INSERT INTO lessons (grade, subject, unit, concept, lesson, start_page, end_page, lesson_text)
        VALUES %s
    """

    with conn.cursor() as cur:
        execute_values(cur, query, rows)
    conn.commit()
    print(f" Inserted {len(rows)} rows from {os.path.basename(file_path)}")


def main():
    folder = input("Enter the path to the folder containing CSV files: ").strip()
    if not os.path.exists(folder):
        print("Folder not found.")
        return

    files = [f for f in os.listdir(folder) if f.endswith(".csv")]
    if not files:
        print(" No CSV files found.")
        return

    print("Available files:")
    for i, f in enumerate(files, 1):
        print(f"{i}. {f}")

    selected = input("Enter the numbers of the files you want to import (e.g., 1 3 4): ").split()
    selected_files = [files[int(i) - 1] for i in selected]

    grade = input("Enter grade manually (e.g., Grade 5): ").strip()
    subject = input("Enter subject manually (default: Science): ").strip() or "Science"

    conn = get_db_connection()
    try:
        for f in selected_files:
            file_path = os.path.join(folder, f)
            insert_lessons_from_csv(file_path, conn, grade, subject)

        with conn.cursor() as cur:
            cur.execute("SELECT COUNT(*) FROM lessons;")
            total = cur.fetchone()[0]
            print(f" Total rows in lessons table: {total}")
    finally:
        conn.close()
        print(" Connection closed.")

if __name__ == "__main__":
    main()
