from typing import Any, Optional

from pydantic import validator
from sqlalchemy.sql.schema import Column
from sqlalchemy.sql.sqltypes import Unicode
from sqlalchemy_utils.primitives.country import Country
from sqlalchemy_utils.types import CountryType
from sqlmodel.main import Field, Relationship

from .base import AbstractBase


__all__ = [
    'StateProvince',
    'City',
    'Street',
    'Address',
]


class StateProvince(AbstractBase, table=True):
    __tablename__ = 'state_province'

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

    # Relationships
    cities: list['City'] = Relationship(
        back_populates='state_province', sa_relationship_kwargs={'lazy': 'selectin'}
    )

    @validator('country', pre=True)
    def country_as_uppercase_string(cls, v: Any) -> str:
        if isinstance(v, Country):
            return v.code
        if isinstance(v, str):
            return v.upper()
        raise TypeError


class City(AbstractBase, table=True):
    __tablename__ = 'city'

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

    # Relationships
    state_province_id: Optional[int] = Field(
        foreign_key='state_province.id', default=None, nullable=False, index=True
    )
    state_province: Optional[StateProvince] = Relationship(
        back_populates='cities', sa_relationship_kwargs={'lazy': 'selectin'}
    )

    streets: list['Street'] = Relationship(
        back_populates='city', sa_relationship_kwargs={'lazy': 'selectin'}
    )


class Street(AbstractBase, table=True):
    __tablename__ = 'street'

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

    # Relationships
    city_id: Optional[int] = Field(
        foreign_key='city.id', default=None, nullable=False, index=True
    )
    city: Optional[City] = Relationship(
        back_populates='streets', sa_relationship_kwargs={'lazy': 'selectin'}
    )

    addresses: list['Address'] = Relationship(
        back_populates='street', sa_relationship_kwargs={'lazy': 'selectin'}
    )


class Address(AbstractBase, table=True):
    __tablename__ = 'address'

    # Fields
    house_number: str = Field(max_length=8, nullable=False)
    supplement: str = Field(max_length=255)

    # Relationships
    street_id: Optional[int] = Field(
        foreign_key='street.id', default=None, nullable=False, index=True
    )
    street: Optional[Street] = Relationship(
        back_populates='addresses', sa_relationship_kwargs={'lazy': 'selectin'}
    )