compub/src/compub/models/companies.py

from datetime import date
from typing import Optional, TYPE_CHECKING

from slugify import slugify
from sqlalchemy.engine.base import Connection
from sqlalchemy.event.api import listens_for
from sqlalchemy.orm.mapper import Mapper
from sqlalchemy.sql.expression import select
from sqlalchemy.sql.functions import count
from sqlalchemy.sql.schema import Column, Index
from sqlalchemy.sql.sqltypes import Unicode
from sqlalchemy_utils.types import CountryType
from sqlmodel.main import Field, Relationship

from compub.utils import multi_max
from .base import AbstractBase, DEFAULT_PK_TYPE as PK
from .geography import Address

if TYPE_CHECKING:
    from .courts import RegisterNumber


__all__ = [
    'LegalForm',
    'LegalFormSubcategory',
    'Industry',
    'Executive',
    'Company',
    'CompanyName',
]


class LegalForm(AbstractBase, table=True):
    __tablename__ = 'legal_form'

    # Fields
    short: str = Field(max_length=32, nullable=False, index=True)
    name: Optional[str] = Field(default=None, max_length=255, sa_column=Column(Unicode(255)))
    country: str = Field(sa_column=Column(CountryType))

    # Relationships
    subcategories: list['LegalFormSubcategory'] = Relationship(
        back_populates='legal_form', sa_relationship_kwargs={'lazy': 'selectin'}
    )

    def __str__(self) -> str:
        return str(self.short)


class LegalFormSubcategory(AbstractBase, table=True):
    __tablename__ = 'legal_form_subcategory'
    __table_args__ = (
        Index('ux_legal_form_subcategory', 'short', 'legal_form_id', unique=True),
    )

    # Fields
    short: str = Field(max_length=32, nullable=False, index=True)
    name: Optional[str] = Field(default=None, max_length=255, sa_column=Column(Unicode(255)))

    # Relationships
    legal_form_id: Optional[PK] = Field(
        foreign_key='legal_form.id', default=None, nullable=False, index=True
    )
    legal_form: Optional[LegalForm] = Relationship(
        back_populates='subcategories', sa_relationship_kwargs={'lazy': 'selectin'}
    )

    companies: list['Company'] = Relationship(
        back_populates='legal_form', sa_relationship_kwargs={'lazy': 'selectin'}
    )

    def __str__(self) -> str:
        return str(self.short)


class CompanyIndustryLink(AbstractBase, table=True):
    __tablename__ = 'company_industry'
    __table_args__ = (
        Index('ux_company_industry', 'company_id', 'industry_id', unique=True),
    )

    # Relationships
    company_id: Optional[PK] = Field(foreign_key='company.id', default=None, nullable=False, primary_key=True)
    industry_id: Optional[PK] = Field(foreign_key='industry.id', default=None, nullable=False, primary_key=True)


class CompanyExecutiveLink(AbstractBase, table=True):
    __tablename__ = 'company_executive'
    __table_args__ = (
        Index('ux_company_executive', 'company_id', 'executive_id', unique=True),
    )

    # Relationships
    company_id: Optional[PK] = Field(foreign_key='company.id', default=None, nullable=False, primary_key=True)
    executive_id: Optional[PK] = Field(foreign_key='executive.id', default=None, nullable=False, primary_key=True)


class Industry(AbstractBase, table=True):
    __tablename__ = 'industry'

    # Fields
    name: str = Field(max_length=255, nullable=False, index=True, sa_column_kwargs={'unique': True})

    # Relationships
    companies: list['Company'] = Relationship(
        back_populates='industries', link_model=CompanyIndustryLink, sa_relationship_kwargs={'lazy': 'selectin'}
    )

    def __str__(self) -> str:
        return str(self.name)


class Executive(AbstractBase, table=True):
    __tablename__ = 'executive'
    __MAX_LENGTH_NAME__: int = 255

    # Fields
    name: str = Field(
        max_length=__MAX_LENGTH_NAME__, sa_column=Column(Unicode(__MAX_LENGTH_NAME__), nullable=False, index=True)
    )

    # Relationships
    companies: list['Company'] = Relationship(
        back_populates='executives', link_model=CompanyExecutiveLink, sa_relationship_kwargs={'lazy': 'selectin'}
    )

    def __str__(self) -> str:
        return str(self.name)


class Company(AbstractBase, table=True):
    __tablename__ = 'company'

    # Fields
    visible: bool = Field(default=True, nullable=False, index=True)
    insolvent: bool = Field(default=False, nullable=False, index=True)
    founding_data: Optional[date]
    liquidation_date: Optional[date]
    # TODO: Get rid of city; implement address properly
    city: Optional[str] = Field(max_length=255, nullable=True, index=True)

    # Relationships
    legal_form_id: Optional[PK] = Field(
        foreign_key='legal_form_subcategory.id', default=None, index=True
    )
    legal_form: Optional[LegalFormSubcategory] = Relationship(
        back_populates='companies', sa_relationship_kwargs={'lazy': 'selectin'}
    )

    address_id: Optional[PK] = Field(
        foreign_key='address.id', default=None, index=True
    )
    address: Optional[Address] = Relationship(
        back_populates='companies', sa_relationship_kwargs={'lazy': 'selectin'}
    )

    industries: list[Industry] = Relationship(
        back_populates='companies', link_model=CompanyIndustryLink, sa_relationship_kwargs={'lazy': 'selectin'}
    )

    executives: list[Executive] = Relationship(
        back_populates='companies', link_model=CompanyExecutiveLink, sa_relationship_kwargs={'lazy': 'selectin'}
    )

    names: list['CompanyName'] = Relationship(
        back_populates='company', sa_relationship_kwargs={'lazy': 'selectin'}
    )

    register_numbers: list['RegisterNumber'] = Relationship(
        back_populates='company', sa_relationship_kwargs={'lazy': 'selectin'}
    )

    def __str__(self) -> str:
        return str(self.current_name or f"<Company {self.id}>")

    @property
    def current_name(self) -> Optional['CompanyName']:
        return multi_max(list(self.names), CompanyName.get_reg_date, 'date_updated', default=None)


class CompanyName(AbstractBase, table=True):
    __tablename__ = 'company_name'
    __table_args__ = (
        Index('ux_company_name_company_id', 'name', 'company_id', unique=True),
    )

    __MAX_LENGTH_NAME__: int = 768
    __MAX_SLUG_LENGTH__: int = 255

    # Fields
    name: str = Field(
        max_length=__MAX_LENGTH_NAME__, sa_column=Column(Unicode(__MAX_LENGTH_NAME__), nullable=False, index=True)
    )
    date_registered: Optional[date]
    slug: Optional[str] = Field(default=None, max_length=__MAX_SLUG_LENGTH__, index=True)

    # Relationships
    company_id: Optional[PK] = Field(
        foreign_key='company.id', default=None, nullable=False, index=True
    )
    company: Optional[Company] = Relationship(
        back_populates='names', sa_relationship_kwargs={'lazy': 'selectin'}
    )

    def __str__(self) -> str:
        return str(self.name)

    def get_reg_date(self) -> date:
        return date(1, 1, 1) if self.date_registered is None else self.date_registered

    @property
    def max_slug_length(self) -> int:
        return self.__MAX_SLUG_LENGTH__


@listens_for(CompanyName, 'before_insert')
def generate_company_name_slug(_mapper: Mapper, connection: Connection, target: CompanyName) -> None:
    if target.slug:
        return
    slug = slugify(target.name)[:(target.max_slug_length - 2)]
    statement = select(count()).select_from(CompanyName).where(CompanyName.slug.startswith(slug))
    num = connection.execute(statement).scalar()
    if num == 0:
        target.slug = slug
    else:
        target.slug = f'{slug}-{str(num + 1)}'