#!/usr/bin/env python3
"""Initialize PostgreSQL database and schema for LightRAG."""

import os
import sys
import asyncio
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker

# Add parent directory to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from app.config import settings
from app.models import Base
from app.db import init_db


async def create_database():
    """Create the database if it doesn't exist."""
    # Connect to default PostgreSQL database
    default_url = settings.DATABASE_URL.rsplit('/', 1)[0] + '/postgres'
    engine = create_engine(default_url, echo=True)

    with engine.connect() as conn:
        conn.execution_options(isolation_level="AUTOCOMMIT")
        db_name = settings.DATABASE_URL.split('/')[-1]

        # Check if database exists
        result = conn.execute(
            text("SELECT 1 FROM pg_database WHERE datname = :db_name"),
            {"db_name": db_name}
        )

        if not result.fetchone():
            print(f"Creating database: {db_name}")
            conn.execute(text(f"CREATE DATABASE {db_name}"))
        else:
            print(f"Database {db_name} already exists")

        conn.commit()

    engine.dispose()


async def init_schema():
    """Initialize database schema."""
    await init_db()
    print("Database schema initialized")


async def main():
    """Main initialization."""
    print(f"Initializing database: {settings.DATABASE_URL}")

    # Create database
    await create_database()

    # Initialize schema
    await init_schema()

    print("Database initialization complete!")


if __name__ == "__main__":
    asyncio.run(main())