#!/usr/bin/env python
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Retrieves customers under a mcc manager and splits them according to their account_type."""

import argparse
import traceback
from typing import Any
from google.ads.searchads360.v0.enums.types.account_level import AccountLevelTypeEnum
from google.ads.searchads360.v0.services.types.search_ads360_service import SearchSearchAds360StreamRequest
from google.api_core import exceptions
from util_searchads360 import SearchAds360Client


def _map_data(data, data_dictionary) -> None:
  """Inserts customer data into its corresponding dictionary using id as the key."""
  if data:
    data_dictionary[data['id']] = data


def _prepare_sub_manager_data(customer) -> dict[str, Any]:
  """Using the customer object, prepares sub-manager info."""
  sub_manager_id = customer.sub_manager_id

  sub_managers_info = {
      'id': sub_manager_id,
      'name': customer.sub_manager_descriptive_name,
      'type': AccountLevelTypeEnum.AccountLevelType.SUB_MANAGER.name,
      'parent_id': customer.manager_id,
  }
  return sub_managers_info


def _prepare_associate_manager_data(customer) -> dict[str, Any] | None:
  """Using the customer object, prepares associate-manager info."""
  associate_manager_id = customer.associate_manager_id
  associate_manager_id_descriptive_name = str(
      customer.associate_manager_descriptive_name
  )
  if not associate_manager_id_descriptive_name or associate_manager_id == 0:
    return None

  associate_manager_info = {
      'id': associate_manager_id,
      'name': associate_manager_id_descriptive_name,
      'type': AccountLevelTypeEnum.AccountLevelType.ASSOCIATE_MANAGER.name,
      'parent_id': customer.sub_manager_id,
  }
  return associate_manager_info


def _prepare_client_customer_data(customer) -> dict[str, Any]:
  """Using the customer object, prepares client-customer info."""
  client_customer_id = customer.id

  # Determine parent ID for the client account
  parent_id = (
      customer.associate_manager_id
      or customer.sub_manager_id
      or customer.manager_id
  )

  client_customer_info = {
      'id': client_customer_id,
      'name': customer.descriptive_name,
      'type': customer.account_type.name,
      'parent_id': parent_id,
  }
  return client_customer_info


def _get_all_accounts_details(
    search_ads_360_service,
    client_customer_id,
    sub_managers_data,
    associate_managers_data,
    client_accounts_data,
) -> None:
  """Fetches details for a client customer and updates hierarchy maps."""
  query = """
      SELECT
        customer.id,
        customer.descriptive_name,
        customer.account_type,
        customer.manager_id,
        customer.manager_descriptive_name,
        customer.sub_manager_id,
        customer.sub_manager_descriptive_name,
        customer.associate_manager_id,
        customer.associate_manager_descriptive_name,
        customer.account_level
      FROM customer"""
  request = SearchSearchAds360StreamRequest()
  request.customer_id = client_customer_id
  request.query = query
  try:
    results_stream = search_ads_360_service.search_stream(request=request)
    for response in results_stream:
      for row in response.results:
        customer = row.customer
        _map_data(_prepare_sub_manager_data(customer), sub_managers_data)
        _map_data(
            _prepare_associate_manager_data(customer), associate_managers_data
        )
        _map_data(_prepare_client_customer_data(customer), client_accounts_data)

  except exceptions.GoogleAPICallError as e:
    print(
        'ERROR: An unexpected error occurred while fetching customer'
        f' information for {client_customer_id}.\nDetails: {e}\n'
    )


def _get_client_customer_ids(search_ads_360_service, customer_id) -> list[str]:
  """Retrieves a list of client customer IDs (if the customer is a manager).

  Args:
    search_ads_360_service: The Search Ads 360 service client.
    customer_id: The customer ID of the mcc manager customer.

  Returns:
    A list of client customer IDs.
  """
  request = SearchSearchAds360StreamRequest()
  query = """
    SELECT
      customer_client.id,
      customer_client.manager
    FROM customer_client
    WHERE customer_client.manager = False
  """
  customer_client_ids = []
  request.customer_id = customer_id
  request.query = query
  try:
    # Issues a search stream request.
    results_stream = search_ads_360_service.search_stream(request=request)
    for response in results_stream:
      for row in response.results:
        customer_client_ids.append(str(row.customer_client.id))

  except exceptions.GoogleAPICallError as e:
    print(
        'ERROR: An unexpected error occurred while fetching client customer IDs'
        f' for {customer_id}.\nDetails: {e}\n'
    )
  return customer_client_ids


def main(service_client, login_customer_id) -> None:
  search_ads_360_service = service_client.get_service()

  sub_managers_data = {}
  associate_managers_data = {}
  client_accounts_data = {}

  # Get all client customer IDs under the login_customer_id (MCC manager)
  customer_client_ids = _get_client_customer_ids(
      search_ads_360_service, login_customer_id
  )

  for customer_client_id in customer_client_ids:
    _get_all_accounts_details(
        search_ads_360_service,
        customer_client_id,
        sub_managers_data,
        associate_managers_data,
        client_accounts_data,
    )


if __name__ == '__main__':
  # SearchAds360Client will read the search-ads-360.yaml configuration file in
  # the home directory if none is specified.
  search_ads_360_client = SearchAds360Client.load_from_file()

  parser = argparse.ArgumentParser(
      description=(
          'Retrieves all customer_client for a mcc manager customer and runs a'
          ' query on each customer_client.'
      )
  )

  # Arguments to provide to run the example.
  parser.add_argument(
      '-l',
      '--login_customer_id',
      type=str,
      required=True,
      help=(
          'The Search Ads 360 MCC manager login customer ID (10 digits, no'
          ' dashes). This can be obtained from the UI or from calling'
          ' ListAccessibleCustomers - '
          'https://developers.google.com/search-ads/reporting/concepts/login-customer-id.'
      ),
  )

  args = parser.parse_args()

  search_ads_360_client.set_ids(args.login_customer_id, args.login_customer_id)

  try:
    main(search_ads_360_client, args.login_customer_id)
  except Exception:  # pylint: disable=broad-except
    traceback.print_exc()