"""
script
To run:
1. make venv
2. . venv/bin/activate
3. python generate_uow_universal.py --env=qa --start-date=20210620 --end-date=20210630 \
--report-name=membersvisittotals_daily --dsp=appreciationengine --version=v1 --licensor=sme

python generate_uow_universal.py --env=dev --start-date=20210601 --end-date=20210630 \
--report-name=charts_weekly_viral --dsp=spotify --version=v1 --licensor=sme --frequency='0 0 * * 4'
( python generate_uow_universal.py -h)
"""
import argparse
import logging
from datetime import datetime, timedelta, timezone

import boto3
import croniter
from db_schema.schemas import slz
from slz_storage.postgres import connection


UOW_MOCK = "{dsp}-{day}-{licensor_name}-{report_name}-{version}"


def run(params: argparse.Namespace):
    logging.basicConfig(format='[CREATE UOW/CS] %(levelname)s:%(message)s', level=logging.INFO)
    logging.info("Start script execution")
    secretsmanager = boto3.client('secretsmanager')
    pg = connection.get_session_from_secret_key(secretsmanager, f'delphi/{params.env}/slz/storage/pg_proxy/user')
    try:
        start_date = params.start_date
        end_date = params.end_date
        licensor_name = params.licensor
        now = datetime.now(tz=timezone.utc)

        report: slz.Report = pg.query(slz.Report).filter(
            slz.Report.report_name == params.report_name).one()
        licensors = {
            item.licensor_name: item.licensor_id for item in pg.query(slz.Licensor).all()
        }
        if not report.data_source.data_source_name == params.dsp:
            raise ValueError('dsp and report does not match')
        already_existed_uow = pg.query(slz.UnitOfWork).filter(
            slz.UnitOfWork.report_date >= start_date.date(),
            slz.UnitOfWork.report_date <= end_date.date(),
            slz.UnitOfWork.report_id == report.report_id,
        ).with_entities(
            slz.UnitOfWork.unit_of_work_code
        ).all()
        already_existed_uow = {res[0] for res in already_existed_uow}

        while start_date <= end_date:
            logging.info("Report date: %s", start_date)

            if params.frequency and not check_date_valid_for_frequency(params.frequency, start_date):
                logging.info('Skip date %s, does not match frequency', start_date.date())
                start_date += timedelta(days=1)
                continue

            unit_of_work_code = UOW_MOCK.format(
                dsp=params.dsp,
                day=start_date.strftime('%Y%m%d'),
                licensor_name=licensor_name,
                report_name=params.report_name,
                version=params.version,
            )
            reprocess_id = ''

            # skip already created
            if unit_of_work_code in already_existed_uow:
                if not params.reprocess:
                    logging.info("Already exists, skip: %s", unit_of_work_code)
                    start_date += timedelta(days=1)
                    continue
                logging.info("Already exists, create reprocess UoW: %s", unit_of_work_code)
                reprocess_id = now.strftime('%Y%m%dT%H%M%S')

            uow_record = slz.UnitOfWork(
                unit_of_work_code=unit_of_work_code,
                reprocess_id=reprocess_id,
                report_date=start_date.date(),
                report_id=report.report_id,
                licensor_id=licensors[licensor_name],
                version='v1',
                activity_status=slz.ActivityStatusEnum.IDLE.value,
                completeness_status=slz.CompletenessStatusEnum.ACTIVE.value,
                unit_of_work_type=slz.UnitOfWorkTypeEnum.BACKFILL.value,
                next_run_at=now + timedelta(minutes=1),
                created_at=now,
                last_updated_at=now,
                is_force_complete=False,
                priority=slz.UnitOfWorkPriorityEnum.PRIORITY_6.value
            )
            logging.info("Add: %s", uow_record.readable)
            try:
                pg.add(uow_record)
                pg.commit()
            except Exception as err:
                logging.error("Error during commit: %s", str(err))
                pg.rollback()
            start_date += timedelta(days=1)

    finally:
        pg.close()
    logging.info("Finish script execution")


def check_date_valid_for_frequency(frequency: str, report_date: datetime) -> bool:
    start_time = report_date - timedelta(days=1)
    rule = croniter.croniter(frequency, start_time)
    next_date = rule.get_next(datetime)
    return next_date.date() == report_date.date()


if __name__ == '__main__':
    def date_type(value):
        try:
            value = datetime.strptime(value, '%Y%m%d')
        except Exception:
            raise argparse.ArgumentTypeError("Not a valid date")
        return value

    def frequency_type(value):
        try:
            croniter.croniter(value)
        except Exception:
            raise argparse.ArgumentTypeError("Not a valid schedule")
        return value


    parser = argparse.ArgumentParser()
    parser.add_argument("--env", type=str, required=True, choices=['dev', 'qa', 'stage', 'prod'])
    parser.add_argument("--dsp", type=str, required=True)
    parser.add_argument("--report-name", type=str, required=True)
    parser.add_argument("--start-date", type=date_type, required=True,
                        help="in format YYYYMMDD"
                        )
    parser.add_argument("--end-date", type=date_type, required=True,
                        help="in format YYYYMMDD"
                        )
    parser.add_argument("--frequency", type=frequency_type, required=False,
                        help="cron schedule expression, for non daily units"
                        )
    parser.add_argument("--licensor", type=str, required=True)
    parser.add_argument("--version", type=str, required=True)
    parser.add_argument(
        '--reprocess', action='store_true',
        help='add this flag to create reprocessing UoW for existing UoW',
    )
    params = parser.parse_args()

    run(params)