import asyncio
from sqlalchemy.ext.asyncio import AsyncSession
from app.database import get_db, init_db
from app.models.db import StaticData
from datetime import datetime
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

async def seed_asteroid_belts():
    await init_db() # Ensure database is initialized

    async for session in get_db(): # Use async for to get the session
        logger.info("Seeding asteroid and Kuiper belt static data...")

        belts_data = [
            {
                "category": "asteroid_belt",
                "name": "Main Asteroid Belt",
                "name_zh": "主小行星带",
                "data": {
                    "innerRadiusAU": 2.2,
                    "outerRadiusAU": 3.2,
                    "count": 1500,
                    "color": "#665544",
                    "size": 0.1,
                    "opacity": 0.4,
                    "heightScale": 0.05,
                    "rotationSpeed": 0.02
                }
            },
            {
                "category": "kuiper_belt",
                "name": "Kuiper Belt",
                "name_zh": "柯伊伯带",
                "data": {
                    "innerRadiusAU": 30,
                    "outerRadiusAU": 50,
                    "count": 2500,
                    "color": "#AABBDD",
                    "size": 0.2,
                    "opacity": 0.3,
                    "heightScale": 0.1,
                    "rotationSpeed": 0.005
                }
            }
        ]

        for belt_item in belts_data:
            # Check if an item with the same category and name already exists
            existing_item = await session.execute(
                StaticData.__table__.select().where(
                    StaticData.category == belt_item["category"],
                    StaticData.name == belt_item["name"]
                )
            )
            if existing_item.scalar_one_or_none():
                logger.info(f"Static data for {belt_item['name']} already exists. Updating...")
                stmt = StaticData.__table__.update().where(
                    StaticData.category == belt_item["category"],
                    StaticData.name == belt_item["name"]
                ).values(
                    name_zh=belt_item["name_zh"],
                    data=belt_item["data"],
                    updated_at=datetime.utcnow()
                )
                await session.execute(stmt)
            else:
                logger.info(f"Adding static data for {belt_item['name']}...")
                static_data_entry = StaticData(
                    category=belt_item["category"],
                    name=belt_item["name"],
                    name_zh=belt_item["name_zh"],
                    data=belt_item["data"]
                )
                session.add(static_data_entry)
        
        await session.commit()
        logger.info("Asteroid and Kuiper belt static data seeding complete.")

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