from sqlalchemy.orm import Session
from watcher.db import Schedule, Service
from watcher.db.base import dbconnect
from watcher.config import settings
from .base import lock_task

import logging

logger = logging.getLogger(__name__)


@lock_task
@dbconnect
def sync_important_schedules(session: Session):
    from yql.api.v1.client import YqlClient
    yql_client = YqlClient(token=settings.YQL_OAUTH_TOKEN)

    service_schedule_db_map = {}
    db_important_schedules = set()
    service_schedule_query = session.query(Service, Schedule).join(Schedule, Schedule.service_id == Service.id)
    for service, schedule in service_schedule_query.all():
        service_schedule_db_map[(service.slug, schedule.slug)] = schedule.id
        if schedule.is_important:
            db_important_schedules.add((service.slug, schedule.slug))

    yt_important_schedules = set()
    query = [
        'USE hahn;',
        'PRAGMA yt.InferSchema = \'1\'; '
        'PRAGMA yt.TmpFolder="//home/abc/watcher_db"; '
        'SELECT service, schedule ',
        'FROM `home/test_assessor/zbp/service_duty_tier`',
    ]
    request = yql_client.query(' '.join(query))
    request.run()
    for table in request.get_results():
        table.fetch_full_data()
        yt_important_schedules |= set(filter(lambda row: row in service_schedule_db_map, table.rows))

    to_add = yt_important_schedules.difference(db_important_schedules)
    to_remove = db_important_schedules.difference(yt_important_schedules)

    to_add_ids = {service_schedule_db_map[slug_pair] for slug_pair in to_add}
    to_remove_ids = {service_schedule_db_map[slug_pair] for slug_pair in to_remove}

    for to_update_ids, value in (to_add_ids, True), (to_remove_ids, False):
        session.query(Schedule).filter(Schedule.id.in_(to_update_ids)).update(
            {Schedule.is_important: value},
            synchronize_session=False,
        )
    logger.info('Finish syncing important schedules')
