"""Bulk create splits for collaborator."""

from functools import partial
import logging
import math
from multiprocessing.pool import ThreadPool
import os
import sys
import time

from owsrequest import request
import pandas as pd
import sqlalchemy

from collaborator import config
from collaborator.api import app
from collaborator.constants import features, service_name
from collaborator.constants.collaborator import CollaboratorType
from collaborator.constants.split_type import SplitTypeId
from collaborator.models.rds.collaborator_persister import CollaboratorPersister
from collaborator.models.rds.split_persister import SplitPersister
from collaborator.utils.typing import ACCOUNT_TYPE_VENDOR, Account, User
from scripts import script_util

COLLABORATOR_NAME = "Collaborator Name"
COLLABORATOR_SPLIT = "Collaborator Split"
COLLABORATOR_SPLIT_TYPE = "Split Type"
COLLABORATOR_ID = "Collaborator ID"
IDENTIFIER = "tuid"
VENDOR_ID = "Vendor ID"

XLSX_EXTENSION = ".xlsx"
UPC = "Release UPC"


def _read_file_and_convert_to_json(file_path, num_rows):
    """Read an excel spreadsheet and convert it to json.

    Args:
        file_path (str): The location of the file
        num_rows (int): number of rows to process
    Returns:
        dict with the parsed file
    """
    _, file_extension = os.path.splitext(file_path)

    nrows = None if num_rows <= 0 else num_rows
    if file_extension == XLSX_EXTENSION:
        df = pd.read_excel(file_path, engine="openpyxl", nrows=nrows)
        return df.where(pd.notnull(df), None).to_dict()

    return pd.read_csv(file_path, nrows=nrows).to_dict()


def _get_products_from_upcs(upcs: set) -> dict:
    """Get product information for a given set of UPCs.

    Args:
        upcs (set): The set of product UPCs
    Returns:
        dict: Mapping of UPC to product information.
    """
    snowflake_conn = script_util.snowflake_connection()

    select_sql = f"""
    SELECT
        p.upc,
        p.release_id,
        t.id
    FROM ORCHARD_APP_REPORTING_V2.ART_RELATIONS_PROD_ART_RELATIONS.RELEASES p
    LEFT JOIN ORCHARD_APP_REPORTING_V2.ART_RELATIONS_PROD_ART_RELATIONS.TRACK t on t.upc = p.upc
    WHERE p.upc IN ({", ".join(f"'{upc}'" for upc in upcs)})
    """

    results = snowflake_conn.execute(sqlalchemy.text(select_sql)).mappings().all()

    tracks_by_upc: dict[str, list[int]] = {}
    for result in results:
        upc = result["upc"]
        tracks_by_upc.setdefault(upc, []).append(result["id"])

    return tracks_by_upc


def _get_existing_splits_from_tuids(tuids: set) -> dict:
    """Get existing track information for a given set of TUIDs.

    Args:
        tuids (set): The set of track TUIDs
    Returns:
        dict: Mapping of TUID to existing track information.
    """
    return SplitPersister.get_for_identifiers(
        split_type_id=SplitTypeId.TRACK, identifiers=tuids
    )


def _create_collaborators(data, ticket_id, currency):
    """Create collaborators if needed.

    Args:
        data (dict): Dictionary of arrays with all information.
        ticket_id (str): The ticket ID to use for created collaborators.
        currency (str): The currency to use for created collaborators.

    Returns:
        None, but updates the input data dictionary in-place with created collaborator IDs.
    """
    total_rows = len(data["Vendor ID"])
    logging.info(f"Found : {total_rows} rows in the file")

    id_name_vendor_map = {
        (
            (
                int(row[COLLABORATOR_ID])
                if not pd.isna(row.get(COLLABORATOR_ID))
                else None
            ),
            (
                str(row[COLLABORATOR_NAME]).strip()
                if not pd.isna(row.get(COLLABORATOR_NAME))
                else None
            ),
            int(row[VENDOR_ID]) if not pd.isna(row.get(VENDOR_ID)) else None,
        )
        for row in [
            {key: data[key][row_index] for key in data.keys()}
            for row_index in range(total_rows)
        ]
    }

    created_collaborators = []
    for collaborator_id, collaborator_name, vendor_id in id_name_vendor_map:
        if not collaborator_id:
            created_collaborator = _create_collaborator(
                {
                    "id": collaborator_id,
                    "collaborator_name": collaborator_name,
                    "account": Account(type=ACCOUNT_TYPE_VENDOR, id=vendor_id),
                    "collaborator_type": CollaboratorType.COLLABORATOR,
                    "currency": currency,
                    "subaccount_id": None,
                    "participant_id": None,
                    "description": None,
                    "internal_id": None,
                    "performance_rights": True,
                    "created_by": ticket_id,
                }
            )
            if created_collaborator:
                created_collaborators.append(created_collaborator)
    for row_index in range(total_rows):
        row = {key: data[key][row_index] for key in data.keys()}
        row_collaborator_data = {
            "collaborator_id": (
                int(row[COLLABORATOR_ID])
                if not pd.isna(row.get(COLLABORATOR_ID))
                else None
            ),
            "collaborator_name": (
                str(row[COLLABORATOR_NAME]).strip()
                if not pd.isna(row.get(COLLABORATOR_NAME))
                else None
            ),
            "vendor_id": (
                int(row[VENDOR_ID]) if not pd.isna(row.get(VENDOR_ID)) else None
            ),
        }
        for created_collaborator in created_collaborators:
            if (
                created_collaborator["name"]
                == row_collaborator_data["collaborator_name"]
                and created_collaborator["vendor_id"]
                == row_collaborator_data["vendor_id"]
            ):
                data[COLLABORATOR_ID][row_index] = created_collaborator["id"]
                break
        if pd.isna(data[COLLABORATOR_ID][row_index]):
            error_msg = (
                f"No collaborator - row {row_index} with data: {row_collaborator_data}"
            )
            logging.warning(error_msg)
            raise ValueError(error_msg)


def _get_product_from_upc(upc: str) -> dict:
    """Get product information for a given UPC.

    Args:
        upc (str): The product UPC
    Returns:
        dict: Product information.
    """
    logging.info(f"  > Finding product for UPC: {upc}")

    response = request.process(
        application=config.SERVICE_NAME,
        service_name=service_name.OWS_PRODUCT,
        path=f"/product/upc/{upc}",
        environment=environment,
        method="GET",
    )

    return response.json()


def _get_track_ids_for_product_id(product_id: int) -> list:
    """Get track IDs associated with a product ID.

    Args:
        product_id (int): The product identifier
    Returns:
        list: List of tracks ID:product ID mappings.
    """
    logging.info(f"     > Finding track IDs for product_ID: {product_id}")

    response = request.process(
        application=config.SERVICE_NAME,
        service_name=service_name.OWS_TRACK,
        path=f"/product/{product_id}/tracks",
        environment=environment,
        method="GET",
    )

    tracks = response.json()
    return [
        {"id": track["tuid"], "product_id": product_id}
        for track in tracks.get("items", [])
    ]


def _has_direct_payments(account_id: int) -> bool:
    """Check if account has direct payments feature enabled."""
    res = request.process(
        application=config.SERVICE_NAME,
        service_name=service_name.OWS_ACCOUNT,
        path=f"/vendor/{account_id}/features",
        environment=environment,
        method="GET",
    )

    feature_controls = res.json().get("items", [])
    return (
        len(
            [
                fc
                for fc in feature_controls
                if int(fc["feature_id"]) == features.DIRECT_PAYMENTS_FEATURE_CONTROL
            ]
        )
        > 0
    )


def _replace_splits_for_tracks(track_splits, ticket_id):
    """Replace splits for given track / collaborators specified.

    Args:
        splits (list): list with splits.
    """
    try:
        splits_by_vendor_and_track = {}
        for vendor_track_splits in track_splits:
            for key in vendor_track_splits:
                vendor_id, identifier = key
                splits_by_vendor_and_track.setdefault(vendor_id, {})[identifier] = (
                    vendor_track_splits[key]
                )

        user = User(type="script", id=ticket_id)
        with app.app_context():
            for vendor_id in splits_by_vendor_and_track:
                split_data = {
                    "vendor_id": vendor_id,
                    "tracks": [
                        {
                            "tuid": identifier,
                            "splits": splits_by_vendor_and_track[vendor_id][identifier],
                        }
                        for identifier in splits_by_vendor_and_track[vendor_id]
                    ],
                }

                has_direct_payments = _has_direct_payments(vendor_id)
                if has_direct_payments:
                    split_data["dp_splits_agreed"] = True

                SplitPersister.replace_track_splits(
                    split_data=split_data,
                    user=user,
                    has_direct_payments=has_direct_payments,
                )
    except Exception as exc:
        logging.exception("Error during split replacement setup")
        raise exc


def _create_collaborator(collaborator: dict):
    """Create collaborators.

    Args:
        collaborator (dict): with collaborator payload
    Returns:
        dict with created or existing collaborator
    """
    try:
        with app.app_context():
            collab_id = collaborator.pop("id", None)

            if collab_id and not math.isnan(collab_id):
                collab_results = CollaboratorPersister.get_by_ids(
                    collaborator_ids=[collab_id], throw_if_not_found=False
                )
                existing = next(iter(collab_results), None)
                if existing:
                    return existing

            existing = CollaboratorPersister.get_by_name(**collaborator)
            if existing:
                return existing

            response = CollaboratorPersister.create_collaborator(**collaborator)
            created_id = response.message["id"]
            logging.info(f"Collaborator with {created_id} created successfully")
            return response.message
    except Exception as exc:
        logging.info("Error during collaborator creation")
        logging.info(f"Exception: {str(exc)}")
        return None


def _build_splits_for_single_track(
    key,
    data,
    source,
    overwrite_existing_splits,
    environment,
    ticket_id,
    existing_splits,
):
    """Build splits for a single track."""
    _, identifier = key

    ingested_splits = [
        {
            "identifier": identifier,
            "split_rate": split[COLLABORATOR_SPLIT],
            "split_type_id": SplitTypeId.TRACK,
            "collaborator_id": split[COLLABORATOR_ID],
            "rate_type": split[COLLABORATOR_SPLIT_TYPE].upper(),
            "source": source,
            "user": None,
            "created_by": ticket_id,
        }
        for split in data[key]
    ]

    existing_splits_for_track = [
        split for split in existing_splits if split["identifier"] == identifier
    ]

    new_splits = []
    for ingested_split in ingested_splits:
        if not any(
            ingested_split["collaborator_id"] == existing_split["collaborator_id"]
            for existing_split in existing_splits_for_track
        ):
            new_splits.append(ingested_split)

    updated_existing_splits = []
    updated_splits = 0
    for existing_split in existing_splits_for_track:
        for ingested_split in ingested_splits:
            if (
                overwrite_existing_splits
                and ingested_split["collaborator_id"]
                == existing_split["collaborator_id"]
            ):
                updated_existing_splits.append(
                    {
                        "id": existing_split["id"],
                        "split_rate": ingested_split["split_rate"],
                        "split_type_id": ingested_split["split_type_id"],
                        "collaborator_id": ingested_split["collaborator_id"],
                        "rate_type": ingested_split["rate_type"],
                        "source": source,
                        "user": None,
                        "updated_by": ticket_id,
                    }
                )
                updated_splits += 1
            else:
                updated_existing_splits.append(existing_split)

    logging.info(
        f" > Splits actions for {identifier} :  "
        f"create {len(new_splits)} splits.  |  update {updated_splits} splits."
    )

    return {key: [*updated_existing_splits, *new_splits]}


def _build_track_splits_from_product_splits(data):
    """Process splits at product level, turning them into track-level splits.

    Args:
        data (dict): Dictionary of rows sorted into buckets by collab_id.
    """
    product_level_upcs = {
        row[UPC] for row in data if row.get(UPC) and not row.get(IDENTIFIER)
    }
    logging.info(f"Found {len(product_level_upcs)} product-level splits to process.")
    tracks_for_product_level_splits_by_upc = _get_products_from_upcs(product_level_upcs)

    final_rows = []
    for row in data:
        tuid = (
            row[IDENTIFIER]
            if IDENTIFIER in row and not pd.isna(row[IDENTIFIER])
            else None
        )
        if tuid:
            final_rows.append(row)
        else:
            upc = row[UPC] if UPC in row and not pd.isna(row[UPC]) else None
            if upc:
                product_tracks_id = tracks_for_product_level_splits_by_upc.get(
                    upc, None
                )
                for track_id in product_tracks_id:
                    new_row = row.copy()
                    new_row[IDENTIFIER] = track_id
                    final_rows.append(new_row)
            else:
                logging.error("No valid track identifier or UPC found in the row.")
                raise ValueError(
                    f"No valid track identifier or UPC found in the row: {row}"
                )
    return final_rows


def run_splits_ingest(
    data,
    ticket_id,
    source,
    environment,
    overwrite_existing_splits=True,
    should_create_collaborators=False,
    currency=None,
):
    """Process payload for splits ingest.

    Args:
        data (dict): Dictionary of arrays with all information.
        source (str): The source of split information.
        environment (str): The operating environment, e.g. QA, PROD.
    """
    total_rows = len(data["Vendor ID"])
    logging.info(f"Found : {total_rows} rows in the file")

    # Create collaborators if needed
    if should_create_collaborators:
        logging.info("Creating collaborators if needed")
        _create_collaborators(data, ticket_id, currency)
        logging.info("Collaborator creation completed.")

    # Process product-level splits into track-level splits if any
    logging.info("Processing product-level splits into track-level splits")
    track_level_splits = _build_track_splits_from_product_splits(
        [
            {key: data[key][row_index] for key in data.keys()}
            for row_index in range(total_rows)
        ]
    )
    logging.info(
        f"Product-level splits processed."
        f"{len(track_level_splits)} product-level splits turned into track-level."
    )

    unique_vendors = set(data[VENDOR_ID].values())
    unique_identifiers = {int(row[IDENTIFIER]) for row in track_level_splits}
    logging.info(
        "Processing rows for:"
        f"\n · {len(unique_vendors)} vendors "
        f"\n · {len(unique_identifiers)} identifiers"
    )

    # Create buckets for every combination of vendor_id and identifier
    vendor_track_buckets = {}
    for row in track_level_splits:
        vendor_track_buckets.setdefault((row[VENDOR_ID], row[IDENTIFIER]), []).append(
            row
        )

    existing_splits = _get_existing_splits_from_tuids(unique_identifiers)

    with ThreadPool(10) as pool:
        processed_vendor_track_splits = pool.map(
            partial(
                _build_splits_for_single_track,
                data=vendor_track_buckets,
                source=source,
                overwrite_existing_splits=overwrite_existing_splits,
                environment=environment,
                ticket_id=ticket_id,
                existing_splits=existing_splits,
            ),
            vendor_track_buckets,
        )

    # Persist splits per vendor and track
    _replace_splits_for_tracks(processed_vendor_track_splits, ticket_id)


if __name__ == "__main__":
    start = time.time()
    logging.basicConfig(level=logging.INFO)

    environment = os.environ.get("Environment", "dev")
    logging.info(f"Got environment: {environment}")

    source = os.environ.get("BUILD_TAG", "collaborator_ingest_script")
    logging.info(f"Got build tag: {source}")

    file_path = sys.argv[1]
    logging.info(f"The file path is: {file_path}")

    ticket_id = os.environ.get("TICKET_ID")
    if ticket_id is None:
        raise Exception("TICKET_ID must be provided.")

    logging.info(f"Got TICKET ID: {ticket_id}")

    num_rows = int(os.environ.get("NUMBER_OF_ROWS_TO_PROCESS", 0))
    logging.info(f"Got number of rows to process: {num_rows}")

    overwrite_existing_splits = bool(os.environ.get("OVERWRITE_EXISTING_SPLITS", True))
    logging.info(f"Should overwrite existing splits: {overwrite_existing_splits}")

    should_create_collaborators = bool(os.environ.get("CREATE_COLLABORATORS", False))
    logging.info(f"Should create collaborators: {should_create_collaborators}")

    currency = os.environ.get("VENDOR_CURRENCY")
    if currency is None:
        raise Exception("VENDOR_CURRENCY must be provided.")
    logging.info(f"currency: {currency}")

    logging.info("Process excel spreadsheet and transform it to dictionary")
    data = _read_file_and_convert_to_json(file_path, num_rows)

    logging.info("Run collaborator and splits ingest")
    run_splits_ingest(
        data,
        ticket_id,
        source,
        environment,
        overwrite_existing_splits,
        should_create_collaborators,
        currency,
    )

    elapsed = time.time() - start
    logging.info("Ingest finished in {0:.2f} seconds.".format(elapsed))