from sqlalchemy import select
from intranet.trip.src.db.tables import person_table, city_table, company_table
from intranet.trip.src.db.utils import json_column


def get_city_column(correlate_table, correlate_field=None, label='city'):
    correlate_field = correlate_field or label + '_id'
    subquery = (
        select([city_table])
        .where(city_table.c.city_id == correlate_table.c[correlate_field])
        .correlate(correlate_table)
    )
    return json_column(label, subquery)


def get_person_column(correlate_table, correlate_field=None, label='person', person_id: int = None):
    correlate_field = correlate_field or label + '_id'
    subquery = (
        select([person_table])
        .where(person_table.c.person_id == correlate_table.c[correlate_field])
    )
    if person_id is not None:
        subquery = subquery.where(correlate_table.c[correlate_field] == person_id)
    return json_column(label, subquery.correlate(correlate_table))


def get_company_column():
    subquery = (
        select([
            company_table,
        ])
        .where(person_table.c.company_id == company_table.c.company_id)
        .correlate(person_table)
    )
    return json_column('company', subquery)


def get_person_with_company_column(correlate_table, correlate_field=None, label='person'):
    correlate_field = correlate_field or label + '_id'
    subquery = (
        select([
            person_table,
            get_company_column(),
        ])
        .where(person_table.c.person_id == correlate_table.c[correlate_field])
        .correlate(correlate_table)
    )
    return json_column(label, subquery)
