import uuid import secrets import string from datetime import datetime, timezone from typing import Optional from sqlalchemy import String, DateTime, ForeignKey from sqlalchemy.orm import Mapped, mapped_column, relationship from sqlalchemy.dialects.postgresql import UUID from app.core.database import Base def generate_invite_code() -> str: """Generate an 8-character alphanumeric invite code.""" alphabet = string.ascii_uppercase + string.digits # Remove confusing characters: 0, O, I, 1 alphabet = alphabet.replace("0", "").replace("O", "").replace("I", "").replace("1", "") return "".join(secrets.choice(alphabet) for _ in range(8)) class InviteCode(Base): __tablename__ = "invite_codes" id: Mapped[uuid.UUID] = mapped_column( UUID(as_uuid=True), primary_key=True, default=uuid.uuid4 ) code: Mapped[str] = mapped_column( String(16), unique=True, nullable=False, index=True, default=generate_invite_code ) created_by_id: Mapped[uuid.UUID] = mapped_column( UUID(as_uuid=True), ForeignKey("users.id"), nullable=False ) used_by_id: Mapped[Optional[uuid.UUID]] = mapped_column( UUID(as_uuid=True), ForeignKey("users.id"), nullable=True ) expires_at: Mapped[Optional[datetime]] = mapped_column( DateTime(timezone=True), nullable=True ) email: Mapped[Optional[str]] = mapped_column(String(255), nullable=True, index=True) assigned_plan: Mapped[str] = mapped_column(String(50), nullable=False, server_default="free") trial_duration_days: Mapped[Optional[int]] = mapped_column(nullable=True) email_sent_at: Mapped[Optional[datetime]] = mapped_column( DateTime(timezone=True), nullable=True ) note: Mapped[Optional[str]] = mapped_column(String(255), nullable=True) created_at: Mapped[datetime] = mapped_column( DateTime(timezone=True), default=lambda: datetime.now(timezone.utc) ) used_at: Mapped[Optional[datetime]] = mapped_column( DateTime(timezone=True), nullable=True ) # Relationships created_by: Mapped["User"] = relationship( "User", foreign_keys=[created_by_id], backref="created_invite_codes" ) used_by: Mapped[Optional["User"]] = relationship( "User", foreign_keys=[used_by_id], backref="used_invite_code" ) @property def is_used(self) -> bool: """Check if the invite code has been used.""" return self.used_by_id is not None @property def is_expired(self) -> bool: """Check if the invite code has expired.""" if self.expires_at is None: return False return datetime.now(timezone.utc) > self.expires_at @property def is_valid(self) -> bool: """Check if the invite code is valid (not used and not expired).""" return not self.is_used and not self.is_expired @property def has_trial(self) -> bool: return self.trial_duration_days is not None and self.trial_duration_days > 0 @property def email_sent(self) -> bool: return self.email_sent_at is not None