import argparse
from pathlib import Path
from typing import List

from ee.db import ObjDb
from ee.part import Part, load_db, save_db, PartDb
from ee.project import Project


class OrderPart(object):
    def __init__(self, order_part: Part, part: Part):
        self.order_part = order_part
        self.part = part


def uri_fn(part: Part):
    return part.uri


def split_parts_by_supplier(project: Project, order_file: Path, part_dbs: List[Path], out_dir: Path):
    parts: ObjDb[Part] = ObjDb[Part]()
    part_by_uri = parts.add_unique_index("uri", uri_fn)

    for part_db in part_dbs:
        for xml in load_db(part_db).iterparts():
            parts.add(Part(xml))

    order_parts: ObjDb[OrderPart] = ObjDb()
    supplier_idx = order_parts.add_index("supplier", lambda op: op.part.supplier)
    for xml in load_db(order_file).iterparts():
        order_part = Part(xml)
        part = part_by_uri.get_single(order_part.get_exactly_one_part_reference().part_uriProp)
        order_parts.add(OrderPart(order_part, part))

    for supplier, parts_for_supplier in supplier_idx.items():
        desc = project.get_supplier_by_uri(supplier)

        print("{}: {}".format(desc.name, len(parts)))
        # supplier_db: ObjDb[Part] = ObjDb[Part]()
        # supplier_db.add_unique_index("uri", uri_fn)

        supplier_descriptor = project.get_supplier_by_uri(supplier)

        db = PartDb()
        for part_for_supplier in parts_for_supplier:
            db.add_entry(part_for_supplier.part.underlying, False)

        save_db(out_dir / "{}.xml".format(supplier_descriptor.key), db, sort=True)


parser = argparse.ArgumentParser()

parser.add_argument("--parts",
                    required=True,
                    metavar="PART DB")

parser.add_argument("--part-db",
                    nargs="*",
                    required=True,
                    metavar="PART DB")

parser.add_argument("--out-dir",
                    metavar="DIR FOR PART DBS")

args = parser.parse_args()

part_db_dirs = [Path(part_db) for part_db in args.part_db]

split_parts_by_supplier(Project.load(), Path(args.parts), part_db_dirs, Path(args.out_dir))