import psycopg2
import os

schema_sql = """
-- Create students table
CREATE TABLE IF NOT EXISTS students (
    id SERIAL PRIMARY KEY,
    student_id VARCHAR(50) UNIQUE NOT NULL,
    student_name VARCHAR(100),
    grade INTEGER NOT NULL,
    language BOOLEAN,
    nationality VARCHAR(20) NOT NULL DEFAULT 'EGYPTIAN'
);

-- Create chat_history table
CREATE TABLE IF NOT EXISTS chat_history (
    id SERIAL PRIMARY KEY,
    student_id VARCHAR(50) NOT NULL,
    role VARCHAR(20) NOT NULL CHECK (role IN ('user', 'assistant', 'system')),
    content TEXT NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    FOREIGN KEY (student_id) REFERENCES students(student_id) ON DELETE CASCADE
);

-- Create indexes for better performance
CREATE INDEX IF NOT EXISTS idx_chat_history_student_id ON chat_history(student_id);
CREATE INDEX IF NOT EXISTS idx_chat_history_created_at ON chat_history(created_at);
CREATE INDEX IF NOT EXISTS idx_students_nationality ON students(nationality);
CREATE INDEX IF NOT EXISTS idx_students_grade ON students(grade);
CREATE INDEX IF NOT EXISTS idx_students_grade_language ON students(grade, language);

-- Insert dummy data for testing
INSERT INTO students (student_id, student_name, grade, language, nationality) VALUES
 -- Arabic Grade 4: One Egyptian, One Saudi
 ('student_001', 'Ahmed Ali', 4, TRUE, 'EGYPTIAN'),
 ('student_002', 'Khalid Al-Rashid', 4, TRUE, 'SAUDI'),

-- Arabic Grade 6: One Egyptian, One Saudi
 ('student_003', 'Fatima Hassan', 6, TRUE, 'EGYPTIAN'),
 ('student_004', 'Nora Al-Zahrani', 6, TRUE, 'SAUDI'),

-- English Grade 5: One Egyptian, One Saudi
 ('student_005', 'Mona Adel', 5, FALSE, 'EGYPTIAN'),
 ('student_006', 'Sara Al-Mutairi', 5, FALSE, 'SAUDI'),

-- English Grade 6: One Egyptian, One Saudi
 ('student_007', 'Omar Youssef', 6, FALSE, 'EGYPTIAN'),
 ('student_008', 'Abdullah Al-Harbi', 6, FALSE, 'SAUDI')

ON CONFLICT (student_id) DO NOTHING;
"""

drop_all_tables_sql = """
DO $$
DECLARE
    rec RECORD;
BEGIN
    -- drop all tables in public schema
    FOR rec IN (SELECT tablename FROM pg_tables WHERE schemaname = 'public') LOOP
        EXECUTE 'DROP TABLE IF EXISTS "' || rec.tablename || '" CASCADE';
    END LOOP;
END $$;
"""

def setup_database(drop_existing_tables: bool = False):
    """
    Sets up the database schema and tables.
    Args:
        drop_existing_tables: If True, drops all existing tables before creating them.
    """
    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_tables:
                print("Dropping all existing tables...")
                cur.execute(drop_all_tables_sql)
                print("All tables dropped.")
            
            print("Setting up schema and inserting data...")
            cur.execute(schema_sql)
            
            print("Database setup complete. Verifying data...")
            
            # Verifications: Select from students and chat_history tables
            print("\nStudents table rows:")
            cur.execute("SELECT * FROM students ORDER BY id;")
            students = cur.fetchall()
            for row in students:
                print(row)
            
            print("\nChat_history table rows:")
            cur.execute("SELECT * FROM chat_history ORDER BY id;")
            chat_history = cur.fetchall()
            for row in chat_history:
                print(row)
                
    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_database(drop_existing_tables=True)
    
    # To run without dropping tables (default)
    #setup_database()