import asyncio
from sqlalchemy import text
from database.database import db_manager
from logger.logger import get_logger
from database.database_models import Base
logger = get_logger(__name__)


async def clean_all_data():
    """
    Clean all data from tables without dropping schema
    """
    await db_manager.initialize()

    async with db_manager.engine.begin() as conn:
        logger.info("Cleaning all tables...")
        
        # Truncate tables in correct order to handle foreign key constraints
        # Using CASCADE will automatically handle dependent tables
        await conn.execute(text("TRUNCATE TABLE messages CASCADE;"))
        await conn.execute(text("TRUNCATE TABLE conversations CASCADE;"))
        await conn.execute(text("TRUNCATE TABLE documents CASCADE;"))
        await conn.execute(text("TRUNCATE TABLE users CASCADE;"))
        
        # Reset sequences
        await conn.execute(text("""
            SELECT setval(pg_get_serial_sequence('users', 'id'), 
                COALESCE((SELECT MAX(id) FROM users), 1), false);
        """))
        await conn.execute(text("""
            SELECT setval(pg_get_serial_sequence('documents', 'id'), 
                COALESCE((SELECT MAX(id) FROM documents), 1), false);
        """))
        
        logger.info("All tables cleaned and sequences reset!")

async def drop_all_tables():
    """
    Drop all database tables.
    """
    await db_manager.initialize()

    async with db_manager.engine.begin() as conn:
        logger.info("Dropping all tables...")
        await conn.run_sync(Base.metadata.drop_all)
        logger.info("All tables dropped successfully!")

async def create_all_tables():
    await db_manager.initialize()

    async with db_manager.engine.begin() as conn:
        logger.info("Creating tables...")
        await conn.run_sync(Base.metadata.create_all)
        logger.info("All tables created successfully!")

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