import argparse
import logging
import pprint
import uuid
from datetime import datetime
from typing import List

import peewee

from actor.recycler.cloud_resource_recycler import aws, gcp
from analyst.libs import consts
from analyst.resource_leak import model
from analyst.resource_leak.model import CloudAsset

logger = logging.getLogger('recycler.cloud_resource_recycle')


def load_and_clean_leaked_assets(batch_id: str, options) -> List[
    aws.CleanResult]:
    # Load leaked CloudAssets
    leaked_assets = load_leaked_assets(options)
    dryrun = options.dryrun

    # Initialize cleaners
    cleaners = {
        'aws/ec2': aws.AWSEC2Cleaner(dryrun),
        'aws/nat': aws.AWSNATCleaner(dryrun),
        'aws/clb': aws.AWSELBv1Cleaner(dryrun),
        'aws/nlb': aws.AWSELBv2Cleaner(dryrun),
        'aws/asg': aws.AWSAutoScalingGroupCleaner(dryrun),
        'aws/node_group': aws.AWSEKSNodeGroupCleaner(dryrun),
        # 'aws/eks': aws.AWSEKSCleaner(dryrun),
        'aws/ebs': aws.AWSEBSCleaner(dryrun),
        'gcp/vm': gcp.GCPVMRecycler(dryrun),
        'gcp/disk': gcp.GCPDiskRecycler(dryrun),
        'gcp/nat': gcp.GCPCloudNATRecycler(dryrun),
        'gcp/lb': gcp.GCPForwardingRuleRecycler(dryrun),
    }

    results = []
    for asset in leaked_assets:
        cleaner = cleaners.get(f'{asset.provider}/{asset.product}')
        if cleaner:
            skip, reason = cleaner.skip(asset)
            if skip:
                logger.info(f'skip to clean {asset.uniq_key()}, reason: {reason}')
                continue

            result = cleaner.clean(asset)
            result.asset = asset
            results.append(result)

            now = datetime.now()
            with model.db.atomic():
                try:
                    clean_record = result.build_db_record4save(now, batch_id)
                    model.CleanHistory.insert(clean_record).execute()
                    if result.result in {consts.CleanResult.Cleaned, }:
                        execute_result = model.CloudAsset.update(
                            status=consts.ResourceUseState.Cleaned,
                        ).where(
                            model.CloudAsset.id == asset.id,
                            model.CloudAsset.status == consts.ResourceUseState.Leaked,
                        ).execute()
                        logger.info(f'mark resource({asset.id}) as cleaned, affected_lines={execute_result}')
                except peewee.IntegrityError as e:
                    logger.warning(f"insert into clean_histories failed: {e}, continue")
        else:
            logger.info(f"No cleaner found for {asset.uniq_key()}")

    return results


def load_leaked_assets(options) -> List[model.CloudAsset]:
    where = [
        CloudAsset.service_type.in_([consts.ServiceType.DEDICATED]),
        CloudAsset.status == consts.ResourceUseState.Leaked,
        CloudAsset.resource_deleted_at.is_null(),
    ]
    if options.envs is not None:
        where.append(CloudAsset.env.in_(options.envs))
    else:
        where.append(CloudAsset.env.in_([consts.Env.Dev, consts.Env.Staging]))
    if options.providers is not None:
        where.append(CloudAsset.provider.in_(options.providers))
    if options.regions is not None:
        where.append(CloudAsset.region.in_(options.regions))
    if options.products is not None:
        where.append(CloudAsset.product.in_(options.products))
    if options.account_ids is not None:
        where.append(CloudAsset.account_id.in_(options.account_ids))
    if options.asset_id is not None:
        where.append(CloudAsset.id == options.asset_id)
    print('Query SQL: ', CloudAsset.select().where(*where).order_by(CloudAsset.id.asc()).sql())
    leaked_resources: list[CloudAsset] = CloudAsset.select().where(*where).order_by(CloudAsset.id.asc())

    ids = [str(each.id) for each in leaked_resources]
    logger.info(f"there are {len(leaked_resources)} leaked resources: {','.join(ids)}")
    return leaked_resources


def save_clean_results(results: List[aws.CleanResult], batch_id: str):
    now = datetime.now()
    to_save = []
    for each in results:
        to_save.append({
            'asset_id': each.asset.id,
            'batch_id': batch_id,
            'result': each.result,
            'reason': each.reason,
            'created_at': now,
        })

    pprint.pprint(to_save)

    for batch in peewee.chunked(to_save, 100):
        with model.db.atomic():
            try:
                model.CleanHistory.insert_many(batch).execute()
            except peewee.IntegrityError as e:
                logger.warning(f"insert into clean_histories failed: {e}, continue")

    logger.info(f"Saved {len(results)} clean history to database for batch {batch_id}")


def run(options):
    model.init_db(model.create_db_instance())

    batch_id = str(uuid.uuid4())
    logger.info(f'run with batch id: {batch_id}')
    results = load_and_clean_leaked_assets(batch_id, options)
    logger.info(f'run with batch id: {batch_id}')


def main():
    parser = argparse.ArgumentParser(
        prog='cloud_resource_recycler',
        description='recycler for leaked cloud resource')
    parser.add_argument('--dryrun', action='store_true')
    parser.add_argument('--envs', nargs='+', default=['dev'], help='which env to clean')
    parser.add_argument('--providers', nargs='+', help='which provider to clean')
    parser.add_argument('--regions', nargs='+', help='which region to clean')
    parser.add_argument('--products', nargs='+', help='which product(s) to clean')
    parser.add_argument('--account_ids', nargs='+', help='which env to clean')
    parser.add_argument('--asset-id', dest='asset_id', type=int, help='the id of asset to clean')
    options = parser.parse_args()
    logger.info(
        f'clean for env({options.envs})/provider({options.providers})'
        f'/account_id({options.account_ids})/region({options.regions})')
    print('options: ', options)

    run(options)


if __name__ == '__main__':
    main()