"""
SPIKE Audit Vendor, Subaccount, CompanyBrand nodes in Neo4j - https://theorchard.atlassian.net/browse/PP-434

# Queries

## Company Brand

```cypher
match(cb:CompanyBrand) return cb.uuid as company_brand_uuid;
```


## Subaccount

```cypher
// Use this query for 9 batches, increasing `skip` by 10k for each query.
match(s:Subaccount)
return s.id, s.uuid order by s.id asc
skip 0 limit 10000;
```

## Vendor

```cypher
// Use this query for 9 batches, increasing `skip` by 10k for each query.
match(v:Vendor)
return v.id, v.uuid order by v.id asc
skip 0 limit 10000;
```
"""

import uuid
import pandas as pd
import os
import sys
import argparse


def invalid_id_test(_id: str):
    try:
        int(_id)
        return False
    except Exception as ex:
        return True


def invalid_uuid_test(_uuid: str):
    try:
        uuid.UUID(_uuid, version=4)
        return False
    except Exception as ex:
        return True


def validate_uuids(df: pd.DataFrame, uuid_field: str):
    # remove extra double-quotes from the neo4j export.
    df[uuid_field] = df[uuid_field].apply(lambda v: v.replace('"', ''))
    # Find invalid values for column
    has_invalid_uuid = df[uuid_field].apply(invalid_uuid_test)

    invalid_uuid_count = has_invalid_uuid.sum()
    if invalid_uuid_count:
        print("\n--- Rows with invalid UUIDs ---")
        print(df[has_invalid_uuid].to_json(indent=2))
        print("---------")

    return invalid_uuid_count


def validate_ids(df: pd.DataFrame, id_field: str):
    # remove extra double-quotes from the neo4j export.
    df[id_field] = df[id_field].apply(lambda v: v.replace('"', ''))
    # Find invalid values for column
    has_invalid_id = df[id_field].apply(invalid_id_test)

    invalid_id_count = has_invalid_id.sum()
    if invalid_id_count:
        print("\n--- Rows with invalid IDs ---")
        print(df[has_invalid_id].to_json(indent=2))
        print("------------\n")

    return invalid_id_count


def audit(fn: str, uuid_field: str, id_field: str = None):
    df = pd.read_csv(fn, dtype='string')

    # convert missing/empty column values to "<missing>"
    df = df.replace(pd.NaT, "<missing>").replace(r'^\s*$', "<missing>", regex=True)

    invalid_id_count, invalid_uuid_count = (0, 0)

    invalid_uuid_count = validate_uuids(df, uuid_field)
    if id_field:
        invalid_id_count = validate_ids(df, id_field)

    print(f"[{os.path.basename(fn)}] {invalid_uuid_count} invalid uuids, {invalid_id_count} invalid ids")


def audit_main(datadir: str):
    vendor_dir = os.path.join(datadir, 'vendor')
    subaccount_dir = os.path.join(datadir, 'subaccount')
    cb_dir = os.path.join(datadir, 'company_brand')

    # vendor
    for fn in sorted(os.listdir(vendor_dir)):
        audit(os.path.join(vendor_dir, fn), uuid_field='v.uuid', id_field="v.id")

    print("\n\n")

    # subaccount
    for fn in sorted(os.listdir(subaccount_dir)):
        audit(os.path.join(subaccount_dir, fn), uuid_field='s.uuid', id_field="s.id")

    print("\n\n")

    # company brand
    for fn in sorted(os.listdir(cb_dir)):
        audit(os.path.join(cb_dir, fn), uuid_field='company_brand_uuid')


def parse_args():
    parser = argparse.ArgumentParser(description="Validate ID and UUID column values for CSV files found in './data'")
    parser.add_argument("--test", default=False, action='store_true', help="Perform a self-check using ./test.csv")
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()

    if args.test:
        print("\n\n ====== TEST ======")
        audit(fn='./test.csv', uuid_field='uuid', id_field='id')
    else:
        if not os.path.isdir('./data'):
            print(f"Did not find '{os.path.abspath('./data')}'. Please run `make data`.")
            sys.exit(1)

        print("\n\n ====== QA ======")
        audit_main(datadir='./data/qa')

        print("\n\n ====== PROD ======")
        audit_main(datadir='./data/prod')