import asyncio
import os
import sys

# Add project root to path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from motor.motor_asyncio import AsyncIOMotorClient
from repos.dao import DAO
from models.Album import Album
from models.Generation import Generation, GenerationStatus
from models.enums import AspectRatios, Quality
from config import settings

# Mock config
# Use the same host as aiws.py but different DB
MONGO_HOST = settings.MONGO_HOST
DB_NAME = "bot_db_test_albums"

async def test_albums():
    print(f"🚀 Starting Album Manual Verification using {MONGO_HOST}...")
    
    # Needs to run inside a loop from main
    client = AsyncIOMotorClient(MONGO_HOST)
    dao = DAO(client, db_name=DB_NAME)
    
    try:
        # 1. Clean up
        await client[DB_NAME]["albums"].drop()
        await client[DB_NAME]["generations"].drop()
        print("✅ Cleaned up test database")

        # 2. Create Album
        album = Album(name="Test Album", description="A test album")
        print("Creating album...")
        album_id = await dao.albums.create_album(album)
        print(f"✅ Created Album: {album_id}")
        
        # 3. Create Generations
        gen1 = Generation(prompt="Gen 1", aspect_ratio=AspectRatios.NINESIXTEEN, quality=Quality.ONEK)
        gen2 = Generation(prompt="Gen 2", aspect_ratio=AspectRatios.NINESIXTEEN, quality=Quality.ONEK)
        
        print("Creating generations...")
        gen1_id = await dao.generations.create_generation(gen1)
        gen2_id = await dao.generations.create_generation(gen2)
        print(f"✅ Created Generations: {gen1_id}, {gen2_id}")
        
        # 4. Add generations to album
        print("Adding generations to album...")
        await dao.albums.add_generation(album_id, gen1_id)
        await dao.albums.add_generation(album_id, gen2_id)
        print("✅ Added generations to album")
        
        # 5. Fetch album and check generation_ids
        album_fetched = await dao.albums.get_album(album_id)
        assert album_fetched is not None
        assert len(album_fetched.generation_ids) == 2
        assert gen1_id in album_fetched.generation_ids
        assert gen2_id in album_fetched.generation_ids
        print("✅ Verified generations in album")
        
        # 6. Fetch generations by IDs via GenerationRepo
        generations = await dao.generations.get_generations_by_ids([gen1_id, gen2_id])
        assert len(generations) == 2
        
        # Ensure ID type match (str vs ObjectId handling in repo)
        gen_ids_fetched = [g.id for g in generations]
        assert gen1_id in gen_ids_fetched
        assert gen2_id in gen_ids_fetched
        print("✅ Verified fetching generations by IDs")
        
        # 7. Remove generation
        print("Removing generation...")
        await dao.albums.remove_generation(album_id, gen1_id)
        album_fetched = await dao.albums.get_album(album_id)
        assert len(album_fetched.generation_ids) == 1
        assert album_fetched.generation_ids[0] == gen2_id
        print("✅ Verified removing generation from album")
        
        print("🎉 Album Verification SUCCESS")
        
    finally:
        # Cleanup client
        client.close()

if __name__ == "__main__":
    try:
        asyncio.run(test_albums())
    except Exception as e:
        print(f"Error: {e}")