#!/usr/bin/env python3
"""Initialize the database schema"""

from database import init_db, engine
from sqlalchemy import text
import sys


def check_database_connection():
    """Check if we can connect to the database"""
    try:
        with engine.connect() as conn:
            conn.execute(text("SELECT 1"))
        print("✓ Database connection successful")
        return True
    except Exception as e:
        print(f"✗ Database connection failed: {e}")
        return False


def main():
    print("Initializing NotebookLlaMa database...")

    if not check_database_connection():
        print("\nPlease ensure:")
        print("1. PostgreSQL is running (docker compose up -d)")
        print("2. Environment variables are set correctly in .env")
        sys.exit(1)

    try:
        init_db()
        print("✓ Database tables created successfully")
        print("\nTables created:")
        print("  - users")
        print("  - documents")
        print("  - notebooks")
        print("  - chat_sessions")
        print("  - chat_messages")
        print("  - document_shares")
        print("\nDatabase initialization complete!")
    except Exception as e:
        print(f"✗ Error creating tables: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()