"""
Cleanup script to handle duplicates in fmea_record table or perform a full reset.
"""
import asyncio
import os
import sys
import argparse

# Ensure the project root is on the path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from sqlalchemy import text
from src.database.core import engine

async def cleanup(mode="deduplicate"):
    async with engine.begin() as conn:
        if mode == "reset":
            print("Full Reset: Dropping and recreating table 'fmea_record'...")
            await conn.execute(text("DROP TABLE IF EXISTS fmea_record CASCADE;"))
            # Note: recreations usually happens via sync script or alembic, 
            # but we can trigger it here if Base is imported.
            from src.database.core import Base
            from src.fmea.model import FMEARecord # Ensure model is loaded
            await conn.run_sync(Base.metadata.create_all)
            print("✓ Table reset successfully.")
            
        elif mode == "deduplicate":
            print("Cleaning mode: Removing duplicates from 'fmea_record'...")
            # SQL to remove duplicates keeping the latest one
            cleanup_sql = """
            DELETE FROM fmea_record
            WHERE id IN (
                SELECT id
                FROM (
                    SELECT id,
                    ROW_NUMBER() OVER (
                        PARTITION BY system, equipment, equipment_tag, failure_mode, failure_cause_mechanism
                        ORDER BY created_at DESC, updated_at DESC
                    ) as row_num
                    FROM fmea_record
                ) t
                WHERE t.row_num > 1
            );
            """
            result = await conn.execute(text(cleanup_sql))
            print(f"✓ Removed duplicates. {result.rowcount} rows deleted.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Clean up FMEA records.")
    parser.add_argument("--mode", choices=["deduplicate", "reset"], default="deduplicate", help="Cleanup mode")
    args = parser.parse_args()
    
    asyncio.run(cleanup(args.mode))