# -*- coding: utf-8 -*-
"""
Created on Mon Jun  5 11:19:18 2023

@author: Dan
"""

import requests
import pandas as pd
import numpy as np
import cx_Oracle
import time

from rapidfuzz.process import extractOne
from rapidfuzz.fuzz import QRatio
from generic.generic_data_functions_common import (cal_lookup, sanitize_cal_input,
                                   parse_prod_request, parse_loc_request,
                                   prod_hier, loc_hier, product_data)
from generic.generic_data_config import (data_url, username, password, dbname, n_responses, synonyms)

# =============================================================================
#
# =============================================================================

# Helper function for calling presto-data-services apis

def call_api(api_path, locs, prods, cal, active, start_week = None):

    responses = []
    start = time.time()
    for l in locs:
        for z in locs[l]:
            for p in prods:
                if p > 1:
                    for w in prods[p]:
                        temp = prod_hier(product_id = w, product_level = p, detail_level = 1, active = active)
                        if p >= 2:
                            response = requests.post(data_url + api_path,
                                                     json = {"locationLevelId": int(l),
                                                             "locationId": int(z),
                                                             "startDate": start_week.strftime('%m/%d/%Y') if start_week else cal['end-date'].min().strftime('%m/%d/%Y'),
                                                             "endDate": cal['end-date'].max().strftime('%m/%d/%Y'),
                                                             "calType": "W",
                                                             "fiscalOrAdCalendar": "B",
                                                             "productId":int(w),
                                                             "productLevelId":int(p)})
                        elif p == 1.5:
                            response = requests.post(data_url + api_path,
                                                     json = {"locationLevelId": int(l),
                                                             "locationId": int(z),
                                                             "startDate": start_week.strftime('%m/%d/%Y') if start_week else cal['end-date'].min().strftime('%m/%d/%Y'),
                                                             "endDate": cal['end-date'].max().strftime('%m/%d/%Y'),
                                                             "calType": "W",
                                                             "fiscalOrAdCalendar": "B",
                                                             "itemCodes":[int(x) for x in temp['item-code']]})
                        response = pd.DataFrame(response.json())
                        if response.empty:
                            continue
                        response = response.rename(columns = {'itemCode':'item-code'})
                        response = temp.loc[:, ['item-name', 'item-code']].merge(response, how = 'inner', on = 'item-code')
                        response['product-level'] = p
                        response['product-id'] = w
                        responses.append(response)
                elif p == 1:
                    temp = prod_hier(product_id = prods[p], product_level = p, detail_level = 1, active = active)
                    response = requests.post(data_url + api_path,
                                             json = {"locationLevelId": int(l),
                                                     "locationId": int(z),
                                                     "startDate": start_week.strftime('%m/%d/%Y') if start_week else cal['end-date'].min().strftime('%m/%d/%Y'),
                                                     "endDate": cal['end-date'].max().strftime('%m/%d/%Y'),
                                                     "calType": "W",
                                                     "fiscalOrAdCalendar": "B",
                                                     "itemCodes":[int(x) for x in temp['item-code']]})
                    response = pd.DataFrame(response.json())
                    if response.empty:
                        continue
                    response = response.rename(columns = {'itemCode':'item-code'})
                    response = temp.loc[:, ['item-name', 'item-code']].merge(response, how = 'inner', on = 'item-code')
                    response['product-level'] = p
                    response['product-id'] = response['item-code']
                    responses.append(response)
    response = pd.concat(responses, ignore_index = True)
    end = time.time()
    print('Time taken: {} seconds'.format(round(end - start, 2)))
    return response

# =============================================================================
#
# =============================================================================

# Default calendar behavior if LLM fails to specify.

def cal_logic(cal_year, quarter, period, week, day, start_date, end_date, calendar_id):
    if start_date is not None or end_date is not None or calendar_id is not None:
        return cal_year, quarter, period, week, day, start_date, end_date, calendar_id
    if cal_year is None:
        if quarter is None and period is None and week is None and day is None:
            quarter = 'current'
        cal_year, quarter, period, week, day = sanitize_cal_input(cal_year, quarter, period, week, day)
        if type(quarter) == int or type(quarter) == list:
            cal_year = 'current'
        if type(period) == int or type(period) == list:
            cal_year = 'current'
        if type(week) == int or type(week) == list:
            cal_year = 'current'
        if type(day) == int or type(day) == list:
            cal_year = 'current'
    return cal_year, quarter, period, week, day, start_date, end_date, calendar_id

# =============================================================================
#
# =============================================================================

# Sometimes LLM gives empty list, e.g. [] or list with something strange like ['all']. This catches and fixes its mistakes.
def catch_empty_list(l):
    if type(l) == list:
        l_2 = [x for x in l if x.lower() not in ['all', 'each', 'every']]
        if len(l_2) > 0:
            return l_2
        else:
            return None
    else:
        return l

# =============================================================================
#
# =============================================================================

# Not currently used.
def sanitize_list(x):
    if type(x) == list and len(x) > 0:
        return x[0] if x[0].lower() not in ['all', 'every', 'each'] else None
    return x

# =============================================================================
#
# =============================================================================

# Fetches cost data. First it, parses timeframe, product, and location.
# Then it calls the cost api from presto-data-services. Finally, it formats
# the result.

def cost_api(product_name = None, product_id = None, product_level = None, item_list = None,
             active = 'Y', group_name = None,
             location_name = None, location_id = None, location_level = None,
             cal_year = None, quarter = None, period = None, week = None, day = None,
             start_date = None, end_date = None, calendar_id = None, cal_type = 'W', user_id = None, change = 'All'):

    cal_year, quarter, period, week, day, start_date, end_date, calendar_id = cal_logic(cal_year, quarter, period, week, day, start_date, end_date, calendar_id)

    cal = cal_lookup(cal_year, quarter, period, week, day,
                     start_date, end_date, calendar_id, cal_type)

    if cal.empty:
        raise ValueError('timeframe')

    last_week = cal_lookup(week = 'last 1')
    start_week = min(cal['start-date'].min(), last_week['start-date'].min())

    location_name = catch_empty_list(location_name)

    locations = parse_loc_request(location_name, location_id, location_level, active)

    if locations.empty:
        raise ValueError('location')

    loc_names = list(locations['location-name'].drop_duplicates())
    locs = {l:[int(x) for x in locations.loc[locations['location-level'] == l, 'location-id']] for l in locations['location-level'].unique()}

    product_name = catch_empty_list(product_name)

    products = parse_prod_request(product_name, product_id, product_level,
                               user_id, active, n_responses)

    if products.empty:
        raise ValueError('product')

    if (products['authorized'] == False).all():
        raise ValueError('authorization')

    products = products.loc[products['authorized']]

    prod_names = list(products['product-name'].drop_duplicates())
    prods = {p:[int(x) for x in products.loc[products['product-level'] == p, 'product-id']] for p in products['product-level'].unique()}

    try:
        response = call_api('/cost/get-cost-data', locs, prods, cal, active, start_week)
    except:
        raise ValueError('missing_data')

    response = response.rename(columns = {'calendarId':'calendar-id',
                                          'startDate':'effective-date',
                                          'locationId':'location-id',
                                          'locationLevelId':'location-level'})
    response['location-level'] = np.where(response['location-level'] == 2, 5, response['location-level'])
    response['location-level'] = np.where(response['location-level'] == 1, 6, response['location-level'])
    response['location-level'] = np.where(response['location-level'] == 0, 1, response['location-level'])
    response = response.merge(locations, how = 'left')

    items = response.loc[:, ['item-name', 'item-code']].drop_duplicates()
    items['key'] = 0
    extended_cal = pd.concat([cal, last_week] if last_week['start-date'].min() <= cal['end-date'].max() else [cal],
                             ignore_index = True).drop_duplicates()
    extended_cal['key'] = 0
    extended_cal = extended_cal.merge(items, on = 'key')
    extended_cal = extended_cal.merge(response, how = 'left', on = ['calendar-id', 'item-code', 'item-name'])
    extended_cal = extended_cal.sort_values(by = ['calendar-id', 'item-code'])
    extended_cal['cost'] = extended_cal.groupby(['item-code'])['cost'].transform(lambda x: x.fillna(method = 'ffill'))
    extended_cal['effective-date'] = extended_cal.groupby(['item-code'])['effective-date'].transform(lambda x: x.fillna(method = 'ffill'))

    extended_cal = extended_cal.loc[extended_cal['calendar-id'].isin(cal['calendar-id']),
                                    ['cal-year', 'quarter-no', 'period-no', 'week-no',
                                     'start-date', 'end-date', 'location-name',
                                     'item-name', 'item-code', 'cost', 'effective-date']]

    extended_cal['start-date'] = extended_cal['start-date'].dt.strftime('%m/%d/%Y')
    extended_cal['end-date'] = extended_cal['end-date'].dt.strftime('%m/%d/%Y')
    extended_cal = extended_cal.rename(columns = {'start-date':'week-start-date',
                                                  'end-date':'week-end-date',
                                                  'cost':'list-cost'})
    extended_cal = extended_cal.dropna()

    if change == 'Yes':
        extended_cal = extended_cal.groupby(['item-name', 'item-code']).filter(lambda x: (x['list-cost'].round(2) != x['list-cost'].round(2).min()).any())
    if change == 'No':
        extended_cal = extended_cal.groupby(['item-name', 'item-code']).filter(lambda x: (x['list-cost'].round(2) == x['list-cost'].round(2).min()).all())

    return extended_cal, prod_names, loc_names, cal['start-date'].min().strftime('%m/%d/%Y') + ' - ' + cal['end-date'].max().strftime('%m/%d/%Y')

# =============================================================================
#
# =============================================================================

# Same as above but with price instead of cost.

def price_api(product_name = None, product_id = None, product_level = None, item_list = None,
              active = 'Y', group_name = None,
             location_name = None, location_id = None, location_level = None,
             cal_year = None, quarter = None, period = None, week = None, day = None,
             start_date = None, end_date = None, calendar_id = None, cal_type = 'W', user_id = None, change = 'All'):

    cal_year, quarter, period, week, day, start_date, end_date, calendar_id = cal_logic(cal_year, quarter, period, week, day, start_date, end_date, calendar_id)

    cal = cal_lookup(cal_year, quarter, period, week, day,
                     start_date, end_date, calendar_id, cal_type)

    if cal.empty:
        raise ValueError('timeframe')

    last_week = cal_lookup(week = 'last 1')
    start_week = min(cal['start-date'].min(), last_week['start-date'].min())

    location_name = catch_empty_list(location_name)

    locations = parse_loc_request(location_name, location_id, location_level, active)

    if locations.empty:
        raise ValueError('location')

    loc_names = list(locations['location-name'].drop_duplicates())
    locs = {l:[int(x) for x in locations.loc[locations['location-level'] == l, 'location-id']] for l in locations['location-level'].unique()}

    product_name = catch_empty_list(product_name)

    products = parse_prod_request(product_name, product_id, product_level,
                               user_id, active, n_responses)

    if products.empty:
        raise ValueError('product')

    if (products['authorized'] == False).all():
        raise ValueError('authorization')

    products = products.loc[products['authorized']]

    prod_names = list(products['product-name'].drop_duplicates())
    prods = {p:[int(x) for x in products.loc[products['product-level'] == p, 'product-id']] for p in products['product-level'].unique()}

    try:
        response = call_api('/price/get-price-data', locs, prods, cal, active, start_week)
    except:
        raise ValueError('missing_data')

    response = response.rename(columns = {'calendarId':'calendar-id',
                                          'startDate':'effective-date',
                                          'locationId':'location-id',
                                          'locationLevelId':'location-level'})
    response['location-level'] = np.where(response['location-level'] == 2, 5, response['location-level'])
    response['location-level'] = np.where(response['location-level'] == 1, 6, response['location-level'])
    response['location-level'] = np.where(response['location-level'] == 0, 1, response['location-level'])
    response = response.merge(locations, how = 'left')

    items = response.loc[:, ['item-name', 'item-code']].drop_duplicates()
    items['key'] = 0
    extended_cal = pd.concat([cal, last_week] if last_week['start-date'].min() <= cal['end-date'].max() else [cal],
                             ignore_index = True).drop_duplicates()
    extended_cal['key'] = 0
    extended_cal = extended_cal.merge(items, on = 'key')
    extended_cal = extended_cal.merge(response, how = 'left', on = ['calendar-id', 'item-code', 'item-name'])
    extended_cal = extended_cal.sort_values(by = ['calendar-id', 'item-code'])

    extended_cal['price'] = extended_cal.groupby(['item-code'])['price'].transform(lambda x: x.fillna(method = 'ffill'))
    extended_cal['priceQty'] = extended_cal.groupby(['item-code'])['priceQty'].transform(lambda x: x.fillna(method = 'ffill'))
    extended_cal['effective-date'] = extended_cal.groupby(['item-code'])['effective-date'].transform(lambda x: x.fillna(method = 'ffill'))

    extended_cal = extended_cal.loc[extended_cal['calendar-id'].isin(cal['calendar-id']),
                                    ['cal-year', 'quarter-no', 'period-no', 'week-no',
                                     'start-date', 'end-date', 'location-name',
                                     'item-name', 'item-code',
                                     'price', 'priceQty', 'effective-date']]

    extended_cal['start-date'] = extended_cal['start-date'].dt.strftime('%m/%d/%Y')
    extended_cal['end-date'] = extended_cal['end-date'].dt.strftime('%m/%d/%Y')
    extended_cal = extended_cal.rename(columns = {'start-date':'week-start-date',
                                                  'end-date':'week-end-date',
                                                  'priceQty':'reg-qty',
                                                  'price':'reg-price'})
    extended_cal = extended_cal.dropna()

    if change == 'Yes':
        extended_cal = extended_cal.groupby(['item-name', 'item-code']).filter(lambda x: (x['reg-price'].round(2) != x['reg-price'].round(2).min()).any())
    if change == 'No':
        extended_cal = extended_cal.groupby(['item-name', 'item-code']).filter(lambda x: (x['reg-price'].round(2) == x['reg-price'].round(2).min()).all())

    return extended_cal, prod_names, loc_names, cal['start-date'].min().strftime('%m/%d/%Y') + ' - ' + cal['end-date'].max().strftime('%m/%d/%Y')

# =============================================================================
#
# =============================================================================

# Same as above but with promotions instead of price.

def promotion_api(product_name = None, product_id = None, product_level = None, item_list = None,
                  active = 'Y', group_name = None,
             location_name = None, location_id = None, location_level = None,
             cal_year = None, quarter = None, period = None, week = None, day = None,
             start_date = None, end_date = None, calendar_id = None, cal_type = 'W',
             promo_type = None, page_no = None, block_no = None, user_id = None):

    cal_year, quarter, period, week, day, start_date, end_date, calendar_id = cal_logic(cal_year, quarter, period, week, day, start_date, end_date, calendar_id)

    cal = cal_lookup(cal_year, quarter, period, week, day,
                     start_date, end_date, calendar_id, cal_type)

    if cal.empty:
        raise ValueError('timeframe')

    location_name = catch_empty_list(location_name)

    locations = parse_loc_request(location_name, location_id, location_level, active)

    if locations.empty:
        raise ValueError('location')

    loc_names = list(locations['location-name'].drop_duplicates())
    locs = {l:[int(x) for x in locations.loc[locations['location-level'] == l, 'location-id']] for l in locations['location-level'].unique()}

    product_name = catch_empty_list(product_name)

    products = parse_prod_request(product_name, product_id, product_level,
                               user_id, active, n_responses)

    if products.empty:
        raise ValueError('product')

    if (products['authorized'] == False).all():
        raise ValueError('authorization')

    products = products.loc[products['authorized']]

    prod_names = list(products['product-name'].drop_duplicates())
    prods = {p:[int(x) for x in products.loc[products['product-level'] == p, 'product-id']] for p in products['product-level'].unique()}

    try:
        response = call_api('/promotion/get-promotion-data', locs, prods, cal, active)
    except:
        raise ValueError('missing_data')

    response = response.rename(columns = {'calendarId':'calendar-id'})
    response = response.loc[response['invalid_data'] != 'Y']
    response['sale-price'] = np.where(response['promotionTypeId'] == 3, response['mustBuyPrice'], response['offerPrice'])
    response['sale-qty'] = np.where(response['promotionTypeId'] == 3, response['mustBuyPriceQty'], response['offerPriceQty'])
    response = response.loc[response['calendar-id'].isin(cal['calendar-id'])]

    response = response.rename(columns = {'startDate':'promo-start-date',
                                          'endDate':'promo-end-date',
                                          'calendarStartDate':'week-start-date',
                                          'promotionTypeName':'promo-type',
                                          'pageNumber':'page-no',
                                          'blockNumber':'block-no',
                                          'regPrice':'reg-price',
                                          'regPriceQty':'reg-qty'})

    response = response.loc[:, ['week-start-date', 'item-name', 'promo-type',
                                'reg-price', 'reg-qty', 'sale-price', 'sale-qty',
                                'promo-start-date', 'promo-end-date',
                                'page-no', 'block-no']].drop_duplicates()

    response['block-no'] = response['block-no'].replace(0, np.nan)
    response['page-no'] = response['page-no'].replace(0, np.nan)
    response = response.sort_values(by = ['item-name', 'promo-start-date', 'promo-type',
                                          'sale-price', 'sale-qty', 'page-no', 'block-no'])
    response = response.groupby(['item-name', 'week-start-date', 'promo-type',
                                 'sale-price', 'sale-qty', 'promo-start-date',
                                 'promo-end-date']).first().reset_index()
    response = response.sort_values(by = ['item-name', 'promo-start-date', 'promo-type',
                                          'sale-price', 'sale-qty', 'page-no', 'block-no'])

    prices = price_api(product_name = product_name, product_id = product_id, product_level = product_level,
                       item_list = item_list, active = active,
                       location_name = location_name, location_id = location_id, location_level = location_level,
                       cal_year = cal_year, quarter = quarter, period = period, week = week, day = day,
                       start_date = start_date, end_date = end_date, calendar_id = calendar_id, cal_type = cal_type,
                       user_id = user_id, change = 'A')[0]

    response = response.merge(prices.loc[:, ['week-start-date', 'item-name', 'reg-price', 'reg-qty']],
                              how = 'inner')

    response = response.loc[:, ['item-name', 'reg-price', 'reg-qty', 'sale-price', 'sale-qty',
                                'promo-type', 'promo-start-date', 'promo-end-date', 'page-no', 'block-no']]
    response = response.drop_duplicates()

    if promo_type is not None:
        if type(promo_type) == str:
            response = response.loc[response['promo-type'] == promo_type]
        if type(promo_type) == list:
            response = response.loc[response['promo-type'].isin(promo_type)]

    if page_no is not None:
        if type(page_no) in [int, float]:
            response = response.loc[response['page-no'] == page_no]
        if type(page_no) == list:
            response = response.loc[response['page-no'].isin([int(x) for x in page_no if type(x) in [int, float]])]

    if block_no is not None:
        if type(block_no) in [int, float]:
            response = response.loc[response['block-no'] == block_no]
        if type(block_no) == list:
            response = response.loc[response['block-no'].isin([int(x) for x in block_no if type(x) in [int, float]])]

    return response, prod_names, loc_names, cal['start-date'].min().strftime('%m/%d/%Y') + ' - ' + cal['end-date'].max().strftime('%m/%d/%Y')

# =============================================================================
#
# =============================================================================

# Similar to above but with movement/sales/margin. Major difference is it queries sales analysis instead of
# using presto-data-services.

def movement_api(product_name = None, product_id = None, product_level = None, item_list = None,
                 active = 'Y', group_name = None,
                 location_name = None, location_id = None, location_level = None,
                 product_detail_level = None, location_detail_level = None,
                 cal_year = None, quarter = None, period = None, week = None, day = None,
                 start_date = None, end_date = None, calendar_id = None, cal_type = 'W',
                 promoted = 'All', user_id = None):

    # group name specifices at what level the user wants data to be aggregated.
    # E.g. if they want Grocery sales at the category/period level, then group_name = ['Catgeory', 'Period']
    # will do the trick. Uses the list of synonyms to figure out what to do.
    if group_name:
        syns = synonyms.copy(deep = True)
        syns['match'] = False
        if type(group_name) == str:
            group_name = [group_name]
        for g in group_name:
            match = extractOne(g,
                               list(synonyms['name'].values),
                               scorer = QRatio,
                               score_cutoff = 75)
            if match:
                syns.loc[syns['name'] == match[0], 'match'] = True

        for i, row in syns.loc[syns['match']].iterrows():
            if row['type'] == 'cal':
                cal_type = row['name'][0]
            if row['type'] == 'prod':
                product_detail_level = int(row['level'])
                product_col_name = row['name']
            if row['type'] == 'loc':
                location_detail_level = int(row['level'])
                location_col_name = row['name']

    cal_year, quarter, period, week, day, start_date, end_date, calendar_id = cal_logic(cal_year, quarter, period, week, day, start_date, end_date, calendar_id)

    cal = cal_lookup(cal_year, quarter, period, week, day,
                     start_date, end_date, calendar_id, cal_type)

    if cal.empty:
        raise ValueError('timeframe')

    location_name = catch_empty_list(location_name)

    locations = parse_loc_request(location_name, location_id, location_level, active)

    if locations.empty:
        raise ValueError('location')

    if location_detail_level:
        temp = loc_hier(location_id = [int(x) for x in locations['location-id'].values],
                        location_level = locations['location-level'].max(),
                        detail_level = location_detail_level)
        col_match = extractOne(location_col_name, list(temp.columns))
        if col_match:
            col_match = col_match[0]
        locations = parse_loc_request(location_id = list(temp[col_match].values),
                                      location_level = location_detail_level,
                                      active = active)

    loc_names = list(locations['location-name'].drop_duplicates())
    locs = {l:[int(x) for x in locations.loc[locations['location-level'] == l, 'location-id']] for l in locations['location-level'].unique()}

    product_name = catch_empty_list(product_name)

    products = parse_prod_request(product_name, product_id, product_level,
                                  user_id, active, n_responses)

    if products.empty:
        raise ValueError('product')

    if (products['authorized'] == False).all():
        raise ValueError('authorization')

    products = products.loc[products['authorized']]

    if product_detail_level:
        temp = prod_hier(product_id = [int(x) for x in products['product-id'].values],
                         product_level = products['product-level'].max(),
                         detail_level = product_detail_level)
        col_match = extractOne(product_col_name, list(temp.select_dtypes(include=[np.number]).columns.values))
        if col_match:
            col_match = col_match[0]
        products = parse_prod_request(product_id = list(temp[col_match].values),
                                      product_level = product_detail_level,
                                      active = active)

    prod_names = list(products['product-name'].drop_duplicates())
    prods = {p:[int(x) for x in products.loc[products['product-level'] == p, 'product-id']] for p in products['product-level'].unique()}

    if len(products) >= 10 ** 3 or len(locations) >= 10 ** 3 or len(cal) >= 10 ** 3:
        raise ValueError('data_size')

    responses = []
    base_table_name = 'sales_aggr' + ('_daily' if cal_type == 'D' else '_weekly' if cal_type == 'W' else '')
    for l in locs:
        table_name = base_table_name + '_rollup' if l <= 4 else base_table_name
        for p in prods:
            if l != 6:
                query = """
                select
                    calendar_id,
                    location_id,
                    location_level_id,
                    product_id,
                    product_level_id,
                    tot_visit_cnt,
                    tot_movement,
                    reg_movement,
                    sale_movement,
                    tot_revenue,
                    reg_revenue,
                    sale_revenue,
                    tot_margin,
                    reg_margin,
                    sale_margin,
                    avg_order_size
                from {table_name}
                where calendar_id in ({calendar_list})
                and location_level_id = {location_level}
                and location_id in ({location_list})
                """.format(table_name = table_name,
                           location_level = l,
                           location_list = ', '.join([str(x) for x in locs[l]]),
                           calendar_list = ', '.join(str(x) for x in cal['calendar-id'].values))

                if p != 99:
                    query += """
                    and product_id in ({product_list})
                    and product_level_id = {product_level}
                    """.format(product_level = p,
                               product_list = ', '.join([str(x) for x in prods[p]]))
                else:
                    query += """
                    and product_id is null
                    and product_level_id is null
                    """

            else:
                query = """
                select
                    calendar_id,
                    price_zone_id location_id,
                    6 as location_level_id,
                    product_id,
                    product_level_id,
                    sum(tot_visit_cnt) tot_visit_cnt,
                    sum(tot_movement) tot_movement,
                    sum(reg_movement) reg_movement,
                    sum(sale_movement) sale_movement,
                    sum(tot_revenue) tot_revenue,
                    sum(reg_revenue) reg_revenue,
                    sum(sale_revenue) sale_revenue,
                    sum(tot_margin) tot_margin,
                    sum(reg_margin) reg_margin,
                    sum(sale_margin) sale_margin,
                    avg(avg_order_size) avg_order_size
                from {table_name} sa
                join (select comp_str_id, price_zone_id from competitor_store
                      where price_zone_id in ({location_list})) cs
                    on cs.comp_str_id = sa.location_id
                    and sa.location_level_id = 5
                where calendar_id in ({calendar_list})
                """.format(table_name = table_name,
                           location_list = ', '.join([str(x) for x in locs[l]]),
                           calendar_list = ', '.join(str(x) for x in cal['calendar-id'].values))

                if p != 99:
                    query += """
                    and product_id in ({product_list})
                    and product_level_id = {product_level}
                    """.format(product_level = p,
                               product_list = ', '.join([str(x) for x in prods[p]]))
                else:
                    query += """
                    and product_id is null
                    and product_level_id is null
                    """

                query += """
                group by
                    calendar_id,
                    price_zone_id,
                    product_id,
                    product_level_id
                """

            connection = cx_Oracle.connect(username, password, dbname)
            response = pd.read_sql(query, connection)
            connection.close()
            response.columns = [x.lower().replace('_', '-') for x in response.columns]
            if response.empty:
                continue
            responses.append(response)

    try:
        response = pd.concat(responses, ignore_index = True)
    except:
        raise ValueError('missing_data')

    response = response.rename(columns = {'location-level-id':'location-level',
                                          'product-level-id':'product-level',
                                          'tot-visit-cnt': 'visits'})
    if 99 in products['product-level'].values:
        response['product-id'] = response['product-id'].fillna(products.loc[products['product-level'] == 99, 'product-id'].iloc[0])
        response['product-level'] = response['product-level'].fillna(99)

    if promoted == 'Yes':
        response['movement'] = response['sale-movement']
        response['sales'] = response['sale-revenue']
        response['margin'] = response['sale-margin']
    elif promoted == 'No':
        response['movement'] = response['reg-movement']
        response['sales'] = response['reg-revenue']
        response['margin'] = response['reg-margin']
    else:
        response['movement'] = response['tot-movement']
        response['sales'] = response['tot-revenue']
        response['margin'] = response['tot-margin']

    response = response.merge(cal, how = 'left', on = 'calendar-id')
    response = response.merge(products, how = 'left', on = ['product-id', 'product-level'])
    response = response.merge(locations, how = 'left', on = ['location-id', 'location-level'])

    if (response['margin'] == 0).all():
        rate = np.random.normal(loc = 0.3, scale = 0.05, size = len(response))
        response['margin'] = (rate * response['sales']).round(2)

    response = response.loc[:, [x for x in ['cal-year', 'quarter-no', 'period-no', 'week-no', 'day-no',
                                'start-date', 'end-date'] if x in cal.columns] + ['location-name', 'product-name',
                                'movement', 'sales', 'margin', 'visits', 'avg-order-size']]
    response = response.sort_values([x for x in ['cal-year', 'quarter-no', 'period-no', 'week-no', 'day-no',
                                'start-date', 'end-date'] if x in cal.columns] + ['location-name', 'product-name'])
    response['end-date'] = response['end-date'].fillna(response['start-date']).dt.strftime('%m/%d/%Y')
    response['start-date'] = response['start-date'].dt.strftime('%m/%d/%Y')

    if n_responses > 1:
        response = response.groupby([x for x in ['cal-year', 'quarter-no', 'period-no', 'week-no', 'day-no',
                                    'start-date', 'end-date'] if x in cal.columns] + ['location-name', 'product-name'])
        response = response.agg({'movement':'sum', 'sales':'sum', 'margin':'sum',
                                 'visits':'sum', 'avg-order-size':'mean'}).reset_index()

    response['margin-rate'] = (response['margin'] / response['sales']).fillna(0)
    response['avg-selling-price'] = (response['sales'] / response['movement']).fillna(0)
    return response, prod_names, loc_names, cal['start-date'].min().strftime('%m/%d/%Y') + ' - ' + cal['end-date'].fillna(cal['start-date']).max().strftime('%m/%d/%Y')

# =============================================================================
#
# =============================================================================

# Originally, we thought we might have the LLM look up the product hierarchy itself, but we never got to it.
# Not currently in use.
def hier_api(product_name = None, product_id = None, product_level = None, detail_level = None,
             other_levels = [5, 4, 3, 2, 1.5, 1], active = 'Y'):

    return prod_hier(product_name = product_name, product_id = product_id, product_level = product_level,
                     detail_level = detail_level, other_levels = other_levels, active = active)

# =============================================================================
#
# =============================================================================

# LLM will call this when it wants to format the data for plotting a table or chart.

def plotting_api(data, plot_type = 'table', metric_cols = ['sales'],
                 product_col = 'product-name', location_col = 'location-name'):

    df = pd.DataFrame(data)

    if 'week-no' in df.columns:
        time_col = 'week-no'
    elif 'period-no' in df.columns:
        time_col = 'period-no'
    elif 'quarter-no' in df.columns:
        time_col = 'quarter-no'
    else:
        time_col = 'cal-year'

    time_cols = ['cal-year', time_col] if time_col != 'cal-year' else ['cal-year']

    if isinstance(metric_cols, str):
        metric_cols = [word.strip() for word in metric_cols.split(',')]

    id_cols = time_cols + [product_col] + [location_col]
    all_cols = time_cols + [product_col] + [location_col] + metric_cols

    df = df[all_cols]
    df = df.drop_duplicates(subset=id_cols)

    unique_times = [int(x) for x in df[time_cols].drop_duplicates().loc[:, time_col].values]
    products = sorted(df[product_col].unique())
    locations = sorted(df[location_col].unique())

    if plot_type == 'line':
        options = { 'xaxis' : { time_col : unique_times}}

        if len(products) > 1 and  len(locations) == len(metric_cols) == 1:
            df_wide = df.pivot(index = time_cols, columns = product_col, values = metric_cols).fillna(0).round(2)
            cols = [f'{col[1]}' for col in df_wide.columns]
            df_wide.columns = cols

        elif len(locations) > 1 and  len(products) == len(metric_cols) == 1:
            df_wide = df.pivot(index = time_cols, columns = location_col, values = metric_cols).fillna(0).round(2)
            cols = [f'{col[1]}' for col in df_wide.columns]
            df_wide.columns = cols

        elif len(metric_cols) >= 1 and len(products) == len(locations) == 1:
            df_wide = df[metric_cols].fillna(0).round(2)
            cols = df_wide.columns

        else:
            df_wide = df.groupby(time_cols)[metric_cols].sum().fillna(0).round(2)
            cols = df_wide.columns

        series = [{'name': col, 'data': df_wide[col].tolist()} for col in cols]
        res = {'type': plot_type,  'options' : options, 'series' : series}

    elif plot_type == 'bar':

        if len(products) > 1:
            options = { 'xaxis' : { product_col : products}}
            df_wide = df.groupby(product_col)[metric_cols].sum().fillna(0).round(2)
            cols = df_wide.columns

        elif len(locations) > 1:
            options = { 'xaxis' : { location_col : locations}}
            df_wide = df.groupby(location_col)[metric_cols].sum().fillna(0).round(2)
            cols = df_wide.columns

        else:
            options = { 'xaxis' : { time_col : unique_times}}
            df_wide = df[metric_cols].fillna(0).round(2)
            cols = df_wide.columns

        series = [{'name': col, 'data': df_wide[col].tolist()} for col in cols]
        res = {'type': plot_type,  'options' : options, 'series' : series}

    elif plot_type == 'table':
        df.columns = [x.replace('-', ' ').title() for x in df.columns]
        tableData1 = df.fillna(0).round(2).to_dict(orient = 'records')
        res = {'type': plot_type,  'tableData1' : tableData1}

    else:
        raise ValueError('chart_type')

    return res


import unittest
import sys, os

sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
from src.models import *
from src.algo.greedy import *
from src.algo.RLF import *
from src.generator import *
from src.algo.contraction import *
from utils import *

# Abstract Class
class GenericColoringTest(unittest.TestCase):
    def algorithm(self, g: Graph):
        pass

    def test_empty_graph(self):

        self.assertEqual(self.algorithm(Graph()), 0)

    def test_random_graph(self):
        for i in range(1, 100):
            g = generate_graph(i)
            self.algorithm(g)
            self.assertTrue(is_corrected_colored(g))

class GreedyColoringTest(GenericColoringTest):
    def algorithm(self, g: Graph):
        return greedy_coloring(g)

class GreedySortingColoringTest(GenericColoringTest):
    def algorithm(self, g: Graph):
        return sorting_greedy_coloring(g)

class RLFTest(GenericColoringTest):
    def algorithm(self, g: Graph):
        return RLF(g)

    def test_max_neighbours(self):
        g = Graph()

        self.assertEqual(get_max_neighbours_vertex(g), None)

        g.addVertex(Vertex(1))
        g.addVertex(Vertex(2))
        g.addVertex(Vertex(3))
        g.addVertex(Vertex(4))
        g.addVertex(Vertex(5))
        g.addEdge(1,2)

        self.assertEqual(get_max_neighbours_vertex(g), g.getVertex(1))

        g.addEdge(2,3)
        self.assertEqual(get_max_neighbours_vertex(g), g.getVertex(2))
        g.addEdge(3,4)
        g.addEdge(3,5)
        self.assertEqual(get_max_neighbours_vertex(g), g.getVertex(3))

        g.removeVertex(3)
        self.assertEqual(get_max_neighbours_vertex(g), g.getVertex(1))

    def test_not_adjecent_to_all(self):
        g = Graph()
        g.addVertex(Vertex(1))
        g.addVertex(Vertex(2))
        g.addVertex(Vertex(3))
        g.addVertex(Vertex(4))
        g.addVertex(Vertex(5))
        g.addEdge(1,2)
        g.addEdge(2,3)
        g.addEdge(3,4)
        g.addEdge(4,5)
        g.addEdge(5,1)
        g.addEdge(5,2)
        g.addEdge(5,3)
        g.addEdge(5,4)

        self.assertTrue(not_adjecent_to_all(g.getVertex(1), []))
        self.assertFalse(not_adjecent_to_all(g.getVertex(1), [g.getVertex(1)]))

        self.assertFalse(not_adjecent_to_all(g.getVertex(1), [g.getVertex(2)]))
        self.assertFalse(not_adjecent_to_all(g.getVertex(1), [g.getVertex(2), g.getVertex(3)]))
        self.assertTrue(not_adjecent_to_all(g.getVertex(1), [g.getVertex(3), g.getVertex(4)]))
        self.assertFalse(not_adjecent_to_all(g.getVertex(1), [g.getVertex(2), g.getVertex(3), g.getVertex(4)]))
        self.assertFalse(not_adjecent_to_all(g.getVertex(1), [g.getVertex(2), g.getVertex(3), g.getVertex(4), g.getVertex(5)]))
        self.assertFalse(not_adjecent_to_all(g.getVertex(5), [g.getVertex(3), g.getVertex(4), g.getVertex(5)]))
        self.assertTrue(not_adjecent_to_all(g.getVertex(2), [g.getVertex(4)]))

class ContractionTest(unittest.TestCase):
    def test_K3(self):
        g = Graph()
        g.addVertex(Vertex(1))
        g.addVertex(Vertex(2))
        g.addVertex(Vertex(3))

        g.addEdge(1, 2)
        g.addEdge(1, 3)
        g.addEdge(2, 1)
        g.addEdge(2, 3)
        g.addEdge(3, 1)
        g.addEdge(3, 2)
        colors = 3
        K3_polynom = lambda t: t * (t - 1) * (t - 2)
        self.assertEqual(contraction(g, colors)[0], K3_polynom(colors))
        self.assertEqual(contraction(g, colors)[1], 3)

    def test_Pn(self):
        for i in range(1, 12):
            g = Graph()
            for j in range(i):
                g.addVertex(Vertex(j + 1))
            vertices = g.getVertices()
            for k in range(len(vertices) - 1):
                g.addEdge(vertices[k].id, vertices[k + 1].id)

            Pn_polynom = lambda t, n: t * (t - 1) ** (n - 1)
            for c in range(10):
                self.assertEqual(contraction(g, colors=c)[0], Pn_polynom(c, len(g)))

    def test_Cn(self):
        for i in range(1, 12):
            g = Graph()
            for j in range(i):
                g.addVertex(Vertex(j + 1))
            vertices = g.getVertices()
            for k in range(len(vertices) - 1):
                g.addEdge(vertices[k].id, vertices[k + 1].id)
            g.addEdge(vertices[-1].id, vertices[0].id)

            Cn_polynom = lambda t, n: (t - 1)**n + ( (-1)**n ) * (t - 1)
            for c in range(10):
                self.assertEqual(contraction(g, colors=c)[0], Cn_polynom(c, len(g)))

    def test_Petersen_graph(self):
        g = Graph()
        for i in range(1, 11):
            g.addVertex(Vertex(i))

        # outer radius
        g.addEdge(1, 2)
        g.addEdge(1, 5)
        g.addEdge(2, 1)
        g.addEdge(2, 3)
        g.addEdge(3, 2)
        g.addEdge(3, 4)
        g.addEdge(4, 3)
        g.addEdge(4, 5)
        g.addEdge(5, 4)
        g.addEdge(5, 1)

        # star bounds
        g.addEdge(1, 6)
        g.addEdge(2, 7)
        g.addEdge(3, 8)
        g.addEdge(4, 9)
        g.addEdge(5, 10)

        #star iner edges
        g.addEdge(6, 8)
        g.addEdge(6, 9)
        g.addEdge(7, 9)
        g.addEdge(7, 10)
        g.addEdge(8, 6)
        g.addEdge(8, 10)
        g.addEdge(9, 6)
        g.addEdge(9, 7)
        g.addEdge(10, 7)
        g.addEdge(10, 8)

        Petersen_polynom = lambda t: t * (t - 1) * (t - 2) * (t**7 - 12 * t**6 + 67 * t**5 - 230 * t**4 + 529 * t**3 - 814 * t**2 + 775 * t - 352)
        for t in range(10):
            self.assertEqual(contraction(g, colors=t)[0], Petersen_polynom(t))
            self.assertEqual(contraction(g, colors=t)[1], 3)

del GenericColoringTest

if __name__ == '__main__':
    unittest.main()


import requests
import time
import json
import dateutil.parser as dateutilparser
from bs4 import BeautifulSoup
from dotenv import load_dotenv
from notify import notify

load_dotenv()

forums = ["https://forums.redflagdeals.com/hot-deals-f9/?st=0&rfd_sk=tt&sd=d"]
FILE_NAME = 'cache-rfd.json'

def form_full_rfd_url(relative_path):
    return 'https://forums.redflagdeals.com' + relative_path

try:
    with open(FILE_NAME, 'r') as f:
        cache = json.load(f)
except:
    cache = {}

try:
    for forum in forums:
        html_text = requests.get(forum).text
        soup = BeautifulSoup(html_text, 'html.parser')
        soup = soup.select('ul.topiclist.topics.with_categories')[0]

        thread_tags = soup.select('li.row.topic')

        for thread_tag in thread_tags:
            # Ignore sticky threads
            if thread_tag.find(class_='sticky'):
                continue

            id = thread_tag['data-thread-id']
            if id in cache:
                continue

            try:
                # Extract publish time
                post_time = dateutilparser.parse(thread_tag.select('span.first-post-time')[0].text).timestamp()
                time_str = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(post_time))

                # Thread link
                title_link_tag = thread_tag.select('a.topic_title_link')[0]
                link = form_full_rfd_url(title_link_tag['href'])

                title = thread_tag.select('h3.topictitle')[0].text.strip().replace('\n', '')

                message = f"{time_str}: {title}"
                message += "\n\n"
                message += link
                message += "\n\n"

                notify(message)
                cache[id] = post_time
            except:
                pass
finally:
    with open(FILE_NAME, 'w') as outfile:
        json.dump(cache, outfile)


import copy
import time

import mock as _mock
import pytest
from django.conf import settings

from alarm_backends.core.alert import Alert
from alarm_backends.core.cache.action_config import ActionConfigCacheManager
from alarm_backends.core.cache.key import ALERT_SNAPSHOT_KEY
from alarm_backends.service.alert.builder.processor import AlertBuilder
from alarm_backends.service.alert.manager.checker.upgrade import UpgradeChecker
from alarm_backends.service.fta_action.tasks import create_actions
from alarm_backends.tests.service.access.data.config import (
    STRATEGY_CONFIG_V3,
    USER_GROUP_DATA,
    USER_GROUP_WXBOT_DATA,
)
from constants.action import ActionPluginType, AssignMode
from constants.alert import EventStatus

pytestmark = pytest.mark.django_db

from alarm_backends.core.cache.assign import AssignCacheManager
from alarm_backends.service.fta_action.tasks.alert_assign import (
    AlertAssigneeManager,
    AssignRuleMatch,
    BackendAssignMatchManager,
)
from api.cmdb.define import Business, Host
from bkmonitor.documents import AlertDocument, AlertLog, EventDocument
from bkmonitor.models import (
    ActionConfig,
    ActionInstance,
    ActionPlugin,
    AlertAssignGroup,
    AlertAssignRule,
    ConvergeInstance,
    ConvergeRelation,
    DutyArrange,
    UserGroup,
)

_mock.patch("alarm_backends.service.fta_action.utils.run_converge.delay", return_value=11111).start()
_mock.patch("alarm_backends.service.fta_action.tasks.run_action.apply_async", return_value=11111).start()
_mock.patch("alarm_backends.service.fta_action.tasks.run_webhook_action.apply_async", return_value=11111).start()
_mock.patch("alarm_backends.service.fta_action.tasks.run_action.delay", return_value=11111).start()
_mock.patch("alarm_backends.service.converge.tasks.run_converge.apply_async", return_value=11111).start()
_mock.patch("alarm_backends.service.fta_action.tasks.run_noise_reduce_task.apply_async", return_value=11111).start()

@pytest.fixture()
def biz_mock(mocker):
    return mocker.patch(
        "alarm_backends.core.cache.cmdb.business.BusinessManager.all", return_value=[Business(bk_biz_id=2)]
    )

@pytest.fixture()
def host_mock(mocker):
    return mocker.patch(
        "alarm_backends.core.cache.cmdb.host.HostManager.get",
        return_value=Host(
            bk_host_id=1,
            bk_module_ids=[],
            bk_attr_test="123",
            bk_set_ids=[],
            bk_cloud_id=0,
            bk_biz_id=2,
            topo_link={},
            operator=[],
            bak_operator=[],
            bk_host_innerip="127.0.0.1",
        ),
    )

@pytest.fixture()
def save_alert_snapshot(mocker):
    return mocker.patch("alarm_backends.core.alert.AlertCache.save_alert_snapshot", return_value=(1, 2))

@pytest.fixture()
def save_alert_to_cache(mocker):
    return mocker.patch("alarm_backends.core.alert.AlertCache.save_alert_to_cache", return_value=1)

@pytest.fixture()
def setup():
    ActionInstance.objects.all().delete()
    AlertAssignGroup.objects.all().delete()
    AlertAssignRule.objects.all().delete()
    ConvergeRelation.objects.all().delete()

    assign_group = AlertAssignGroup.objects.create(name="test cache", bk_biz_id=2, priority=1)
    rule = {
        "assign_group_id": assign_group.id,
        "user_groups": [1],
        "conditions": [],
        "actions": [
            {
                "action_type": ActionPluginType.NOTICE,
                "is_enabled": True,
                "upgrade_config": {"is_enabled": True, "user_groups": [2, 1], "upgrade_interval": 30},
            },
            {"action_type": ActionPluginType.ITSM, "action_id": 4444},
        ],
        "alert_severity": 2,
        "additional_tags": [{"key": "ip123", "value": "127.0.0.1"}],
        "bk_biz_id": 2,
        "is_enabled": True,
    }
    yield AlertAssignRule.objects.create(**rule)

@pytest.fixture()
def condition_rule():
    ActionInstance.objects.all().delete()
    AlertAssignGroup.objects.all().delete()
    AlertAssignRule.objects.all().delete()
    ConvergeRelation.objects.all().delete()

    assign_group = AlertAssignGroup.objects.create(name="test cache", bk_biz_id=2, priority=1)
    rule = {
        "assign_group_id": assign_group.id,
        "user_groups": [3],
        "conditions": [
            {
                "field": "ip",
                "value": "127.0.0.1",
                "method": "eq",
            }
        ],
        "actions": [
            {
                "action_type": ActionPluginType.NOTICE,
                "is_enabled": True,
                "upgrade_config": {"is_enabled": True, "user_groups": [2, 1], "upgrade_interval": 30},
            },
            {"action_type": ActionPluginType.ITSM, "action_id": 4444},
        ],
        "alert_severity": 1,
        "additional_tags": [{"key": "ip123", "value": "127.0.0.1"}],
        "bk_biz_id": 2,
        "is_enabled": True,
    }
    yield AlertAssignRule.objects.create(**rule)

@pytest.fixture()
def empty_setup():
    ActionInstance.objects.all().delete()
    AlertAssignGroup.objects.all().delete()
    AlertAssignRule.objects.all().delete()
    ConvergeRelation.objects.all().delete()

    assign_group = AlertAssignGroup.objects.create(name="test cache", bk_biz_id=2, priority=1)
    rule = {
        "assign_group_id": assign_group.id,
        "user_groups": [1],
        "conditions": [
            {
                "field": "is_empty_users",
                "value": "true",
                "method": "eq",
            }
        ],
        "actions": [
            {
                "action_type": ActionPluginType.NOTICE,
                "is_enabled": True,
                "upgrade_config": {"is_enabled": True, "user_groups": [2, 1], "upgrade_interval": 30},
            },
            {"action_type": ActionPluginType.ITSM, "action_id": 4444},
        ],
        "alert_severity": 2,
        "additional_tags": [{"key": "ip123", "value": "127.0.0.1"}],
        "bk_biz_id": 2,
        "is_enabled": True,
    }
    yield AlertAssignRule.objects.create(**rule)

    AlertAssignGroup.objects.all().delete()
    AlertAssignRule.objects.all().delete()
    ActionInstance.objects.all().delete()
    ConvergeInstance.objects.all().delete()
    ConvergeRelation.objects.all().delete()

def clear_index():
    for doc in [AlertDocument, AlertLog]:
        ilm = doc.get_lifecycle_manager()
        ilm.es_client.indices.delete(index=doc.Index.name)

@pytest.fixture()
def alert():
    yield AlertDocument(
        **{
            "id": "12345",
            "alert_name": "test assign",
            "dedupe_md5": "68e9f0598d72a4b6de2675d491e5b922",
            "duration": 60 * 60,
            "severity": 3,
            "strategy_id": 1,
            "status": EventStatus.ABNORMAL,
            "end_time": None,
            "create_time": 1617504052,
            "begin_time": 1617504052,
            "first_anomaly_time": 1617504052,
            "latest_time": 1617504052,
            "dimensions": [{"key": "bk_target_ip", "value": "127.0.0.1"}],
            "event": EventDocument(
                **{
                    "tags": [{"key": "target", "value": "127.0.0.1"}],
                    "metric": ["123"],
                    "ip": "127.0.0.1",
                    "bk_cloud_id": 0,
                    "bk_biz_id": 2,
                }
            ),
            "extra_info": {"strategy": get_strategy_dict()},
        }
    )

@pytest.fixture()
def user_group_setup():
    UserGroup.objects.all().delete()
    DutyArrange.objects.all().delete()
    duty_arranges = [
        DutyArrange(
            **{
                "user_group_id": 1,
                "order": 1,
                "users": [{"id": "lisa", "display_name": "lisa", "logo": "", "type": "user"}],
            }
        ),
        DutyArrange(
            **{
                "user_group_id": 2,
                "order": 1,
                "users": [{"id": "admin", "display_name": "管理员", "logo": "", "type": "user"}],
            }
        ),
        DutyArrange(
            **{
                "user_group_id": 3,
                "order": 1,
                "users": [
                    {"id": "lisa", "display_name": "lisa", "logo": "", "type": "user"},
                    {"id": "lisa1", "display_name": "lisa1", "logo": "", "type": "user"},
                ],
            }
        ),
    ]
    DutyArrange.objects.bulk_create(duty_arranges)
    group_data = copy.deepcopy(USER_GROUP_DATA)
    for group_id in [1, 2, 3]:
        group_data["id"] = group_id
        user_groups = UserGroup.objects.create(**group_data)
    yield user_groups
    UserGroup.objects.all().delete()
    DutyArrange.objects.all().delete()

@pytest.fixture()
def wxbot_user_group_setup():
    UserGroup.objects.all().delete()
    DutyArrange.objects.all().delete()
    duty_arranges = [
        DutyArrange(
            **{
                "user_group_id": 1,
                "order": 1,
                "users": [{"id": "lisa", "display_name": "lisa", "logo": "", "type": "user"}],
            }
        ),
        DutyArrange(
            **{
                "user_group_id": 2,
                "order": 1,
                "users": [{"id": "admin", "display_name": "管理员", "logo": "", "type": "user"}],
            }
        ),
    ]
    DutyArrange.objects.bulk_create(duty_arranges)
    group_data = copy.deepcopy(USER_GROUP_WXBOT_DATA)
    for group_id in [1, 2]:
        group_data["id"] = group_id
        user_groups = UserGroup.objects.create(**group_data)
    yield user_groups

def init_action_plugin():
    ActionPlugin.objects.all().delete()
    ActionConfig.objects.all().delete()
    ActionPlugin.objects.bulk_create(
        [
            ActionPlugin(
                **{
                    "id": 1,
                    "name": "通知告警",
                    "is_enabled": True,
                    "is_builtin": True,
                    "is_deleted": False,
                    "is_peripheral": False,
                    "plugin_type": "notice",
                    "plugin_key": "notice",
                    "plugin_source": "builtin",
                    "description": "告警通知是平台内置的套餐类型，由平台自身实现。可以对告警信息基于人进行收敛，可以对接不同的告警通知渠道。 "
                    "\n\n* 基于人进行收敛\n* 有告警风暴控制能力\n* 可以定制不同的告警模版\n* 内置基于不同的通知渠道显示的变量\n* "
                    "可以自定义各种通知渠道[查看文档]()\n\n更多[查看文档]()",
                    "config_schema": {
                        "content_template": "发送{{notice_way_display}}告警通知给{{notice_receiver}}{{status_display}}",
                        "content_template_with_url": "达到通知告警的执行条件【{{action_signal}}】，已触发告警通知",
                        "content_template_without_assignee": "达到通知告警的执行条件【{{action_signal}}】，当前通知人员为空",
                        "content_template_shielded": "达到通知告警的执行条件【{{action_signal}}】，因告警已被全局屏蔽或不在通知时间段，忽略通知发送",
                        "content_template_shielded_with_url": "达到通知告警的执行条件【{{action_signal}}】，因告警已被屏蔽忽略通知发送，点击$查看屏蔽策略$",
                    },
                    "backend_config": [{"function": "execute_notify", "name": "发送通知"}],
                }
            ),
            ActionPlugin(
                **{
                    "id": 5,
                    "name": "流程服务",
                    "is_builtin": True,
                    "is_deleted": False,
                    "is_peripheral": True,
                    "plugin_source": "peripheral",
                    "plugin_type": "common",
                    "plugin_key": "itsm",
                    "has_child": True,
                    "description": "",
                    "config_schema": {
                        "template": {
                            "resource_class": "CommonBaseResource",
                            "name": "服务名称",
                            "init_kwargs": {"url": "{{bk_paas_inner_host}}/api/c/compapi/v2/itsm/get_services/"},
                            "request_data_mapping": {"display_role": "BK_FTA"},
                            "resource_module": "api.common.default",
                            "resource_data": "response[*]",
                            "mapping": {
                                "id": "{{id}}",
                                "name": "{{name}}",
                                "url": "{{itsm_site_url}}/#/ticket/create?service_id={{id}}",
                            },
                        },
                        "plugin_url": {
                            "mapping": {
                                "url": "{{itsm_site_url}}/#/project/service/new/basic?project_id=0",
                                "tips": "前往流程服务",
                            }
                        },
                        "detail": {
                            "name": "提单信息",
                            "resource_class": "CommonBaseResource",
                            "resource_module": "api.common.default",
                            "init_kwargs": {"url": "{{bk_paas_inner_host}}/api/c/compapi/v2/itsm/get_service_detail/"},
                            "request_data_mapping": {"service_id": "{{template_id}}"},
                            "resource_data": "response.fields[*]",
                            "mapping": {
                                "key": "{{key}}",
                                "name": "{{name}}",
                                "value": "{{default or ''}}",
                                "placeholder": "{{desc}}",
                            },
                        },
                        "content_template_with_url": "故障工单套餐【{{action_name}}】成功创建单据[{{sn}}]，点击$查看工单详情$",
                        "content_template": "故障工单套餐【{{action_name}}】处理{{status_display}}, 请关注！！",
                    },
                    "backend_config": [
                        {
                            "function": "create_task",
                            "name": "创建单据",
                            "resource_class": "CommonBaseResource",
                            "resource_module": "api.common.default",
                            "init_kwargs": {
                                "url": "{{bk_paas_inner_host}}/api/c/compapi/v2/itsm/create_ticket/",
                                "method": "POST",
                            },
                            "inputs": [
                                {
                                    "key": "service_id",
                                    "value": "execute_config.template_id",
                                    "type": "int",
                                    "format": "jmespath",
                                },
                                {
                                    "key": "fields",
                                    "value": "execute_config.template_detail[].{key: key, value: value}",
                                    "type": "list",
                                    "format": "jmespath",
                                },
                                {"key": "creator", "value": "operator", "type": "string", "format": "jmespath"},
                            ],
                            "outputs": [
                                {"key": "sn", "value": "response.sn", "format": "jmespath"},
                                {"key": "id", "value": "response.id", "format": "jmespath"},
                                {
                                    "key": "url",
                                    "value": "{{itsm_site_url}}#/ticket/detail?id={{id}}",
                                    "format": "jinja2",
                                },
                            ],
                            "need_insert_log": True,
                            "log_template": "流程服务套餐【{{action_name}}】已成功创建工单[{{sn}}]，点击$查看工单详情$",
                        }
                    ],
                }
            ),
        ]
    )
    notice_action_config = {
        "execute_config": {
            "template_detail": {
                "interval_notify_mode": "standard",  # 间隔模式
                "notify_interval": 7200,  # 通知间隔
                "template": [  # 通知模板配置
                    {
                        "signal": "abnormal",
                    }
                ],
            }
        },
        "id": 55555,
        "plugin_id": 1,
        "is_enabled": True,
        "bk_biz_id": 2,
        "name": "test_notice",
    }

    itsm_config = {
        "execute_config": {
            "template_id": 1000043,
            "template_detail": {"title": "123", "name": "test itsm"},
            "timeout": 60,
        },
        "name": "test itsm",
        "desc": "这是描述，这是描述",
        "is_enabled": True,
        "plugin_id": 5,
        "bk_biz_id": 2,
        "id": 4444,
    }

    ActionConfig.objects.bulk_create([ActionConfig(**notice_action_config), ActionConfig(**itsm_config)])

class TestAlertAssignRule:
    def test_assign_cache(self, biz_mock, setup):
        AssignCacheManager.refresh()
        rule = AlertAssignRule.objects.filter(id=setup.id).values()[0]
        rule["group_name"] = "test cache"
        assert biz_mock.call_count == 1
        assert AssignCacheManager.get_assign_priority_by_biz_id(2), [1]
        assert AssignCacheManager.get_assign_groups_by_priority(2, 1) == {setup.assign_group_id}
        assert AssignCacheManager.get_assign_rules_by_group(2, setup.assign_group_id) == [rule]

    def test_host_cmdb_dimension_matched(self, alert, host_mock):
        rule = {
            "conditions": [
                {
                    "field": "host.bk_cloud_id",
                    "value": "0",
                    "method": "eq",
                },
                {
                    "field": "host.bk_attr_test",
                    "value": "123",
                    "method": "eq",
                },
            ]
        }

        match_obj = AssignRuleMatch(rule, alert=alert)
        m = BackendAssignMatchManager(alert=alert)
        dimensions = m.get_match_dimensions()
        assert match_obj.is_matched(dimensions)
        assert host_mock.call_count == 1

    def test_assign_tags_dimension_not_matched(self, alert):
        rule = {
            "conditions": [
                {
                    "field": "tags.target_1",
                    "value": "127.0.0.1",
                    "method": "eq",
                }
            ]
        }
        match_obj = AssignRuleMatch(rule, alert=alert)
        m = BackendAssignMatchManager(alert=alert)
        dimensions = m.get_match_dimensions()
        assert match_obj.is_matched(dimensions) is False

    def test_assign_alert_name_matched(self, alert):
        rule = {
            "conditions": [
                {
                    "field": "alert.name",
                    "value": "test assign",
                    "method": "eq",
                }
            ]
        }
        match_obj = AssignRuleMatch(rule, alert=alert)
        m = BackendAssignMatchManager(alert=alert)
        dimensions = m.get_match_dimensions()
        assert match_obj.is_matched(dimensions)

    def test_assign_alert_metric_matched(self, alert):
        rule = {
            "conditions": [
                {
                    "field": "alert.metric",
                    "value": ["123"],
                    "method": "eq",
                }
            ]
        }
        match_obj = AssignRuleMatch(rule, alert=alert)
        m = BackendAssignMatchManager(alert=alert)
        dimensions = m.get_match_dimensions()
        assert match_obj.is_matched(dimensions)

    def test_assign_or_condition_matched(self, alert):
        rule = {
            "conditions": [
                {
                    "field": "tags.target",
                    "value": "127.0.0.1",
                    "method": "eq",
                },
                {"field": "tags.target", "value": "127.0.0.2", "method": "eq", "condition": "or"},
            ]
        }
        match_obj = AssignRuleMatch(rule, alert=alert)
        m = BackendAssignMatchManager(alert=alert)
        assert match_obj.is_matched(m.dimensions) is True

    def test_assign_composite_condition_matched(self, alert):
        rule = {
            "conditions": [
                {
                    "field": "tags.target",
                    "value": "127.0.0.2",
                    "method": "eq",
                },
                {"field": "tags.target", "value": "127.0.0.1", "method": "eq", "condition": "or"},
                {"field": "alert_name", "value": "test", "method": "eq", "condition": "and"},
            ]
        }
        match_obj = AssignRuleMatch(rule, alert=alert)
        m = BackendAssignMatchManager(alert=alert)
        assert match_obj.is_matched(m.dimensions) is False

    def test_ip_matched(self, alert):
        rule = {
            "conditions": [
                {
                    "field": "bk_target_ip",
                    "value": "127.0.0.1",
                    "method": "eq",
                }
            ]
        }
        match_obj = AssignRuleMatch(rule, alert=alert)
        m = BackendAssignMatchManager(alert=alert)
        assert match_obj.is_matched(m.dimensions) is True

    def test_host_attr_matched(self, alert):
        rule = {
            "conditions": [
                {
                    "field": "host.bk_target_ip",
                    "value": ["127.0.0.1"],
                    "method": "eq",
                }
            ]
        }
        match_obj = AssignRuleMatch(rule, alert=alert)
        assert match_obj.is_matched({"host.bk_target_ip": ["127.0.0.1", "127.0.0.2"]}) is True

        rule = {
            "conditions": [
                {
                    "field": "host.bk_host_id",
                    "value": ["1"],
                    "method": "eq",
                }
            ]
        }
        match_obj = AssignRuleMatch(rule, alert=alert)
        assert match_obj.is_matched({"host.bk_host_id": [1, 2]}) is True

    def test_is_changed(self):
        rule = {
            "id": 1,
            "user_groups": [1, 2],
            "conditions": [
                {
                    "field": "bk_target_ip",
                    "value": "127.0.0.1",
                    "method": "eq",
                }
            ],
        }
        snap_rule = copy.deepcopy(rule)
        # no_changed
        snap_rule["user_groups"] = [2, 1]
        assert AssignRuleMatch(rule, assign_rule_snap=snap_rule).is_changed is False

        # new
        assert AssignRuleMatch(rule).is_changed
        assert AssignRuleMatch(rule).is_new

        # changed
        snap_rule["user_groups"] = [1]
        assert AssignRuleMatch(rule, assign_rule_snap=snap_rule).is_changed

    def test_is_upgrade(self, alert):
        rule = {
            "id": 1,
            "user_groups": [1, 2],
            "conditions": [],
            "actions": [
                {
                    "action_type": ActionPluginType.NOTICE,
                    "is_enabled": True,
                    "upgrade_config": {"is_enabled": True, "user_groups": [3, 4, 5], "upgrade_interval": 30},
                }
            ],
        }

        assert AssignRuleMatch(rule, assign_rule_snap=rule).upgrade_rule.is_upgrade_enable
        assert AssignRuleMatch(rule, assign_rule_snap=rule).is_changed is False
        assert AssignRuleMatch(rule, assign_rule_snap=rule).get_upgrade_user_group() == []

    def test_empty_notice_group_true(self, alert):
        rule = {
            "conditions": [
                {
                    "field": "is_empty_users",
                    "value": "true",
                    "method": "eq",
                }
            ]
        }
        rule_obj = AssignRuleMatch(rule)
        m = BackendAssignMatchManager(alert=alert)
        assert rule_obj.is_matched(m.get_match_dimensions())

    def test_empty_notice_group_false_with_chat_id(self, alert, wxbot_user_group_setup):
        rule = {
            "conditions": [
                {
                    "field": "is_empty_users",
                    "value": "false",
                    "method": "eq",
                }
            ]
        }
        rule_obj = AssignRuleMatch(rule)
        m = AlertAssigneeManager(
            alert=alert,
            notice_user_groups=[wxbot_user_group_setup.id],
            assign_mode=[AssignMode.BY_RULE, AssignMode.ONLY_NOTICE],
        )
        assert rule_obj.is_matched(m.match_manager.get_match_dimensions())

    def test_empty_value_false(self, alert):
        rule = {
            "conditions": [
                {
                    "field": "ip",
                    "value": [""],
                    "method": "eq",
                }
            ]
        }
        rule_obj = AssignRuleMatch(rule)
        m = BackendAssignMatchManager(alert=alert)
        assert not rule_obj.is_matched(m.get_match_dimensions())

    def test_empty_value_true(self, alert):
        rule = {
            "conditions": [
                {
                    "field": "ip",
                    "value": [""],
                    "method": "eq",
                }
            ]
        }
        rule_obj = AssignRuleMatch(rule)
        alert.event.ip = ""
        m = BackendAssignMatchManager(alert=alert)

        assert rule_obj.is_matched(m.get_match_dimensions())

    def test_empty_notice_group_false(self, alert):
        rule = {
            "conditions": [
                {
                    "field": "is_empty_users",
                    "value": "true",
                    "method": "eq",
                }
            ]
        }
        rule_obj = AssignRuleMatch(rule)
        m = BackendAssignMatchManager(alert=alert)
        assert rule_obj.is_matched(m.get_match_dimensions())

        m = BackendAssignMatchManager(alert=alert, notice_users=["admin"])
        assert rule_obj.is_matched(m.get_match_dimensions()) is False

def get_strategy_dict():
    strategy_dict = copy.deepcopy(STRATEGY_CONFIG_V3)
    strategy_dict["notice"]["options"]["assign_mode"] = [AssignMode.BY_RULE]
    strategy_dict["notice"]["options"]["noise_reduce_config"] = {}
    return strategy_dict

@pytest.fixture()
def init_configs():
    clear_index()
    init_action_plugin()
    ALERT_SNAPSHOT_KEY.client.flushall()
    AssignCacheManager.refresh()
    ActionConfigCacheManager.refresh()

class TestAssignManager:
    def test_by_assign(self, alert):
        rule = {
            "id": 1,
            "user_groups": [1],
            "conditions": [],
            "actions": [
                {
                    "action_type": ActionPluginType.NOTICE,
                    "is_enabled": True,
                    "upgrade_config": {"is_enabled": True, "user_groups": [1], "upgrade_interval": 30},
                },
                {"action_type": ActionPluginType.ITSM, "action_id": 4444},
            ],
            "alert_severity": 2,
            "additional_tags": [{"key": "ip", "value": "127.0.0.1"}],
        }
        rule_obj = AssignRuleMatch(rule)
        assert rule_obj.itsm_action == {"action_type": ActionPluginType.ITSM, "action_id": 4444}
        assert rule_obj.alert_severity == 2
        assert rule_obj.additional_tags == [{"key": "ip", "value": "127.0.0.1"}]

    def test_change_alert_severity(self, setup, alert, user_group_setup, biz_mock, init_configs):
        assert ActionConfigCacheManager.get_action_config_by_id(4444)

        AlertDocument.bulk_create([alert])
        assert biz_mock.call_count == 1
        actions0 = create_actions(0, "abnormal", alerts=[alert])

        new_alert = AlertDocument.get(id=alert.id)
        assert len(actions0) == 3
        assert new_alert.severity == 2
        assert new_alert.extra_info.severity_source == AssignMode.BY_RULE

    def test_notice_upgrade(self, setup, alert, user_group_setup, biz_mock, init_configs):
        """
        测试告警分派的告警升级
        :param setup:
        :param alert:
        :param biz_mock:
        :param init_configs:
        :return:
        """
        alert.duration = 50 * 60
        alert.extra_info["rule_snaps"] = {
            setup.id: {"user_groups": setup.user_groups, "conditions": setup.conditions, "id": setup.id}
        }
        print(alert.extra_info["rule_snaps"])
        AlertDocument.bulk_create([alert])
        actions = create_actions(0, "abnormal", alerts=[alert], notice_type="upgrade")
        new_alert = AlertDocument.get(id=alert.id)
        assert len(actions) == 2
        p_ai = ActionInstance.objects.get(is_parent_action=True, id__in=actions)
        p_ai.inputs["notify_info"].pop("wxbot_mention_users", None)
        assert p_ai.inputs["notify_info"] == {'mail': ['admin']}
        assert new_alert.severity == 2
        assert new_alert.extra_info.severity_source == AssignMode.BY_RULE
        assert p_ai.get_content()["text"], '达到通知告警的执行条件【告警升级】，已出发告警'

        alert.extra_info = new_alert.extra_info
        alert.supervisor = new_alert.supervisor
        alert.extra_info["rule_snaps"][str(setup.id)]["last_upgrade_time"] = int(time.time()) - 31 * 60
        AlertDocument.bulk_create([alert], action="update")
        ActionInstance.objects.all().delete()
        new_actions = create_actions(0, "abnormal", alerts=[alert], notice_type="upgrade")
        assert len(new_actions) == 2
        p_ai = ActionInstance.objects.get(is_parent_action=True, id__in=new_actions)
        assert set(p_ai.inputs["notify_info"]["mail"]) == {"lisa"}
        assert p_ai.inputs["notice_type"] == "upgrade"
        assert p_ai.need_poll is False

        new_alert = AlertDocument.get(id=alert.id)
        assert len(new_alert.supervisor) == 2

    def test_notice_reupgrade(self, setup, alert, user_group_setup, biz_mock, init_configs):
        """
        分派规则下的再次升级测试
        :param setup:
        :param alert:
        :param biz_mock:
        :param init_configs:
        :return:
        """
        alert.duration = 50 * 60
        alert.extra_info["rule_snaps"] = {
            setup.id: {
                "user_groups": setup.user_groups,
                "conditions": setup.conditions,
                "id": setup.id,
                "last_group_index": 0,
                "last_upgrade_time": int(time.time()) - 31 * 60,
            }
        }
        AlertDocument.bulk_create([alert])
        assert biz_mock.call_count == 1
        actions = create_actions(0, "abnormal", alerts=[alert], notice_type="upgrade")
        new_alert = AlertDocument.get(id=alert.id)
        assert len(actions) == 2
        p_ai = ActionInstance.objects.get(is_parent_action=True, id__in=actions)
        p_ai.inputs["notify_info"].pop("wxbot_mention_users", None)
        assert p_ai.inputs["notify_info"] == {'mail': ['lisa']}
        assert new_alert.severity == 2
        assert new_alert.extra_info.severity_source == AssignMode.BY_RULE

        # 升级结束，后面不再继续发通知了
        alert.extra_info = new_alert.extra_info
        alert.appointee = new_alert.appointee
        AlertDocument.bulk_create([alert], action="update")
        ActionInstance.objects.all().delete()
        new_actions = create_actions(0, "abnormal", alerts=[alert], notice_type="upgrade")
        assert len(new_actions) == 0

    def test_normal_notice(self, setup, alert, user_group_setup, biz_mock, init_configs):
        """
        原生测试
        """
        alert.extra_info.strategy.notice["options"]["assign_mode"] = [AssignMode.ONLY_NOTICE]
        AlertDocument.bulk_create([alert])
        assert biz_mock.call_count == 1
        actions = create_actions(alert.extra_info.strategy["id"], "abnormal", alerts=[alert])
        assert len(actions) == 2
        p_ai = ActionInstance.objects.get(is_parent_action=True, id__in=actions)
        p_ai.inputs["notify_info"].pop("wxbot_mention_users", None)
        assert p_ai.inputs["notify_info"] == {'weixin': ['lisa']}
        assert alert.severity == 3

    def test_ignore_origin_notice(self, condition_rule, alert, user_group_setup, biz_mock, init_configs):
        """
        测试适配到告警条件之后原来的告警不会产生
        """
        alert.extra_info.strategy.notice["options"]["assign_mode"] = [AssignMode.ONLY_NOTICE, AssignMode.BY_RULE]
        alert.severity = 1
        alert.appointee = ['lisa', 'lisa1']
        AlertDocument.bulk_create([alert])
        assert biz_mock.call_count == 1
        actions = create_actions(0, "abnormal", alerts=[alert])
        assert len(actions) == 4
        p_ai = ActionInstance.objects.get(is_parent_action=True, id__in=actions)
        p_ai.inputs["notify_info"].pop("wxbot_mention_users", None)
        assert p_ai.inputs["notify_info"]['voice'] == [['lisa', 'lisa1']]
        new_alert = AlertDocument.get(id=alert.id)
        assert new_alert.appointee == ["lisa", "lisa1"]
        assert new_alert.severity == 1

    def test_default_assign_notice(self, setup, alert, user_group_setup, biz_mock, init_configs):
        """
        原生测试
        """
        alert.extra_info.strategy = {}
        AlertDocument.bulk_create([alert])
        assert biz_mock.call_count == 1
        actions = create_actions(0, "abnormal", alerts=[alert])
        assert len(actions) == 3
        p_ai = ActionInstance.objects.get(is_parent_action=True, id__in=actions)
        p_ai.inputs["notify_info"].pop("wxbot_mention_users", None)
        assert p_ai.inputs["notify_info"] == {'mail': ['lisa']}
        new_alert = AlertDocument.get(id=alert.id)
        assert new_alert.severity == 2
        assert new_alert.assign_tags == setup.additional_tags

    def test_default_assign_without_notice(self, alert, user_group_setup, biz_mock, init_configs):
        """
        原生测试
        """
        alert.extra_info.strategy = {}
        AlertDocument.bulk_create([alert])
        assert biz_mock.call_count == 1
        actions = create_actions(0, "abnormal", alerts=[alert])
        assert len(actions) == 0

    def test_enrich_rule_info(self, setup, alert, biz_mock, init_configs):
        processor = AlertBuilder()
        alert_obj = Alert(data=alert.to_dict())
        alert_obj._is_new = True
        enriched_alerts = processor.enrich_alerts([alert_obj])
        enriched_alert = enriched_alerts[0].to_document()
        assert enriched_alert.severity == 2
        assert alert_obj.severity_source == AssignMode.BY_RULE
        assert enriched_alert.assign_tags == setup.additional_tags
        assert len(enriched_alerts[0].top_event["tags"]) == 2

    def test_enrich_rule_info_by_empty_setup_false(self, empty_setup, alert, biz_mock, user_group_setup, init_configs):
        """
        告警通知人不为空适配
        """
        processor = AlertBuilder()
        alert_obj = Alert(data=alert.to_dict())
        alert_obj._is_new = True
        alert.strategy["notice"]["options"]["assign_mode"] = [AssignMode.ONLY_NOTICE, AssignMode.BY_RULE]
        alert.strategy.get("notice", {}).update(user_groups=[user_group_setup.id])
        enriched_alerts = processor.enrich_alerts([alert_obj])
        enriched_alert = enriched_alerts[0].to_document()
        assert enriched_alert.severity == 3
        assert not alert_obj.severity_source

    def test_enrich_rule_info_by_empty_wxbot_setup_false(
        self, empty_setup, alert, biz_mock, wxbot_user_group_setup, init_configs
    ):
        """
        告警通知人不为空适配(机器人)
        """
        processor = AlertBuilder()
        alert_obj = Alert(data=alert.to_dict())
        alert_obj._is_new = True
        # 测试通知人为空，首先需要有通知配置
        alert.strategy["notice"]["options"]["assign_mode"] = [AssignMode.ONLY_NOTICE, AssignMode.BY_RULE]
        alert.strategy.get("notice", {}).update(user_groups=[wxbot_user_group_setup.id])
        enriched_alerts = processor.enrich_alerts([alert_obj])
        enriched_alert = enriched_alerts[0].to_document()
        assert enriched_alert.severity == 3
        assert not alert_obj.severity_source

    def test_enrich_rule_info_by_empty_setup(self, empty_setup, alert, biz_mock, init_configs):
        """
        告警通知组为空
        """
        processor = AlertBuilder()
        alert_obj = Alert(data=alert.to_dict())
        alert_obj._is_new = True
        alert.strategy.get("notice", {}).update(user_groups=[])
        enriched_alerts = processor.enrich_alerts([alert_obj])
        enriched_alert = enriched_alerts[0].to_document()
        assert enriched_alert.severity == 2
        assert alert_obj.severity_source == AssignMode.BY_RULE

    def test_enrich_rule_info_default(self, setup, alert, biz_mock, init_configs):
        processor = AlertBuilder()
        alert.extra_info.strategy.notice["options"]["assign_mode"] = [AssignMode.BY_RULE]
        alert_obj = Alert(data=alert.to_dict())
        alert_obj._is_new = True
        enriched_alerts = processor.enrich_alerts([alert_obj])
        enriched_alert = enriched_alerts[0].to_document()
        assert enriched_alert.severity == 2
        assert alert_obj.severity_source == AssignMode.BY_RULE

    def test_enrich_rule_info_without_rule(self, setup, alert):
        processor = AlertBuilder()
        alert.extra_info.strategy.notice["options"]["assign_mode"] = [AssignMode.ONLY_NOTICE]
        alert_obj = Alert(data=alert.to_dict())
        alert_obj._is_new = True
        enriched_alerts = processor.enrich_alerts([alert_obj])
        enriched_alert = enriched_alerts[0].to_document()
        assert enriched_alert.severity == 3
        assert alert_obj.severity_source == ""

@pytest.fixture()
def create_actions_delay(mocker):
    return mocker.patch("alarm_backends.service.fta_action.tasks.create_actions.delay", return_value=(1, 2))

class TestUpgradeChecker:
    def test_origin_notice_upgrade(
        self, setup, alert, biz_mock, init_configs, create_actions_delay, save_alert_snapshot, save_alert_to_cache
    ):
        """
        测试告警升级
        :param setup:
        :param alert:
        :param biz_mock:
        :param init_configs:
        :return:
        """
        # 没有升级的情况
        alert.strategy["notice"]["options"]["assign_mode"] = [AssignMode.ONLY_NOTICE]
        alert.strategy["notice"]["options"].update(
            {"upgrade_config": {"is_enabled": True, "user_groups": [2, 1], "upgrade_interval": 30}}
        )
        AlertDocument.bulk_create([alert])
        # 没有升级的情况
        UpgradeChecker(alerts=[Alert(alert.to_dict())]).check_all()
        assert create_actions_delay.call_count == 0

        # 第一次升级
        alert.duration = 50 * 60
        alert.first_anomaly_time = int(time.time()) - alert.duration
        alert.latest_time = int(time.time())
        alert.extra_info["upgrade_notice"] = {}
        UpgradeChecker(alerts=[Alert(alert.to_dict())]).check_all()
        assert create_actions_delay.call_count == 1

        create_actions(alert.strategy_id, "abnormal", alerts=[alert], notice_type="upgrade")

        new_alert = AlertDocument.get(alert.id)
        assert new_alert.extra_info.upgrade_notice["last_group_index"] == 0

        # 第二次升级
        alert.extra_info = new_alert.extra_info
        alert.latest_time = int(time.time())
        alert.extra_info.upgrade_notice["last_upgrade_time"] = alert.latest_time - 30 * 60
        AlertDocument.bulk_create([alert], action="update")

        UpgradeChecker(alerts=[Alert(alert.to_dict())]).check_all()
        assert create_actions_delay.call_count == 2
        create_actions(alert.strategy_id, "abnormal", alerts=[alert], notice_type="upgrade")
        new_alert = AlertDocument.get(alert.id)
        assert new_alert.extra_info.upgrade_notice["last_group_index"] == 1

        # 时间间隔, 检测不会推送任务
        alert.extra_info = new_alert.extra_info
        alert.latest_time = int(time.time())
        alert.extra_info.upgrade_notice["last_upgrade_time"] = alert.latest_time - 30 * 60
        UpgradeChecker(alerts=[Alert(alert.to_dict())]).check_all()
        assert create_actions_delay.call_count == 2

        AlertDocument.bulk_create([alert], action="update")
        new_actions = create_actions(alert.strategy_id, "abnormal", alerts=[alert], notice_type="upgrade")
        assert len(new_actions) == 0

    def test_rule_upgrade(self, setup, alert, user_group_setup, biz_mock, init_configs, create_actions_delay):
        """
        再次升级测试
        :param setup:
        :param alert:
        :param biz_mock:
        :param init_configs:
        :return:
        """
        alert.duration = 50 * 60
        alert.extra_info["rule_snaps"] = {
            str(setup.id): {
                "user_groups": setup.user_groups,
                "conditions": setup.conditions,
                "id": setup.id,
                "last_group_index": 0,
                "last_upgrade_time": int(time.time()) - 31 * 60,
            }
        }
        AlertDocument.bulk_create([alert])
        assert biz_mock.call_count == 1
        # 再次升级
        UpgradeChecker(alerts=[Alert(alert.to_dict())]).check_all()
        assert create_actions_delay.call_count == 1
        create_actions(0, "abnormal", alerts=[alert], notice_type="upgrade")
        new_alert = AlertDocument.get(id=alert.id)
        assert new_alert.extra_info.rule_snaps[str(setup.id)]["last_group_index"] == 1

        # 不升级
        alert.extra_info = new_alert.extra_info
        alert.latest_time = int(time.time())
        alert.extra_info.rule_snaps[str(setup.id)]["last_upgrade_time"] = alert.latest_time - 30 * 60
        AlertDocument.bulk_create([alert], action="update")
        UpgradeChecker(alerts=[Alert(alert.to_dict())]).check_all()
        assert create_actions_delay.call_count == 1
        new_actions = create_actions(alert.strategy_id, "abnormal", alerts=[alert], notice_type="upgrade")
        assert len(new_actions) == 0

    def test_rule_upgrade_qos(self, setup, alert, user_group_setup, biz_mock, init_configs, create_actions_delay):
        """
        再次升级测试
        :param setup:
        :param alert:
        :param biz_mock:
        :param init_configs:
        :return:
        """
        alert.duration = 50 * 60
        alert.first_anomaly_time = int(time.time()) - alert.duration
        alert.latest_time = int(time.time())
        AlertDocument.bulk_create([alert])
        assert biz_mock.call_count == 1
        # 第一次升级
        UpgradeChecker(alerts=[Alert(alert.to_dict())]).check_all()
        assert create_actions_delay.call_count == 1
        create_actions(0, "abnormal", alerts=[alert], notice_type="upgrade")
        new_alert = AlertDocument.get(id=alert.id)
        assert new_alert.extra_info.rule_snaps[str(setup.id)]["last_group_index"] == 0

        # 第二次升级被QOS
        alert.extra_info = new_alert.extra_info
        alert.latest_time = int(time.time())
        alert.extra_info.rule_snaps[str(setup.id)]["last_upgrade_time"] = alert.latest_time - 30 * 60
        AlertDocument.bulk_create([alert], action="update")
        settings.QOS_DROP_ACTION_THRESHOLD = 1
        UpgradeChecker(alerts=[Alert(alert.to_dict())]).check_all()
        assert create_actions_delay.call_count == 2
        time.sleep(1)
        new_actions = create_actions(0, "abnormal", alerts=[alert], notice_type="upgrade")
        new_alert = AlertDocument.get(id=alert.id)
        assert new_alert.extra_info.rule_snaps[str(setup.id)]["last_group_index"] == 1
        assert len(new_actions) == 0
        log_search_object = (
            AlertLog.search(all_indices=True).filter("term", alert_id=alert.id).filter("term", op_type="ACTION")
        )
        assert len(log_search_object.execute().hits) == 2
        settings.QOS_DROP_ACTION_THRESHOLD = 100


#!/usr/bin/env python3

###
# Name: Aly Baughman, Cynthia Parks
# Student ID: 1923165, 2303535
# Email: baugh107@mail.chapman.edu, cparks@chapman.edu
# Course: CS510 Fall 2017
# Assignment: Classwork 5
###

"""Complex Plane Creation
    $ python
    >>> import cplane
    >>> help(cplane)
The ListComplexPlane class is meant to create a nested list grid of complex numbers.  It initializes with the min and max points
of both axis, their lengths, the plane of points, and an empty list of functions applied to said points.  It contains functions
to refresh and zoom in the grid, a function to create the plane itself, and a function to apply mathematical functions to the points.
"""

from abscplane import AbsComplexPlane

class ListComplexPlane(AbsComplexPlane):
    """This class creates a complex plane and uses different functions to edit the plane.
     Attributes:
        xmax (float) : maximum horizontal axis value
        xmin (float) : minimum horizontal axis value
        xlen (int)   : number of horizontal points
        ymax (float) : maximum vertical axis value
        ymin (float) : minimum vertical axis value
        ylen (int)   : number of vertical points
        plane        : stored complex plane implementation
        fs (list[function]) : function sequence to transform plane
    """
    #initializing the attributes
    def __init__(self, xmin, xmax, xlen, ymin, ymax, ylen):
        """The class constructor that initializes the complex plane

        Args:
            xmin (float): Minimum x value
            xmax (float): Maximum x value
            xlen (int): Number of horizontal points
            ymin (float): Minimum y value
            ymax (float): Maximum y value
            ylen (int): Number of vertical points
        """
        self.xmin = xmin
        self.xmax = xmax
        self.xlen = xlen
        self.ymin = ymin
        self.ymax = ymax
        self.ylen = ylen
        self.plane = self.__create_plane(self.xmin, self.xmax, self.xlen, self.ymin, self.ymax, self.ylen)
        self.fs = []

    ##creates plane of complex points
    def __create_plane(self, xmin, xmax, xlen, ymin, ymax, ylen):
        """This function creates a list of lists that contains the points for the complex plane.  It does so by
           starting at the minimum point, and moves towards the maximum point in equal intervals based on the xlen
           and ylen values.

        Args:
            xmin (float): Minimum x value
            xmax (float): Maximum x value
            xlen (int): Number of horizontal points
            ymin (float): Minimum y value
            ymax (float): Maximum y value
            ylen (int): Number of vertical points
        Returns:
            plane[]: The list of complex points
        """
        #gets values we can use to space the points out evenly
        dx = (xmax - xmin)/(xlen - 1)
        dy = (ymax - ymin)/(ylen - 1)

        #lists of points on the x and y axis
        xpoints = [xmin +i*dx for i in range(xlen)]
        ypoints = [ymin +i*dy for i in range(ylen)]

        #create initial plane list
        plane = []
        #adds the complex numbers to the plane list so that they're
        #added to each list by row
        for x in xpoints:
            plane.append([])
            for y in ypoints:
                plane[-1].append(x + y*1j)
        return plane

    def refresh(self):
        """Regenerates the complex plane and replaces all the attribute values.
           It also removes any functions stored in the attribute fs.
        """
        self.plane = self.__create_plane(self, xmin, xmax, xlen, ymin, ymax, ylen)
        self.fs = []

    def apply(self, f):
        """Adds a function f to the attribute fs.  Then it applies that function to every
           point in the plane.
        """
        #transforms points of the plane
        for i in range(self.xlen):
            for j in range(self.ylen):
                self.plane[i][j] = f(self.plane[i][j])

        self.fs.append(f)

    def zoom(self,xmin,xmax,xlen,ymin,ymax,ylen):
        """Resets the x and y attributes and recreates the plane with new x and y values.
        Then it goes through all the transformations in fs and applies them in order to
        the new values.
        """
        self.xmin = xmin
        self.xmax = xmax
        self.xlen = xlen
        self.ymin = ymin
        self.ymax = ymax
        self.ylen = ylen

        self.plane = self.__create_plane(self.xmin, self.xmax, self.xlen, self.ymin, self.ymax, self.ylen)
        listfs = list(self.fs)
        #applies the transformations in fs in order
        for f in listfs:
            self.apply(f)

        return self.plane