from typing import List, Optional
from datetime import datetime
from bson import ObjectId
from motor.motor_asyncio import AsyncIOMotorClient
from models.Environment import Environment


class EnvironmentRepo:
    def __init__(self, client: AsyncIOMotorClient, db_name="bot_db"):
        self.collection = client[db_name]["environments"]

    async def create_env(self, env: Environment) -> Environment:
        env_dict = env.model_dump(exclude={"id"})
        res = await self.collection.insert_one(env_dict)
        env.id = str(res.inserted_id)
        return env

    async def get_env(self, env_id: str) -> Optional[Environment]:
        res = await self.collection.find_one({"_id": ObjectId(env_id)})
        if not res:
            return None
        res["id"] = str(res.pop("_id"))
        return Environment(**res)

    async def get_character_envs(self, character_id: str) -> List[Environment]:
        cursor = self.collection.find({"character_id": character_id})
        envs = []
        async for doc in cursor:
            doc["id"] = str(doc.pop("_id"))
            envs.append(Environment(**doc))
        return envs

    async def update_env(self, env_id: str, update_data: dict) -> bool:
        update_data["updated_at"] = datetime.utcnow()
        res = await self.collection.update_one(
            {"_id": ObjectId(env_id)},
            {"$set": update_data}
        )
        return res.modified_count > 0

    async def delete_env(self, env_id: str) -> bool:
        res = await self.collection.delete_one({"_id": ObjectId(env_id)})
        return res.deleted_count > 0

    async def add_asset(self, env_id: str, asset_id: str) -> bool:
        res = await self.collection.update_one(
            {"_id": ObjectId(env_id)},
            {
                "$addToSet": {"asset_ids": asset_id},
                "$set": {"updated_at": datetime.utcnow()}
            }
        )
        return res.modified_count > 0

    async def add_assets(self, env_id: str, asset_ids: List[str]) -> bool:
        res = await self.collection.update_one(
            {"_id": ObjectId(env_id)},
            {
                "$addToSet": {"asset_ids": {"$each": asset_ids}},
                "$set": {"updated_at": datetime.utcnow()}
            }
        )
        return res.modified_count > 0

    async def remove_asset(self, env_id: str, asset_id: str) -> bool:
        res = await self.collection.update_one(
            {"_id": ObjectId(env_id)},
            {
                "$pull": {"asset_ids": asset_id},
                "$set": {"updated_at": datetime.utcnow()}
            }
        )
        return res.modified_count > 0