import psycopg2
import os
from dotenv import load_dotenv

load_dotenv()

def setup_mcq_table(drop_existing_table: bool = False):
    """
    Sets up the mcq_questions table in the database.
    """
    try:
        conn = psycopg2.connect(
            host=os.getenv("POSTGRES_HOST", "localhost"),
            port=os.getenv("POSTGRES_PORT", "5432"),
            user=os.getenv("POSTGRES_USER"),
            password=os.getenv("POSTGRES_PASSWORD"),
            dbname=os.getenv("POSTGRES_DB")
        )
        conn.autocommit = True
        
        with conn.cursor() as cur:
            if drop_existing_table:
                print("Dropping existing mcq_questions table...")
                cur.execute("DROP TABLE IF EXISTS mcq_questions CASCADE;")
                print("Table dropped.")

            print("Creating mcq_questions table...")
            # --- THIS IS THE UPDATED TABLE SCHEMA ---
            cur.execute("""
                CREATE TABLE IF NOT EXISTS mcq_questions (
                    id SERIAL PRIMARY KEY,
                    grade INTEGER NOT NULL,
                    is_arabic BOOLEAN NOT NULL, -- <-- ADDED THIS LINE
                    subject TEXT NOT NULL,
                    unit TEXT NOT NULL,
                    concept TEXT NOT NULL,
                    question_text TEXT NOT NULL,
                    correct_answer TEXT NOT NULL,
                    wrong_answer_1 TEXT NOT NULL,
                    wrong_answer_2 TEXT NOT NULL,
                    wrong_answer_3 TEXT NOT NULL,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                );
            """)
            
            print("Creating indexes on mcq_questions table...")
            # --- THIS IS THE UPDATED INDEX ---
            cur.execute("""
                CREATE INDEX IF NOT EXISTS idx_mcq_topic 
                ON mcq_questions(grade, is_arabic, subject, unit, concept); -- <-- ADDED is_arabic
            """)

            print("MCQ table setup complete.")
                
    except Exception as e:
        print(f"An error occurred during MCQ table setup: {e}")
    finally:
        if 'conn' in locals() and conn:
            conn.close()
            print("Database connection closed.")

if __name__ == "__main__":
    # To apply the changes, it's best to drop and recreate the table.
    # Be careful if you have existing data you want to keep!
    print("Creating MCQ table...")
    setup_mcq_table(drop_existing_table=False)