import pandas as pd
import numpy as np
from datetime import date, datetime
import json

import logging
from calendar import monthrange

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# Define revenue categories in one place for easy maintenance
REVENUE_CATEGORIES = ['serverless', 'dedicated', 'support', 'starter', 'essential', 'premium', 'lake', 'byoc']
CREDIT_TYPES = ['program', 'redeem', 'customer_success']

# Finance-approved, month-scoped exceptions for CRM orders whose source fields
# are not yet usable by the generic allocator. Keep every key exact so this does
# not make other Draft or out-of-term orders eligible.
SPECIAL_CP_ALLOCATION_OVERRIDES = {
    (
        '2026-07',
        '1372813089209222473',
        '801RC00000OFxDIYA1',
    ): {
        'order_number': 'JP-OF-25021801',
    },
}

# Generate order_columns dynamically based on revenue categories
def generate_order_columns():
    base_columns = [
        'accumulate_revenue_to_org_tenant_id',
        'order_name',
        'account_number',
        'account_name',
        'order_service_start_date',
        'order_service_end_date',
        'order_is_terminated',
        'order_termination_date',
        'real_order_termination_date',
        'need_new_order',
        'bill_month',
        'reason',
        'revenue_amount',
        'credit_amount',
        'accumulated_revenue',
        'accumulated_revenue_before_adjust',
        'accumulated_credit_amount',
        'adjusted_revenue_amount',
        'revenue_breakdown_ratio',
    ]

    # Add category-specific columns
    category_columns = []
    for category in REVENUE_CATEGORIES:
        category_columns.extend([
            f'{category}_amount',
            f'adjusted_{category}_amount',
        ])

    # Add fixed columns
    fixed_columns = [
        'adjusted_shortfall_amount',
        'order_net_fees',
        'deferred_revenue',
        'new_deferred_revenue',
        'is_consider_credit_during_adj_rev_breakdown',
        'is_order_closing',
        'billing_allocation_amount',
        'end_of_contract_shortfall_amount',
    ]

    return base_columns + category_columns + fixed_columns

order_columns = generate_order_columns()


def coalesce(val, val_to_return):
    if val is None or np.isnan(val):
        return val_to_return
    return val


def calculate_actual_closing_revenue(
        order_total_fees,
        accumulated_actual_revenue_before_billing,
        billing_allocation_amount):
    """Return the current-period actual RevRec amount and its closing shortfall.

    The current billing allocation is actual revenue for this bill month.  A
    commitment that closes in the month must then recognise the rest of its
    total fee after only prior-period actual revenue has been deducted.
    """
    remaining_actual_revenue = max(
        float(order_total_fees) - float(accumulated_actual_revenue_before_billing),
        0.0,
    )
    billing_allocation_amount = min(
        max(float(billing_allocation_amount), 0.0),
        remaining_actual_revenue,
    )
    shortfall_amount = remaining_actual_revenue - billing_allocation_amount

    return remaining_actual_revenue, shortfall_amount


def _is_missing_scalar(val):
    if val is None:
        return True
    if isinstance(val, float) and np.isnan(val):
        return True
    if val is pd.NaT:
        return True
    try:
        return bool(pd.isna(val))
    except (TypeError, ValueError):
        return False


def format_optional_date_str(val):
    """ISO date strings for BigQuery STRING; downstream expects string (e.g. Go bigquery.NullString)."""
    if _is_missing_scalar(val):
        return None
    if isinstance(val, pd.Timestamp):
        return val.strftime('%Y-%m-%d')
    if isinstance(val, datetime):
        return val.date().isoformat()
    if isinstance(val, date):
        return val.isoformat()
    if isinstance(val, str):
        s = val.strip()
        if not s:
            return None
        for fmt in ('%Y-%m-%d', '%Y/%m/%d', '%Y-%m-%d %H:%M:%S', '%Y/%m/%d %H:%M:%S'):
            try:
                return datetime.strptime(s, fmt).date().isoformat()
            except ValueError:
                continue
        return s
    if isinstance(val, (int, np.integer)):
        s = str(int(val)).zfill(8)
        if len(s) == 8 and s.isdigit():
            return f'{s[:4]}-{s[4:6]}-{s[6:8]}'
        return str(int(val))
    return str(val)


def create_on_demand_order(
        tenant_id,
        bill_month,
        start_date,
        end_date,
        **category_amounts):
    """
    Create an on-demand order with dynamic revenue category amounts.
    category_amounts should contain keys like: serverless_amount, dedicated_amount, etc.
    """
    # Get total amount from category_amounts or default parameter
    amount = category_amounts.pop('amount', 0)

    o = process_order({})
    o['need_new_order'] = True
    o['accumulate_revenue_to_org_tenant_id'] = tenant_id
    o['order_name'] = '__new_order__'
    o['order_service_start_date'] = start_date
    o['order_service_end_date'] = end_date

    o['bill_month'] = bill_month
    o['revenue_amount'] = amount
    o['order_net_fees'] = amount
    o['accumulated_revenue_before_adjust'] = 0.0

    # Set category amounts dynamically
    for category in REVENUE_CATEGORIES:
        amount_key = f'{category}_amount'
        o[amount_key] = category_amounts.get(amount_key, 0.0)

    # Create logging message dynamically
    category_breakdown = ' '.join([f'{cat}:{category_amounts.get(f"{cat}_amount", 0.0)}' for cat in REVENUE_CATEGORIES])
    logger.info('creating on demand order for {} at {}: {} with revenue breakdowns: {}'.format(
        tenant_id, bill_month, amount, category_breakdown))

    return o

def append_json_string(existing_json_str, new_json_str):
    existing_is_missing = _is_missing_scalar(existing_json_str) or existing_json_str == ''
    new_is_missing = _is_missing_scalar(new_json_str) or new_json_str == ''
    if existing_is_missing or new_is_missing:
        if existing_is_missing and new_is_missing:
            return None
        return new_json_str if existing_is_missing else existing_json_str
    existing_list = json.loads(existing_json_str)
    new_list = json.loads(new_json_str)
    combined_list = existing_list + new_list
    return json.dumps(combined_list)

def prioritise_orders(orders, billing_start_date, billing_end_date):
    # We need to prioritise orders based on the following rules:
    # 1. Orders that are terminated this month will have highest priority: order_is_terminated
    # 2. Orders that end naturally this month will have higher priority: order_service_end_date between billing_start_date and billing_end_date
    # 3. Then, orders with smaller start date will have higher priority: order_service_start_date
    # 4. Then, orders with a reason field will have higher priority: reason
    # We will sort the orders based on these rules

    if isinstance(orders, pd.DataFrame):
        # Create sorting keys for DataFrame
        # First key: order_is_terminated (True first)
        # Second key: ends this month (True first)
        # Third key: order_service_start_date (ascending)
        # Fourth key: has reason (True first)

        orders = orders.copy()

        orders['_ends_this_month'] = (
            (orders['order_service_end_date'] >= billing_start_date) &
            (orders['order_service_end_date'] <= billing_end_date)
        )

        orders['_has_reason'] = orders['reason'].notna() & (orders['reason'] != '')

        # Sort the DataFrame
        sorted_orders = orders.sort_values(
            by=['order_is_terminated', '_ends_this_month', 'order_service_start_date', '_has_reason'],
            ascending=[False, False, True, False]
        )

        # Drop temporary columns
        sorted_orders = sorted_orders.drop(columns=['_ends_this_month', '_has_reason'])

        return sorted_orders

    elif isinstance(orders, list):
        # For handling lists of dictionaries (like in tests)
        def sort_key(order):
            is_terminated = bool(order.get('order_is_terminated') or False)
            end_date = order.get('order_service_end_date')
            start_date = order.get('order_service_start_date') or date.min
            has_reason = bool(order.get('reason', ''))

            ends_this_month = (
                end_date and billing_start_date and billing_end_date and
                end_date >= billing_start_date and end_date <= billing_end_date
            )

            # Return tuple for sorting: (terminated, ends_this_month, start_date, has_reason)
            # Negate first two booleans because True sorts after False by default
            return (-int(is_terminated), -int(ends_this_month), start_date, -int(has_reason))

        return sorted(orders, key=sort_key)

    return orders

def break_billing_segments(df):
    # The source view exposes the default billing window as M/D/YYYY strings,
    # while order dates arrive as DATE values.  Normalise before comparing the
    # two in process(), and keep the segment rows on the same representation.
    df = df.copy()
    for column in ('billing_start_date', 'billing_end_date'):
        df[column] = pd.to_datetime(df[column], errors='raise').dt.date

    # Here we fetch all rows that have
    # segment_a_start_date, segment_a_end_date
    # segment_b_start_date, segment_b_end_date
    # segment_c_start_date, segment_c_end_date
    # that are not empty.

    # We want to split this row into multiple rows (depending on how many segment data we have).
    # Changes:
    #  start_date = segment_start_date
    #  end_date = segment_end_date
    #  total_amount = segment_amount
    # Rest of rows remain the same.
    segments = ['a', 'b', 'c']
    dfs = []

    for segment in segments:
        mask = (
            pd.notna(df[f'segment_{segment}_start_date']) &
            (df[f'segment_{segment}_start_date'] != '') &
            pd.notna(df[f'segment_{segment}_end_date']) &
            (df[f'segment_{segment}_end_date'] != '')
        )
        segment_df = df[mask].copy()
        segment_ratio = segment_df[f'segment_{segment}_amount'] / segment_df['total_amount']

        segment_df['billing_start_date'] = pd.to_datetime(segment_df[f'segment_{segment}_start_date']).dt.date
        segment_df['billing_end_date'] = pd.to_datetime(segment_df[f'segment_{segment}_end_date']).dt.date
        # We also need to set the breakdown columns by calculating the ratio of each segment amount to the total amount
        for category in REVENUE_CATEGORIES:
            raw_col = f'raw_{category}_amount'
            if raw_col in segment_df.columns:
                segment_df[raw_col] = segment_df[raw_col] * segment_ratio

        segment_df['total_amount'] = segment_df[f'segment_{segment}_amount']
        dfs.append(segment_df)

    df_original = df[~df.index.isin(pd.concat(dfs).index)]
    df_combined = pd.concat([df_original] + dfs, ignore_index=True)

    return df_combined

def process_order(order):
    o = {}
    for col in order_columns:
        if col in order:
            o[col] = order[col]
        else:
            o[col] = None
    return o

def model(dbt, session):
    return None

def build_order_sql(order_source, order_date_str, billing_date_str):
    """Build the order snapshot query with prior-period actual CP balances."""
    return f"""
        WITH order_snapshot AS (
            SELECT *
            FROM {order_source}
            WHERE date_day = '{order_date_str}'
        ),
        actual_revenue_before_billing AS (
            SELECT
                revrec.`order` AS order_id,
                SUM(COALESCE(
                    SAFE_CAST(NULLIF(revrec.actual_revenue_amount, '') AS FLOAT64),
                    SAFE_CAST(NULLIF(revrec.revenue_amount, '') AS FLOAT64),
                    0
                )) AS accumulated_actual_revenue_before_billing,
                SUM(COALESCE(
                    SAFE_CAST(NULLIF(revrec.actual_revenue_amount, '') AS FLOAT64),
                    SAFE_CAST(NULLIF(revrec.adjusted_revenue_amount, '') AS FLOAT64),
                    0
                )) AS accumulated_actual_adjusted_revenue_before_billing,
                SUM(COALESCE(
                    SAFE_CAST(NULLIF(revrec.credit_amount, '') AS FLOAT64),
                    0
                )) AS accumulated_actual_credit_amount_before_billing
            FROM `gcp-digital-bi.data_foundation_finance_final.data_foundation_finance_monthly_cloud_revrec` revrec
            INNER JOIN order_snapshot orders
                ON orders.order_id = revrec.`order`
            WHERE revrec.deleted_at IS NULL
                AND SAFE.PARSE_DATE('%Y-%m', revrec.bill_month) < DATE '{billing_date_str}-01'
                AND SAFE_CAST(revrec.billing_end_date AS DATE)
                    BETWEEN orders.order_service_start_date AND orders.order_service_end_date
                AND (
                    LOWER(TRIM(COALESCE(revrec.revenue_recognition_type, ''))) = 'actual'
                    OR (
                        TRIM(COALESCE(revrec.revenue_recognition_type, '')) = ''
                        AND revrec.estimate IS NULL
                    )
                )
            GROUP BY revrec.`order`
        )
        SELECT
            orders.*,
            COALESCE(actual.accumulated_actual_revenue_before_billing, 0)
                AS accumulated_actual_revenue_before_billing,
            COALESCE(actual.accumulated_actual_adjusted_revenue_before_billing, 0)
                AS accumulated_actual_adjusted_revenue_before_billing,
            COALESCE(actual.accumulated_actual_credit_amount_before_billing, 0)
                AS accumulated_actual_credit_amount_before_billing
        FROM order_snapshot orders
        LEFT JOIN actual_revenue_before_billing actual
            ON orders.order_id = actual.order_id
    """


def local_run():
    from google.cloud import bigquery
    # Create a FileHandler and set the filename
    file_handler = logging.FileHandler('my_application.log')

    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    file_handler.setFormatter(formatter)

    logger.addHandler(file_handler)


    # Specify whether this is a prod run or a test run
    # During prod runs, no further change is needed - only need to check if on-demand orders are not created yet
    # During test runs, you need to set month_to_test to the month you want to test
    is_prod_run = True

    # When True, only process billing rows where actual_revenue_recognition_date is populated (recognised revenue).
    # When False, only process billing rows where it is NOT populated (predicted revenue).
    is_recognised_revenue = True

    if is_prod_run:
        today = datetime.now().date()
        previous_month = today.month - 1 if today.month > 1 else 12
        previous_month_year = today.year if today.month > 1 else today.year - 1
        billing_date = datetime(previous_month_year, previous_month, 1).date()
        order_snapshot_date = today

    else:
        month_to_test = '2026-03'
        test_tenant_ids = None #['1372813089191271282']

        billing_date = datetime.strptime(month_to_test, '%Y-%m')
        order_snapshot_date = datetime(billing_date.year, billing_date.month, 3)

    billing_date_str = billing_date.strftime('%Y-%m')
    order_date_str = order_snapshot_date.strftime('%Y-%m-%d')

    client = bigquery.Client(project="gcp-digital-bi")
    if is_prod_run:
        billing_sql = f"""
            SELECT *
            FROM `gcp-digital-bi.dwh_finance.salesforce_org_monthly_revrec`
            WHERE billing_month = '{billing_date_str}'
        """
        order_source = '`gcp-digital-bi.dwh_finance.salesforce_order_daily_view`'
    else:
        billing_sql = f"""
            SELECT *
            FROM `gcp-digital-bi-dev.sandbox_dwh_finance_ywang.salesforce_org_monthly_revrec`
            WHERE billing_month = '{billing_date_str}'
        """
        order_source = '`gcp-digital-bi-dev.sandbox_dwh_finance_ywang.salesforce_order_daily_view`'

    # CP allocation in both estimate and actual modes must not treat estimates
    # as consumed commitment. Legacy rows (before recognition type existed) are
    # final actual revenue; newer rows must explicitly be actual. The current
    # bill month is deliberately excluded because process() allocates its billing.
    order_sql = build_order_sql(order_source, order_date_str, billing_date_str)
    df_orders = client.query_and_wait(order_sql).to_dataframe()
    df_billing = client.query_and_wait(billing_sql).to_dataframe()

    df_orders = df_orders[
        (df_orders['date_day'] == order_snapshot_date) &
        (df_orders['order_service_start_date'] <= order_snapshot_date)
    ]

    if not is_prod_run:
        # Further filter out previously created on-demand orders
        # These are identified by service_start_date = start of month, service_end_date = end of month
        df_orders = df_orders[
            ~((df_orders['order_service_start_date'] == datetime(billing_date.year, billing_date.month, 1).date()) &
              (df_orders['order_service_end_date'] == datetime(billing_date.year, billing_date.month, monthrange(billing_date.year, billing_date.month)[1]).date())
            )
        ]

        if test_tenant_ids:
            df_billing = df_billing[
            df_billing['accumulate_revenue_to_org_tenant_id'].isin(test_tenant_ids) |
            df_billing['tenant_id'].isin(test_tenant_ids)
        ]

    if is_recognised_revenue:
        df_billing = df_billing[df_billing['actual_revenue_recognition_date'].notna() & (df_billing['actual_revenue_recognition_date'] != '')]
        revenue_mode = 'recognised'
    else:
        df_billing = df_billing[df_billing['actual_revenue_recognition_date'].isna() | (df_billing['actual_revenue_recognition_date'] == '')]
        revenue_mode = 'predicted'

    logger.info(f'revenue mode: {revenue_mode}, billing rows after filter: {len(df_billing)}')

    df_org_output, df_order_output = process(
        df_orders=df_orders,
        df_billing=df_billing,
        is_recognised_revenue=is_recognised_revenue,
    )

    job_config = bigquery.LoadJobConfig(
        write_disposition="WRITE_TRUNCATE",
    )
    if is_prod_run:
        # We load the dataframes into both prod and a backup table
        job = client.load_table_from_dataframe(df_org_output, "gcp-digital-bi-dev.sandbox_dwh_ywang.monthly_revrec_collections", job_config=job_config)
        job.result()  # Wait for the job to complete.
        job = client.load_table_from_dataframe(df_order_output, "gcp-digital-bi-dev.sandbox_dwh_ywang.monthly_revrec_orders", job_config=job_config)
        job.result()  # Wait for the job to complete.
        month_suffix = billing_date.strftime('%Y_%m')
        job = client.load_table_from_dataframe(df_org_output, f"gcp-digital-bi-dev.sandbox_dwh_ywang.monthly_revrec_collections_{month_suffix}_{revenue_mode}", job_config=job_config)
        job.result()  # Wait for the job to complete.
        job = client.load_table_from_dataframe(df_order_output, f"gcp-digital-bi-dev.sandbox_dwh_ywang.monthly_revrec_orders_{month_suffix}_{revenue_mode}", job_config=job_config)
        job.result()  # Wait for the job to complete.

    else:
        job = client.load_table_from_dataframe(df_org_output, "gcp-digital-bi-dev.sandbox_ywang.revrec_org_test", job_config=job_config)
        job.result()  # Wait for the job to complete.
        job = client.load_table_from_dataframe(df_order_output, "gcp-digital-bi-dev.sandbox_ywang.revrec_order_test", job_config=job_config)
        job.result()  # Wait for the job to complete.

def process(df_orders, df_billing, is_recognised_revenue=False):
    # TODO: TEMP: break billing by segments into individual rows
    df_billing = break_billing_segments(df_billing)

    # Keep only Activated orders plus the exact CRM records named by an override.
    # Service window, tenant, and override month are evaluated per billing row.
    globally_overridden_order_ids = {
        order_id
        for _, _, order_id in SPECIAL_CP_ALLOCATION_OVERRIDES
    }
    if 'order_id' in df_orders.columns:
        globally_overridden_order_mask = df_orders['order_id'].isin(
            globally_overridden_order_ids
        )
    else:
        globally_overridden_order_mask = pd.Series(False, index=df_orders.index)
    df_orders = df_orders[
        (df_orders['order_product_family'] == 'TiDB Cloud') &
        (
            (df_orders['order_status'] == 'Activated') |
            globally_overridden_order_mask
        )
    ]

    # For each billing recorrd, we need to know:
    #   1. How much billing occurred for that billing period
    #   2. What is the deferred revenue of order(s) for that billing period
    #   3. Whether existing order(s) end/terminate during that billing period

    all_orders = list()
    org_info = dict()

    closed_orders = set()

    for billing in df_billing.to_dict('records'):

        logger.info("processing row {}".format(str(billing['tenant_id'])))

        tenant_id = billing['accumulate_revenue_to_org_tenant_id']
        org_billing = billing['total_amount']
        bill_month = billing['billing_month']
        bill_account_id = billing['account_id']
        year, month = map(int, bill_month.split('-'))
        month_first_day = datetime(year, month, 1).date()
        month_last_day = datetime(year, month, monthrange(year, month)[1]).date()
        billing_start_date = billing['billing_start_date']
        billing_end_date = billing['billing_end_date']
        remaining_billing = org_billing

        def stripe_stuff(o, b, key, col_name):
            return append_json_string(o.get(key, {}).get(col_name, None), b[col_name])


        # Collect org info
        org_info[(tenant_id, bill_month)] = {
            'month_first_day': month_first_day,
            'month_last_day': month_last_day,
            'org_total_billing': org_info.get((tenant_id, bill_month), {}).get('org_total_billing', 0.0) + org_billing,
            'account_surrogate_key': billing['customer_surrogate_key'],
            'account_number': billing['account_number_for_netsuite'],
            'account_name': billing['customer_name'],
            'region': billing['region'],
            'payment_type': billing['payment_type'],
            'wire_transfer_amount': billing['wire_transfer_amount'], # to sum
            'marketplace_amount': billing['marketplace_amount'],
            'stripe_amount': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_amount'),
            'stripe_tax_rate': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_tax_rate'),
            'stripe_amount_including_tax': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_amount_including_tax'),
            'stripe_invoice_number': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_invoice_number'),
            'stripe_invoice': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_invoice'),
            'stripe_invoice_paid': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_invoice_paid'),
            'stripe_invoice_finalized_date': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_invoice_finalized_date'),
            'stripe_invoice_paid_date': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_invoice_paid_date'),
            'collections_from': '' if (_is_missing_scalar(billing['stripe_amount']) or billing['stripe_amount'] == '') else 'Stripe',
            'invoice_amount': billing['invoice_amount'],
            'is_salesforce_customer': billing['is_salesforce_customer'],
            'actual_revenue_recognition_date': billing.get('actual_revenue_recognition_date'),
        }

        # Calculate total credit amount dynamically
        all_credit_fields = []
        for category in REVENUE_CATEGORIES:
            for credit_type in CREDIT_TYPES:
                all_credit_fields.append(f'raw_{credit_type}_credit_amount_{category}')

        org_info[(tenant_id, bill_month)]['credit_amount'] = (
            org_info.get((tenant_id, bill_month), {}).get('credit_amount', 0.0) +
            sum([coalesce(billing.get(field), 0.0) for field in all_credit_fields])
        )

        # Calculate category-specific credit amounts dynamically
        for category in REVENUE_CATEGORIES:
            category_credit_fields = [f'raw_{credit_type}_credit_amount_{category}' for credit_type in CREDIT_TYPES]
            credit_key = f'{category}_credit_amount'
            org_info[(tenant_id, bill_month)][credit_key] = (
                org_info.get((tenant_id, bill_month), {}).get(credit_key, 0.0) +
                sum([coalesce(billing.get(field), 0.0) for field in category_credit_fields])
            )

        logger.info("processing billing record for org: {} at {}".format(tenant_id, bill_month))
        bill_orders = []

        # We only examine orders that are active between the billing start/end dates
        # We also filter out orders that are not terminated this month
        normal_order_mask = (
            (df_orders['order_status'] == 'Activated') &
            (df_orders['order_service_start_date'] <= billing_end_date) &
            (df_orders['order_service_end_date'] >= billing_start_date) &
            (df_orders['accumulate_revenue_to_org_tenant_id'] == tenant_id) & (
                (df_orders['order_is_terminated'] == False) | (
                    (df_orders['order_is_terminated'] == True) &
                    (df_orders['order_termination_date'] > billing_end_date)
                )
            )
        )

        special_overrides = {
            order_id: override
            for (override_month, override_tenant_id, order_id), override
            in SPECIAL_CP_ALLOCATION_OVERRIDES.items()
            if override_month == bill_month and override_tenant_id == tenant_id
        }
        if special_overrides and 'order_id' in df_orders.columns:
            special_order_mask = df_orders['order_id'].isin(special_overrides)
        else:
            special_order_mask = pd.Series(False, index=df_orders.index)

        _df_orders = df_orders[normal_order_mask | special_order_mask].copy()
        if 'order_id' in _df_orders.columns:
            for order_id, override in special_overrides.items():
                override_mask = _df_orders['order_id'] == order_id
                _df_orders.loc[
                    override_mask,
                    'accumulate_revenue_to_org_tenant_id',
                ] = tenant_id
                _df_orders.loc[override_mask, 'order_name'] = override['order_number']

        _df_orders = prioritise_orders(_df_orders, billing_start_date, billing_end_date)

        # Then, we loop through all orders for the tenant, and assign correct revenue to each order
        # The priority of orders which billing is applied to is as below:
        # 1. Orders that finish this month (order_service_end_date)
        # 2. Orders that were created earlier (created_date)
        # last. newly created on demand order(s)
        for _, order in _df_orders.iterrows():
            logger.info("order net fees: {}, accumulated revenue b4 adjust: {}, accumulated_revenue: {}".format(order['order_net_fees'], order['accumulated_revenue_before_adjust'], order['accumulated_revenue']))

            order['bill_month'] = bill_month
            order['order_net_fees'] = float(round(order['order_net_fees'], 4))

            is_order_closing = (
                (order['order_service_end_date'] >= billing_start_date and order['order_service_end_date'] <= billing_end_date) or
                (order['order_is_terminated'] == True and order['real_order_termination_date'] == billing_end_date)
            )
            is_commitment_plan = (
                order.get('cloud_order_type') == 'Commitment Plan'
            )
            if is_commitment_plan:
                # The order daily view aggregates both actual and estimated
                # Revenue__c rows. For any CP allocation (estimate or actual),
                # only prior periods' actual revenue may reduce the commitment
                # balance; otherwise estimate + actual makes a continuing CP
                # look fully consumed and sends billing to a new on-demand order.
                accumulated_actual_revenue = coalesce(
                    order.get('accumulated_actual_revenue_before_billing'),
                    0.0,
                )
                order['accumulated_revenue_before_adjust'] = accumulated_actual_revenue
                order['accumulated_revenue'] = coalesce(
                    order.get('accumulated_actual_adjusted_revenue_before_billing'),
                    accumulated_actual_revenue,
                )
                order['accumulated_credit_amount'] = coalesce(
                    order.get('accumulated_actual_credit_amount_before_billing'),
                    0.0,
                )

            deferred_revenue = order['order_net_fees'] - order['accumulated_revenue_before_adjust']

            if deferred_revenue < 0.02:
                logger.warning('Order {} does not have deferred revenue, skipping..'.format(order['order_name']))
                continue


            # If order ends this month, we set the order_end flag
            if is_order_closing:
                logger.info('order {} ends this month.'.format(order['order_name']))
                closed_orders.add((order['order_name'], tenant_id, bill_month))

                # order['revenue_amount'] = deferred_revenue
                # remaining_billing -= deferred_revenue
                # if remaining_billing <= 0:
                #     # Assign full breakdown amounts
                #     for category in REVENUE_CATEGORIES:
                #         raw_col = f'raw_{category}_amount'
                #         category_col = f'{category}_amount'
                #         order[category_col] = billing.get(raw_col, 0.0)
                # else:
                #     # Proportionally assign breakdown amounts
                #     for category in REVENUE_CATEGORIES:
                #         raw_col = f'raw_{category}_amount'
                #         category_col = f'{category}_amount'
                #         order[category_col] = float(deferred_revenue / org_billing) * billing.get(raw_col, 0.0)

                # # Create logging message dynamically
                # category_breakdown = ' '.join([f'{cat}: {order.get(f"{cat}_amount", 0)}' for cat in REVENUE_CATEGORIES])
                # logger.info('order {} ends this month, assigning all revenue to this month. Revenue: {}, breakdowns: {}'.format(
                #     order['order_name'], order['revenue_amount'], category_breakdown))
            # else:
            if True:
                if deferred_revenue > remaining_billing:
                    order['revenue_amount'] = remaining_billing
                    if org_billing == 0:
                        # Set all category amounts to 0
                        for category in REVENUE_CATEGORIES:
                            order[f'{category}_amount'] = 0
                    else:
                        # Proportionally assign breakdown amounts
                        for category in REVENUE_CATEGORIES:
                            raw_col = f'raw_{category}_amount'
                            category_col = f'{category}_amount'
                            order[category_col] = float(remaining_billing / org_billing) * billing.get(raw_col, 0.0)
                        remaining_billing = 0
                else:
                    order['revenue_amount'] = deferred_revenue
                    # Proportionally assign breakdown amounts
                    for category in REVENUE_CATEGORIES:
                        raw_col = f'raw_{category}_amount'
                        category_col = f'{category}_amount'
                        order[category_col] = float(deferred_revenue / org_billing) * billing.get(raw_col, 0.0)
                    remaining_billing -= deferred_revenue

                # Create logging message dynamically
                category_breakdown = ' '.join([f'{cat}: {order.get(f"{cat}_amount", 0)}' for cat in REVENUE_CATEGORIES])
                logger.info('order {} assigning revenue to this month. Revenue: {}, breakdowns: {}'.format(
                    order['order_name'], order['revenue_amount'], category_breakdown))

            if remaining_billing < 0:
                remaining_billing = 0


            order['order_net_fees'] = float(round(order['order_net_fees'], 2))

            bill_orders.append(process_order(order))

            if remaining_billing <= 0:
                logger.info('all billing has been allocated for org: {} at {}'.format(tenant_id, bill_month))
                break

        logger.info('total billing: {}'.format(str(org_billing)))
        logger.info('remaining billing: {}'.format(str(remaining_billing)))


        # If there are still money to be reconciled, then we need to create an on-demand order, for the bill, for that month
        if remaining_billing > 0.02:
            logger.info('on demand order created for bill {} at {}'.format(tenant_id, bill_month))

            # Calculate category amounts dynamically
            category_amounts = {}
            for category in REVENUE_CATEGORIES:
                raw_col = f'raw_{category}_amount'
                amount_col = f'{category}_amount'
                category_amounts[amount_col] = float(remaining_billing / org_billing) * billing.get(raw_col, 0.0)

            category_amounts['amount'] = remaining_billing
            bill_orders.append(create_on_demand_order(tenant_id, bill_month, month_first_day, month_last_day, **category_amounts))

        all_orders.extend(bill_orders)

    merged_orders = []

    # Finally, we loop through all orders and combine order records that have been billed across multiple bills
    # We use (order_name, tenant_id, bill_month) as keys to merge
    for tenant_id, bill_month in list(org_info.keys()):
        _orders = prioritise_orders(
            [x for x in all_orders if x['accumulate_revenue_to_org_tenant_id'] == tenant_id and x['bill_month'] == bill_month],
            org_info[(tenant_id, bill_month)]['month_first_day'],
            org_info[(tenant_id, bill_month)]['month_last_day']
            )

        n_o_key = ('__new_order__', tenant_id, bill_month)
        merged = {}
        excess = 0.0

        for order in _orders:
            key = (order['order_name'], tenant_id, bill_month)
            logger.info('processing order {}'.format(str(key)))
            if key not in merged:
                logger.info('adding order {}...'.format(str(key)))
                merged[key] = order
            else:
                logger.info('order {} is duplicated, merging...'.format(str(key)))
                merged_order = merged[key]
                merged_rev = merged_order['revenue_amount'] + order['revenue_amount']

                # Calculate merged category amounts dynamically
                merged_category_amounts = {}
                for category in REVENUE_CATEGORIES:
                    category_col = f'{category}_amount'
                    merged_category_amounts[category_col] = merged_order.get(category_col, 0) + order.get(category_col, 0)

                deferred_revenue = merged_order['order_net_fees'] - merged_order['accumulated_revenue_before_adjust']

                if merged_order['need_new_order']:
                    merged_order['revenue_amount'] = merged_rev
                    # Update all category amounts
                    for category in REVENUE_CATEGORIES:
                        category_col = f'{category}_amount'
                        merged_order[category_col] = merged_category_amounts[category_col]
                    merged_order['order_net_fees'] += order['order_net_fees']

                # For existing orders, revenue is capped at deferred_revenue
                # Excess revenue amounts will be moved to new orders
                elif merged_rev > deferred_revenue:
                    excess += merged_rev - deferred_revenue
                    merged_order['revenue_amount'] = deferred_revenue
                    # Proportionally adjust category amounts
                    for category in REVENUE_CATEGORIES:
                        category_col = f'{category}_amount'
                        merged_order[category_col] = float(deferred_revenue / merged_rev) * merged_category_amounts[category_col]

                else:
                    merged_order['revenue_amount'] = merged_rev
                    # Update all category amounts
                    for category in REVENUE_CATEGORIES:
                        category_col = f'{category}_amount'
                        merged_order[category_col] = merged_category_amounts[category_col]

            # If an order has been completed - either by being terminated / ends this month, or by being fully billed
            # We do not care about credit calculations for these orders
            if abs(merged[key]['order_net_fees'] - merged[key]['revenue_amount'] + merged[key]['accumulated_revenue_before_adjust']) <= 0.01 or key in closed_orders:
                merged[key]['is_consider_credit_during_adj_rev_breakdown'] = False
            else:
                merged[key]['is_consider_credit_during_adj_rev_breakdown'] = True

            if excess > 0:
                if n_o_key in merged:
                    logger.info('moving excess revenue to new order...')
                    merged[n_o_key]['revenue_amount'] += excess
                    merged[n_o_key]['order_net_fees'] += excess
                    # Update category amounts proportionally
                    for category in REVENUE_CATEGORIES:
                        category_col = f'{category}_amount'
                        merged[n_o_key][category_col] += float(excess / merged_rev) * merged_category_amounts[category_col]
                    logger.info(excess)
                else:
                    logger.info('creating new order for excess...')
                    # Calculate category amounts for new order
                    excess_category_amounts = {}
                    for category in REVENUE_CATEGORIES:
                        category_col = f'{category}_amount'
                        excess_category_amounts[category_col] = float(excess / merged_rev) * merged_category_amounts[category_col]

                    excess_category_amounts['amount'] = excess
                    _new_order = create_on_demand_order(tenant_id, bill_month, month_first_day, month_last_day, **excess_category_amounts)
                    merged[n_o_key] = _new_order

            # We do not care about credit for on demand orders
            if key == n_o_key:
                merged[key]['credit_amount'] = 0.0
                # Set all category credit amounts to 0
                for category in REVENUE_CATEGORIES:
                    merged[key][f'{category}_credit_amount'] = 0.0
                merged[key]['is_consider_credit_during_adj_rev_breakdown'] = False
            else:
                merged[key]['credit_amount'] = org_info[(tenant_id, bill_month)]['credit_amount']
                # Set category credit amounts
                for category in REVENUE_CATEGORIES:
                    credit_key = f'{category}_credit_amount'
                    merged[key][credit_key] = org_info[(tenant_id, bill_month)].get(credit_key, 0.0)

                # Set credit amount for org to 0 to avoid double counting
                org_info[(tenant_id, bill_month)]['credit_amount'] = 0.0
                for category in REVENUE_CATEGORIES:
                    credit_key = f'{category}_credit_amount'
                    org_info[(tenant_id, bill_month)][credit_key] = 0.0

        # Snapshot billing allocation before end-of-contract bump
        for key in merged.keys():
            merged[key]['is_order_closing'] = key in closed_orders
            merged[key]['billing_allocation_amount'] = merged[key]['revenue_amount']
            merged[key]['end_of_contract_shortfall_amount'] = 0.0

        # Then, we loop through orders that end in that billing month, and set their revenue to be fully recognised if not already
        for key in merged.keys():
            if key in closed_orders:
                _o = merged[key]
                deferred_revenue = _o['order_net_fees'] - _o['accumulated_revenue_before_adjust']
                if _o['revenue_amount'] < deferred_revenue:
                    if is_recognised_revenue and _o.get('cloud_order_type') == 'Commitment Plan':
                        revenue_amount, shortfall_amount = calculate_actual_closing_revenue(
                            _o['order_net_fees'],
                            _o['accumulated_revenue_before_adjust'],
                            _o['billing_allocation_amount'],
                        )
                        _o['revenue_amount'] = revenue_amount
                        _o['end_of_contract_shortfall_amount'] = shortfall_amount
                    else:
                        _o['revenue_amount'] = deferred_revenue
                        _o['end_of_contract_shortfall_amount'] = deferred_revenue - _o['billing_allocation_amount']
                    logger.info(f"order {_o['order_name']} ends this month, revenue adjusted to deferred revenue: {_o['revenue_amount']}, shortfall: {_o['end_of_contract_shortfall_amount']}")

        merged_orders.extend(merged.values())

    # Once revenue has been calculated for each order, we can now calculate adjusted values for the month.
    # TODO: This may need to be moved across scopes
    for order in merged_orders:
        revenue_amount = coalesce(order['revenue_amount'], 0)
        credit_amount = coalesce(order['credit_amount'], 0)
        acc_credit = coalesce(order['accumulated_credit_amount'], 0)
        acc_revenue = coalesce(order['accumulated_revenue_before_adjust'], 0)
        acc_adjusted_revenue = coalesce(order['accumulated_revenue'], 0)
        net_fees = coalesce(order['order_net_fees'], 0)
        deferred_revenue = net_fees - acc_revenue
        adjusted_deferred_revenue = net_fees - acc_adjusted_revenue
        is_consider_credit_during_adj_rev_breakdown = order['is_consider_credit_during_adj_rev_breakdown']

        # Get category amounts and credit amounts
        category_amounts = {}
        category_credit_amounts = {}
        for category in REVENUE_CATEGORIES:
            category_amounts[category] = coalesce(order.get(f'{category}_amount'), 0)
            category_credit_amounts[category] = coalesce(order.get(f'{category}_credit_amount'), 0)

        logger.debug('order {}: rev: {}, credit: {}, acc_revenue_before_adjust: {}, acc_credit: {}, net_fees: {}, acc_rev: {}, deferred_revenue: {}, adj_deferred_revenue: {}'.format(
            order['order_name'], revenue_amount, credit_amount, acc_revenue, acc_credit, net_fees, acc_revenue, deferred_revenue, adjusted_deferred_revenue
        ))

        # Create dynamic debug message for category amounts
        category_debug = ' '.join([f'{cat}: {category_amounts[cat]}' for cat in REVENUE_CATEGORIES])
        logger.debug('order {}: {}'.format(order['order_name'], category_debug))

        # Create dynamic debug message for category credit amounts
        category_credit_debug = ' '.join([f'{cat}_credit: {category_credit_amounts[cat]}' for cat in REVENUE_CATEGORIES])
        logger.debug('order {}: {}'.format(order['order_name'], category_credit_debug))

        logger.debug('order {}: is_consider_credit_during_adj_rev_breakdown: {}'.format(order['order_name'], is_consider_credit_during_adj_rev_breakdown))

        adjusted_revenue_amount = (((acc_revenue + revenue_amount + credit_amount + acc_credit) * net_fees) / (net_fees + credit_amount + acc_credit) ) - net_fees + adjusted_deferred_revenue

        # Calculate other adjusted values
        def adjusted_values(val, credit_val):
            if is_consider_credit_during_adj_rev_breakdown:
                if revenue_amount + credit_amount == 0:
                    return 0.0
                return (val + credit_val) * adjusted_revenue_amount / (revenue_amount + credit_amount)
            else:
                if revenue_amount == 0:
                    return 0.0
                return val * adjusted_revenue_amount / revenue_amount

        # Calculate adjusted amounts for each category dynamically
        adjusted_category_amounts = {}
        for category in REVENUE_CATEGORIES:
            adjusted_amount = adjusted_values(category_amounts[category], category_credit_amounts[category])
            adjusted_category_amounts[category] = adjusted_amount
            order[f'adjusted_{category}_amount'] = adjusted_amount

        # Calculate shortfall amount
        total_adjusted_categories = sum(adjusted_category_amounts.values())
        adjusted_shortfall_amount = adjusted_revenue_amount - total_adjusted_categories

        # Calculate deferred revenues
        order['deferred_revenue'] = net_fees - acc_adjusted_revenue
        order['new_deferred_revenue'] = net_fees - (adjusted_revenue_amount + acc_adjusted_revenue)

        # Assign adjusted amounts
        order['adjusted_revenue_amount'] = adjusted_revenue_amount
        order['adjusted_shortfall_amount'] = adjusted_shortfall_amount

        # Create dynamic logging messages
        adjusted_breakdown = ' '.join([f'{cat}: {adjusted_category_amounts[cat]}' for cat in REVENUE_CATEGORIES])
        original_breakdown = ' '.join([f'{cat}: {category_amounts[cat]}' for cat in REVENUE_CATEGORIES])

        logger.info('{} {} adjusted: rev: {}, {}, shortfall: {}'.format(
            order['accumulate_revenue_to_org_tenant_id'], order['order_name'],
            adjusted_revenue_amount, adjusted_breakdown, adjusted_shortfall_amount))
        logger.info('{} {} original: rev: {}, {}'.format(
            order['accumulate_revenue_to_org_tenant_id'], order['order_name'],
            revenue_amount, original_breakdown))
    # Now, we have a clear view of:
    # - Each new order to be created
    # - revrec infromation for each existing order

    # Temp: We reorganise the records into existing row format for integration with Xiaopeng's code
    # TODO: Refactor with downstream systems to reduce clutter
    rows = []

    org_data = []
    order_data = []

    for key in list(org_info.keys()):
        tenant_id, bill_month = key
        org = org_info[key]
        cur_orders = [o for o in merged_orders if o['accumulate_revenue_to_org_tenant_id'] == tenant_id and o['bill_month'] == bill_month]

        if len(cur_orders) == 0:
            logger.warning("{} is missing matchable orders.".format(str(key)))
            continue

        _org = {
            "bill_month": bill_month,
            "adjust_tenant_id": tenant_id,
            "start_date": org['month_first_day'],
            "end_date": org['month_last_day'],
            "account_surrogate_key": org['account_surrogate_key'],
            "account_id": org['account_number'],
            "account_name": org['account_name'],
            "account_region": org['region'],
            "payment_type": org['payment_type'],
            "marketplace_amount": org['marketplace_amount'],
            "stripe_amount": org['stripe_amount'],
            "wire_transfer_amount": org['wire_transfer_amount'],
            "total_amount": org['org_total_billing'],
            "stripe_tax_rate": org['stripe_tax_rate'],
            "stripe_amount_including_tax": org['stripe_amount_including_tax'],
            "stripe_invoice_number": org['stripe_invoice_number'],
            "stripe_invoice": org['stripe_invoice'],
            "stripe_invoice_paid": org['stripe_invoice_paid'],
            "stripe_invoice_finalized_date": org['stripe_invoice_finalized_date'],
            "stripe_invoice_paid_date": org['stripe_invoice_paid_date'],
            "invoice_amount": org['invoice_amount'],
            "actual_revenue_recognition_date": format_optional_date_str(org.get('actual_revenue_recognition_date')),
        }

        org_data.append(_org)

        for order in cur_orders:
            tenant_id = order['accumulate_revenue_to_org_tenant_id']
            bill_month = order['bill_month']

            account_id = org['account_number']
            account_name = org['account_name']
            need_new_order = order.get('need_new_order', False)
            is_new_order = (
                False if _is_missing_scalar(need_new_order)
                else bool(need_new_order)
            )
            if not is_new_order:
                crm_account_id = order.get('account_number')
                crm_account_name = order.get('account_name')
                if (
                    not _is_missing_scalar(crm_account_id)
                    and str(crm_account_id).strip()
                ):
                    account_id = crm_account_id
                if (
                    not _is_missing_scalar(crm_account_name)
                    and str(crm_account_name).strip()
                ):
                    account_name = crm_account_name

            if order.get('order_name', None) == '__new_order__':
                order_number = bill_month.replace('-', '_') + '__ON_DEMAND__' + org['account_number'].replace('-', '_')
            else:
                order_number = order.get('order_name', None)

            _order = {
                "tenant_id": tenant_id,
                "bill_month": bill_month,
                "is_new_order": order['need_new_order'],
                "is_new_salesforce_order": org['is_salesforce_customer'] and order['need_new_order'],
                "account_surrogate_key": org['account_surrogate_key'],
                "account_id": account_id,
                "account_name": account_name,
                "account_region": org['region'],
                "payment_type": org['payment_type'],
                "stripe_tax_rate": org['stripe_tax_rate'],
                "collections_from": org['collections_from'],
                "order_number": order_number,
                "order_start_date": order.get('order_service_start_date', None),
                "order_end_date": order.get('order_service_end_date', None),
                "order_billing_start_date": org['month_first_day'],
                "order_billing_end_date": org['month_last_day'],
                "order_deferred_revenue_usd": order.get('deferred_revenue'),
                "order_termination_date": format_optional_date_str(
                    order.get('order_termination_date', None)
                ),
                "order_total_fees_usd": order.get('order_net_fees', None),
                "order_accumulated_revenue_before_adjust": order.get('accumulated_revenue_before_adjust'),
                "order_accumulated_credit_amount": order.get('accumulated_credit_amount'),
                "revenue_amount": order.get('revenue_amount', None),
                "credit_amount": order.get('credit_amount', None),
                "adjusted_revenue_amount": order.get('adjusted_revenue_amount', None),
                "adjusted_shortfall_amount": order.get('adjusted_shortfall_amount', None),
                "is_order_closing": order.get('is_order_closing', False),
                "billing_allocation_amount": order.get('billing_allocation_amount', None),
                "end_of_contract_shortfall_amount": order.get('end_of_contract_shortfall_amount', None),
                "actual_revenue_recognition_date": format_optional_date_str(org.get('actual_revenue_recognition_date')),
            }

            # Dynamically add revenue breakdown amounts
            for category in REVENUE_CATEGORIES:
                _order[f"{category}_revenue_amount"] = order.get(f'{category}_amount', None)
                _order[f"adjusted_{category}_amount"] = order.get(f'adjusted_{category}_amount', None)

            order_data.append(_order)

    df_org = pd.DataFrame.from_dict(org_data)
    df_order = pd.DataFrame.from_dict(order_data)
    # All-NULL columns otherwise autodetect as INT64 in load_table_from_dataframe; Go expects STRING.
    if 'order_termination_date' in df_order.columns:
        df_order['order_termination_date'] = df_order['order_termination_date'].astype(pd.StringDtype())
    if 'actual_revenue_recognition_date' in df_order.columns:
        df_order['actual_revenue_recognition_date'] = df_order['actual_revenue_recognition_date'].astype(pd.StringDtype())
    if 'actual_revenue_recognition_date' in df_org.columns:
        df_org['actual_revenue_recognition_date'] = df_org['actual_revenue_recognition_date'].astype(pd.StringDtype())
    for column in [
        'stripe_amount',
        'stripe_tax_rate',
        'stripe_amount_including_tax',
        'stripe_invoice_number',
        'stripe_invoice',
        'stripe_invoice_paid',
        'stripe_invoice_finalized_date',
        'stripe_invoice_paid_date',
    ]:
        if column in df_org.columns:
            df_org[column] = df_org[column].astype(pd.StringDtype())

    return df_org, df_order

if __name__ == "__main__":
    local_run()