import psycopg2
import os
import sys

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
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 data (Preserving OLD students and adding NEW ones up to Grade 12)
INSERT INTO students (student_id, student_name, grade, language, nationality) VALUES
 -- THE ORIGINAL STUDENTS (Do not change IDs)
 ('student_001', 'Ahmed Ali', 4, TRUE, 'EGYPTIAN'),
 ('student_002', 'Khalid Al-Rashid', 4, TRUE, 'SAUDI'),
 ('student_003', 'Fatima Hassan', 6, TRUE, 'EGYPTIAN'),
 ('student_004', 'Nora Al-Zahrani', 6, TRUE, 'SAUDI'),
 ('student_005', 'Mona Adel', 5, FALSE, 'EGYPTIAN'),
 ('student_006', 'Sara Al-Mutairi', 5, FALSE, 'SAUDI'),
 ('student_007', 'Omar Youssef', 6, FALSE, 'EGYPTIAN'),
 ('student_008', 'Abdullah Al-Harbi', 6, FALSE, 'SAUDI'),

 -- NEW STUDENTS: PREP & SECONDARY (Grades 7-12)
 -- Grade 7 (Prep 1)
 ('student_009', 'Mostafa Mahmoud', 7, TRUE, 'EGYPTIAN'),
 ('student_010', 'Yasmine Fahmy', 7, FALSE, 'EGYPTIAN'),
 
 -- Grade 8 (Prep 2)
 ('student_011', 'Hassan Kareem', 8, TRUE, 'EGYPTIAN'),
 ('student_012', 'Laila Ibrahim', 8, FALSE, 'EGYPTIAN'),

 -- Grade 9 (Prep 3)
 ('student_013', 'Zaid Al-Qahtani', 9, TRUE, 'SAUDI'),
 ('student_014', 'Mariam Soliman', 9, FALSE, 'EGYPTIAN'),

 -- Grade 10 (Secondary 1)
 ('student_015', 'Tarek Hegazy', 10, TRUE, 'EGYPTIAN'),
 ('student_016', 'Salma Ezzat', 10, FALSE, 'EGYPTIAN'),

 -- Grade 11 (Secondary 2)
 ('student_017', 'Faisal Al-Otaibi', 11, TRUE, 'SAUDI'),
 ('student_018', 'Habiba Ahmed', 11, FALSE, 'EGYPTIAN'),

 -- Grade 12 (Secondary 3)
 ('student_019', 'Youssef Mansour', 12, TRUE, 'EGYPTIAN'),
 ('student_020', 'Jana Wael', 12, FALSE, 'EGYPTIAN')

ON CONFLICT (student_id) DO NOTHING;
"""

drop_all_tables_sql = """
DO $$
DECLARE
    rec RECORD;
BEGIN
    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):
    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("Setting up schema and inserting hardcoded data...")
            cur.execute(schema_sql)
            
            print("Database setup complete.")
            cur.execute("SELECT student_id, student_name, grade FROM students ORDER BY grade ASC;")
            for row in cur.fetchall():
                print(row)
                
    except Exception as e:
        print(f"An error occurred: {e}")
    finally:
        if 'conn' in locals() and conn:
            conn.close()

if __name__ == "__main__":
    _drop_env = os.getenv("DROP_TEST_SCHEMA", "False")
    drop_existing_tables = str(_drop_env).strip().lower() in ("1", "true", "yes", "y", "t")
    setup_database(drop_existing_tables=drop_existing_tables)