models.py

from __future__ import annotations

import enum
from datetime import datetime
from typing import Any, Dict

from sqlalchemy import (
    DDL,
    JSON,
    Boolean,
    Column,
    Date,
    DateTime,
    Enum,
    Float,
    ForeignKey,
    Index,
    Integer,
    String,
    Time,
    UniqueConstraint,
    event,
)
from sqlalchemy.dialects.postgresql import TSVECTOR
from sqlalchemy.exc import IntegrityError, SQLAlchemyError
from sqlalchemy.ext.declarative.api import DeclarativeMeta, declarative_base
from sqlalchemy.orm import Session, relationship
from sqlalchemy.sql.schema import CheckConstraint

import app.routers.salary.config as SalaryConfig
from app.config import PSQL_ENVIRONMENT
from app.dependencies import logger
from app.internal.privacy import PrivacyKinds

Base: DeclarativeMeta = declarative_base()


class UserFeature(Base):
    __tablename__ = "user_feature"

    id = Column(Integer, primary_key=True, index=True)
    feature_id = Column("feature_id", Integer, ForeignKey("features.id"))
    user_id = Column("user_id", Integer, ForeignKey("users.id"))

    is_enable = Column(Boolean, default=False)


class User(Base):
    __tablename__ = "users"

    id = Column(Integer, primary_key=True, index=True)
    username = Column(String, unique=True, nullable=False)
    email = Column(String, unique=True, nullable=False)
    password = Column(String, nullable=False)
    full_name = Column(String)
    description = Column(String, default="Happy new user!")
    avatar = Column(String, default="profile.png")
    telegram_id = Column(String, unique=True)
    is_active = Column(Boolean, default=False)
    is_active_exercise = Column(Boolean, default=False)
    disabled = Column(Boolean, default=False, nullable=False)
    privacy = Column(String, default="Private", nullable=False)
    is_manager = Column(Boolean, default=False)
    language_id = Column(Integer, ForeignKey("languages.id"))
    target_weight = Column(Float, nullable=True)

    owned_events = relationship(
        "Event",
        cascade="all, delete",
        back_populates="owner",
    )
    events = relationship(
        "UserEvent",
        cascade="all, delete",
        back_populates="participants",
    )
    salary_settings = relationship(
        "SalarySettings",
        cascade="all, delete",
        back_populates="user",
    )
    comments = relationship("Comment", back_populates="user")
    tasks = relationship(
        "Task", cascade="all, delete", back_populates="owner")

    features = relationship("Feature", secondary=UserFeature.__tablename__)
    oauth_credentials = relationship(
        "OAuthCredentials",
        cascade="all, delete",
        back_populates="owner",
        uselist=False,
    )

    notes = relationship("Note", back_populates="creator")

    def __repr__(self):
        return f"<User {self.id}>"

    @staticmethod
    async def get_by_username(db: Session, username: str) -> User:
        """query database for a user by unique username"""
        return db.query(User).filter(User.username == username).first()


class UserExercise(Base):
    """
    Table name user exercise
    Save when user start his exercise
    """
    __tablename__ = "user_exercise"

    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(Integer, ForeignKey("users.id"))
    start_date = Column(DateTime, nullable=False)


class Feature(Base):
    __tablename__ = "features"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)
    route = Column(String, nullable=False)
    creator = Column(String, nullable=True)
    description = Column(String, nullable=False)

    users = relationship("User", secondary=UserFeature.__tablename__)


class Event(Base):
    __tablename__ = "events"

    id = Column(Integer, primary_key=True, index=True)
    title = Column(String, nullable=False)
    start = Column(DateTime, nullable=False)
    end = Column(DateTime, nullable=False)
    content = Column(String)
    location = Column(String, nullable=True)
    latitude = Column(String, nullable=True)
    longitude = Column(String, nullable=True)
    vc_link = Column(String, nullable=True)
    is_google_event = Column(Boolean, default=False)
    color = Column(String, nullable=True)
    all_day = Column(Boolean, default=False)
    invitees = Column(String)
    privacy = Column(String, default=PrivacyKinds.Public.name, nullable=False)
    emotion = Column(String, nullable=True)
    image = Column(String, nullable=True)
    availability = Column(Boolean, default=True, nullable=False)

    owner_id = Column(Integer, ForeignKey("users.id"))
    category_id = Column(Integer, ForeignKey("categories.id"))

    owner = relationship("User", back_populates="owned_events")
    participants = relationship(
        "UserEvent",
        cascade="all, delete",
        back_populates="events",
    )
    shared_list = relationship(
        "SharedList",
        uselist=False,
        back_populates="event",
    )
    comments = relationship("Comment", back_populates="event")

    # PostgreSQL
    if PSQL_ENVIRONMENT:
        events_tsv = Column(TSVECTOR)
        __table_args__ = (
            Index("events_tsv_idx", "events_tsv", postgresql_using="gin"),
        )

    def __repr__(self):
        return f"<Event {self.id}>"


class UserEvent(Base):
    __tablename__ = "user_event"

    id = Column(Integer, primary_key=True, index=True)
    user_id = Column("user_id", Integer, ForeignKey("users.id"))
    event_id = Column("event_id", Integer, ForeignKey("events.id"))

    events = relationship("Event", back_populates="participants")
    participants = relationship("User", back_populates="events")

    def __repr__(self):
        return f"<UserEvent ({self.participants}, {self.events})>"


class Language(Base):
    __tablename__ = "languages"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, unique=True, nullable=False)


class Category(Base):
    __tablename__ = "categories"

    __table_args__ = (UniqueConstraint("user_id", "name", "color"),)
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)
    color = Column(String, nullable=False)
    user_id = Column(Integer, ForeignKey("users.id"), nullable=False)

    @staticmethod
    def create(
        db_session: Session,
        name: str,
        color: str,
        user_id: int,
    ) -> Category:
        try:
            category = Category(name=name, color=color, user_id=user_id)
            db_session.add(category)
            db_session.flush()
            db_session.commit()
            db_session.refresh(category)
        except (SQLAlchemyError, IntegrityError) as e:
            logger.error(f"Failed to create category: {e}")
            raise e
        else:
            return category

    def to_dict(self) -> Dict[str, Any]:
        return {c.name: getattr(self, c.name) for c in self.__table__.columns}

    def __repr__(self) -> str:
        return f"<Category {self.id} {self.name} {self.color}>"


class PSQLEnvironmentError(Exception):
    pass


# PostgreSQL
if PSQL_ENVIRONMENT:
    trigger_snippet = DDL(
        """
    CREATE TRIGGER ix_events_tsv_update BEFORE INSERT OR UPDATE
    ON events
    FOR EACH ROW EXECUTE PROCEDURE
    tsvector_update_trigger(events_tsv,'pg_catalog.english','title','content')
    """,
    )

    event.listen(
        Event.__table__,
        "after_create",
        trigger_snippet.execute_if(dialect="postgresql"),
    )


class InvitationStatusEnum(enum.Enum):
    UNREAD = 0
    ACCEPTED = 1
    DECLINED = 2


class MessageStatusEnum(enum.Enum):
    UNREAD = 0
    READ = 1


class Invitation(Base):
    __tablename__ = "invitations"

    id = Column(Integer, primary_key=True, index=True)
    creation = Column(DateTime, default=datetime.now, nullable=False)
    status = Column(
        Enum(InvitationStatusEnum),
        default=InvitationStatusEnum.UNREAD,
        nullable=False,
    )

    recipient_id = Column(Integer, ForeignKey("users.id"))
    event_id = Column(Integer, ForeignKey("events.id"))
    recipient = relationship("User")
    event = relationship("Event")

    def decline(self, session: Session) -> None:
        """declines the invitation."""
        self.status = InvitationStatusEnum.DECLINED
        session.merge(self)
        session.commit()

    def accept(self, session: Session) -> None:
        """Accepts the invitation by creating an
        UserEvent association that represents
        participantship at the event."""

        association = UserEvent(
            user_id=self.recipient.id,
            event_id=self.event.id,
        )
        self.status = InvitationStatusEnum.ACCEPTED
        session.merge(self)
        session.add(association)
        session.commit()

    def __repr__(self):
        return f"<Invitation ({self.event.owner} to {self.recipient})>"


class Message(Base):
    __tablename__ = "messages"

    id = Column(Integer, primary_key=True, index=True)
    body = Column(String, nullable=False)
    link = Column(String)
    creation = Column(DateTime, default=datetime.now, nullable=False)
    status = Column(
        Enum(MessageStatusEnum),
        default=MessageStatusEnum.UNREAD,
        nullable=False,
    )

    recipient_id = Column(Integer, ForeignKey("users.id"))
    recipient = relationship("User")

    def mark_as_read(self, session):
        self.status = MessageStatusEnum.READ
        session.merge(self)
        session.commit()

    def __repr__(self):
        return f"<Message {self.id}>"


class UserSettings(Base):
    __tablename__ = "user_settings"

    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(Integer, ForeignKey("users.id"))
    music_on = Column(Boolean, default=False, nullable=False)
    music_vol = Column(Integer, default=None)
    sfx_on = Column(Boolean, default=False, nullable=False)
    sfx_vol = Column(Integer, default=None)
    primary_cursor = Column(String, default="default", nullable=False)
    secondary_cursor = Column(String, default="default", nullable=False)
    video_game_releases = Column(Boolean, default=False)


class AudioTracks(Base):
    __tablename__ = "audio_tracks"

    id = Column(Integer, primary_key=True, index=True)
    title = Column(String, nullable=False, unique=True)
    is_music = Column(Boolean, nullable=False)


class UserAudioTracks(Base):
    __tablename__ = "user_audio_tracks"

    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(Integer, ForeignKey("users.id"))
    track_id = Column(Integer, ForeignKey("audio_tracks.id"))


class OAuthCredentials(Base):
    __tablename__ = "oauth_credentials"

    id = Column(Integer, primary_key=True, index=True)
    token = Column(String)
    refresh_token = Column(String)
    token_uri = Column(String)
    client_id = Column(String)
    client_secret = Column(String)
    expiry = Column(DateTime)

    user_id = Column(Integer, ForeignKey("users.id"))
    owner = relationship("User", back_populates=__tablename__, uselist=False)


class SalarySettings(Base):
    # Code revision required after categories feature is added
    # Code revision required after holiday times feature is added
    # Code revision required after Shabbat times feature is added
    __tablename__ = "salary_settings"

    user_id = Column(
        Integer,
        ForeignKey("users.id"),
        primary_key=True,
    )
    # category_id = Column(
    #     Integer, ForeignKey("categories.id"), primary_key=True,
    # )
    category_id = Column(
        Integer,
        primary_key=True,
    )
    wage = Column(
        Float,
        nullable=False,
        default=SalaryConfig.MINIMUM_WAGE,
    )
    off_day = Column(
        Integer,
        CheckConstraint("0<=off_day<=6"),
        nullable=False,
        default=SalaryConfig.SATURDAY,
    )
    # holiday_category_id = Column(
    #     Integer, ForeignKey("holiday_categories.id"), nullable=False,
    #     default=SalaryConfig.ISRAELI_JEWISH,
    # )
    holiday_category_id = Column(
        Integer,
        nullable=False,
        default=SalaryConfig.ISRAELI_JEWISH,
    )
    regular_hour_basis = Column(
        Float,
        nullable=False,
        default=SalaryConfig.REGULAR_HOUR_BASIS,
    )
    night_hour_basis = Column(
        Float,
        nullable=False,
        default=SalaryConfig.NIGHT_HOUR_BASIS,
    )
    night_start = Column(
        Time,
        nullable=False,
        default=SalaryConfig.NIGHT_START,
    )
    night_end = Column(
        Time,
        nullable=False,
        default=SalaryConfig.NIGHT_END,
    )
    night_min_len = Column(
        Time,
        nullable=False,
        default=SalaryConfig.NIGHT_MIN_LEN,
    )
    first_overtime_amount = Column(
        Float,
        nullable=False,
        default=SalaryConfig.FIRST_OVERTIME_AMOUNT,
    )
    first_overtime_pay = Column(
        Float,
        nullable=False,
        default=SalaryConfig.FIRST_OVERTIME_PAY,
    )
    second_overtime_pay = Column(
        Float,
        nullable=False,
        default=SalaryConfig.SECOND_OVERTIME_PAY,
    )
    week_working_hours = Column(
        Float,
        nullable=False,
        default=SalaryConfig.WEEK_WORKING_HOURS,
    )
    daily_transport = Column(
        Float,
        CheckConstraint(f"daily_transport<={SalaryConfig.MAXIMUM_TRANSPORT}"),
        nullable=False,
        default=SalaryConfig.STANDARD_TRANSPORT,
    )

    user = relationship("User", back_populates="salary_settings")

    # category = relationship("Category", back_populates="salary_settings")
    # holiday_category =relationship("HolidayCategory",
    #                                back_populates="salary_settings")

    def __repr__(self):
        return f"<SalarySettings ({self.user_id}, {self.category_id})>"


class WikipediaEvents(Base):
    __tablename__ = "wikipedia_events"

    id = Column(Integer, primary_key=True, index=True)
    date_ = Column(String, nullable=False)
    wikipedia = Column(String, nullable=False)
    events = Column(JSON, nullable=True)
    date_inserted = Column(DateTime, default=datetime.utcnow)


class CoronaStats(Base):
    __tablename__ = "corona_stats"

    id = Column(Integer, primary_key=True, index=True)
    date_ = Column(DateTime, nullable=False)
    date_inserted = Column(DateTime, default=datetime.utcnow)
    vaccinated = Column(Integer, nullable=False)
    vaccinated_total = Column(Integer, nullable=False)
    vaccinated_population_perc = Column(Integer, nullable=False)
    vaccinated_second_dose = Column(Integer, nullable=False)
    vaccinated_second_dose_total = Column(Integer, nullable=False)
    vaccinated_second_dose_perc = Column(Float, nullable=False)


class Quote(Base):
    __tablename__ = "quotes"

    id = Column(Integer, primary_key=True, index=True)
    text = Column(String, nullable=False)
    author = Column(String)


class Country(Base):
    __tablename__ = "countries"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False, unique=True)
    timezone = Column(String, nullable=False)


class Comment(Base):
    __tablename__ = "comments"

    id = Column(Integer, primary_key=True)
    user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
    event_id = Column(Integer, ForeignKey("events.id"), nullable=False)
    content = Column(String, nullable=False)
    time = Column(DateTime, nullable=False)

    user = relationship("User", back_populates="comments")
    event = relationship("Event", back_populates="comments")

    def __repr__(self):
        return f"<Comment {self.id}>"


class Zodiac(Base):
    __tablename__ = "zodiac-signs"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)
    start_month = Column(Integer, nullable=False)
    start_day_in_month = Column(Integer, nullable=False)
    end_month = Column(Integer, nullable=False)
    end_day_in_month = Column(Integer, nullable=False)

    def __repr__(self):
        return (
            f"<Zodiac "
            f"{self.name} "
            f"{self.start_day_in_month}/{self.start_month}-"
            f"{self.end_day_in_month}/{self.end_month}>"
        )


class SharedListItem(Base):
    __tablename__ = "shared_list_item"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)
    amount = Column(Float, nullable=False)
    participant = Column(String, nullable=True)
    notes = Column(String, nullable=True)
    shared_list_id = Column(Integer, ForeignKey("shared_list.id"))

    shared_list = relationship("SharedList", back_populates="items")

    def __repr__(self):
        return (
            f"<Item {self.id}: {self.name} "
            f"Amount: {self.amount} "
            f"Participant: {self.participant} "
            f"Notes: {self.notes})>"
        )


class SharedList(Base):
    __tablename__ = "shared_list"

    id = Column(Integer, primary_key=True, index=True)
    event_id = Column(String, ForeignKey("events.id"))
    title = Column(String, nullable=True)

    items = relationship("SharedListItem", back_populates="shared_list")
    event = relationship("Event", back_populates="shared_list")

    def __repr__(self):
        return f"<Shared list {self.id}: " f"Items: {self.items}>"


class Joke(Base):
    __tablename__ = "jokes"

    id = Column(Integer, primary_key=True, index=True)
    text = Column(String, nullable=False)


class InternationalDays(Base):
    __tablename__ = "international_days"

    id = Column(Integer, primary_key=True, index=True)
    day = Column(Integer, nullable=False)
    month = Column(Integer, nullable=False)
    international_day = Column(String, nullable=False)


class Task(Base):
    __tablename__ = "tasks"

    id = Column(Integer, primary_key=True, index=True)
    title = Column(String, nullable=False)
    description = Column(String, nullable=False)
    is_done = Column(Boolean, default=False)
    is_important = Column(Boolean, nullable=False)
    date = Column(Date, nullable=False)
    time = Column(Time, nullable=False)
    owner_id = Column(Integer, ForeignKey("users.id"))

    owner = relationship("User", back_populates="tasks")


# insert language data


def insert_data(target, session: Session, **kw):
    """insert language data
    Credit to adrihanu   https://stackoverflow.com/users/9127249/adrihanu
    https://stackoverflow.com/questions/17461251"""
    session.execute(
        target.insert(),
        {"id": 1, "name": "English"},
        {"id": 2, "name": "עברית"},
    )


event.listen(Language.__table__, "after_create", insert_data)


class Note(Base):
    __tablename__ = "notes"

    id = Column(Integer, primary_key=True, index=True)
    title = Column(String, nullable=False)
    description = Column(String)
    timestamp = Column(DateTime, default=datetime.utcnow, nullable=False)
    user_id = Column(Integer, ForeignKey("users.id"))

    creator = relationship("User", back_populates="notes")

    def __repr__(self) -> str:
        return f"<Note {self.id} {self.title}>"