from api.db.connection import engine, SessionLocal
from api.models.subscription import Subscription
from api.models.base import Base
from datetime import datetime, timedelta

async def init_db():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)

async def set_subscription(user_id: int, uuid: str, until: datetime):
    async with SessionLocal() as session:
        sub = await session.get(Subscription, user_id)
        if sub:
            sub.uuid = uuid
            sub.until = until
        else:
            sub = Subscription(user_id=user_id, uuid=uuid, until=until)
            session.add(sub)
        await session.commit()

async def get_subscription(user_id: int):
    async with SessionLocal() as session:
        sub = await session.get(Subscription, user_id)
        if sub:
            return {"uuid": sub.uuid, "until": sub.until}
        return None

async def extend_subscription(user_id: int, months: int):
    sub = await get_subscription(user_id)
    if sub:
        new_until = max(sub["until"], datetime.now()) + timedelta(days=30*months)
        await set_subscription(user_id, sub["uuid"], new_until)
        return new_until
    return None