import factory
from wbfdm.factories.instruments import InstrumentFactory

from wbportfolio.models import ProductGroup, ProductGroupRepresentant


class ProductGroupFactory(InstrumentFactory):
    class Meta:
        model = ProductGroup

    type = factory.Faker("random_element", elements=[x[0] for x in ProductGroup.ProductGroupType.choices])
    category = factory.Faker("random_element", elements=[x[0] for x in ProductGroup.ProductGroupCategory.choices])
    umbrella = factory.Faker("company")
    management_company = factory.SubFactory("wbcore.contrib.directory.factories.entries.CompanyFactory")
    depositary = factory.SubFactory("wbcore.contrib.directory.factories.entries.CompanyFactory")
    transfer_agent = factory.SubFactory("wbcore.contrib.directory.factories.entries.CompanyFactory")
    administrator = factory.SubFactory("wbcore.contrib.directory.factories.entries.CompanyFactory")
    investment_manager = factory.SubFactory("wbcore.contrib.directory.factories.entries.CompanyFactory")
    auditor = factory.SubFactory("wbcore.contrib.directory.factories.entries.CompanyFactory")
    paying_agent = factory.SubFactory("wbcore.contrib.directory.factories.entries.CompanyFactory")


class ProductGroupRepresentantFactory(factory.django.DjangoModelFactory):
    class Meta:
        model = ProductGroupRepresentant

    product_group = factory.SubFactory(ProductGroupFactory)
    representant = factory.SubFactory("wbcore.contrib.directory.factories.entries.CompanyFactory")
    country = factory.SubFactory("wbcore.contrib.geography.factories.CountryFactory")
