from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from datetime import datetime, timedelta
from typing import Optional
from tg_bot.infrastructure.database.models import UserModel


class UserService:

    def __init__(self, session: AsyncSession):
        self.session = session

    async def get_user_by_telegram_id(self, telegram_id: int) -> Optional[UserModel]:
        result = await self.session.execute(
            select(UserModel).filter_by(telegram_id=str(telegram_id))
        )
        return result.scalar_one_or_none()

    async def get_or_create_user(
        self,
        telegram_id: int,
        username: str = "",
        first_name: str = "",
        last_name: str = ""
    ) -> UserModel:
        user = await self.get_user_by_telegram_id(telegram_id)
        if not user:
            user = UserModel(
                telegram_id=str(telegram_id),
                username=username,
                first_name=first_name,
                last_name=last_name
            )
            self.session.add(user)
            await self.session.commit()
        else:
            user.username = username
            user.first_name = first_name
            user.last_name = last_name
            await self.session.commit()
        return user

    async def update_user_questions(self, telegram_id: int) -> bool:
        user = await self.get_user_by_telegram_id(telegram_id)
        if user:
            user.questions_used += 1
            await self.session.commit()
            return True
        return False

    async def activate_premium(self, telegram_id: int) -> bool:
        try:
            user = await self.get_user_by_telegram_id(telegram_id)
            if user:
                user.is_premium = True
                if user.premium_until and user.premium_until > datetime.now():
                    user.premium_until = user.premium_until + timedelta(days=30)
                else:
                    user.premium_until = datetime.now() + timedelta(days=30)
                await self.session.commit()
                return True
            else:
                return False
        except Exception as e:
            print(f"Error activating premium: {e}")
            await self.session.rollback()
            return False