import asyncio
import logging
import re
import urllib
from dataclasses import dataclass, field
from time import sleep
from typing import Dict, Iterable, List, Union

from okta.client import Client as OktaClient
from okta.models import Group, GroupProfile, User, UserProfile, UserStatus

from datahub.configuration import ConfigModel
from datahub.ingestion.api.common import PipelineContext
from datahub.ingestion.api.source import Source, SourceReport
from datahub.ingestion.api.workunit import MetadataWorkUnit
from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import (
    CorpGroupSnapshot,
    CorpUserSnapshot,
)
from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
from datahub.metadata.schema_classes import (  # GroupMembershipClass,
    CorpGroupInfoClass,
    CorpUserInfoClass,
    GroupMembershipClass,
)

logger = logging.getLogger(__name__)


class OktaConfig(ConfigModel):

    # Required: Domain of the Okta deployment. Example: dev-33231928.okta.com
    okta_domain = "dev-44231988.okta.com"
    # Required: An API token generated from Okta.
    okta_api_token = "00be4R_M2MzDqXawbWgfKGpKee0kuEOfX1RCQSRx00"

    # Optional: Whether to ingest users, groups, or both.
    ingest_users: bool = True
    ingest_groups: bool = True
    ingest_group_membership: bool = True

    # Optional: Customize the mapping to DataHub Username from an attribute appearing in the Okta User
    # profile. Reference: https://developer.okta.com/docs/reference/api/users/
    okta_profile_to_username_attr: str = "login"
    okta_profile_to_username_regex: str = "([^@]+)"

    # Optional: Customize the mapping to DataHub Group from an attribute appearing in the Okta Group
    # profile. Reference: https://developer.okta.com/docs/reference/api/groups/
    okta_profile_to_group_name_attr: str = "name"
    okta_profile_to_group_name_regex: str = "(.*)"

    # Optional: Include deprovisioned or suspended Okta users in the ingestion.
    include_deprovisioned_users = False
    include_suspended_users = False

    # Optional: Page size for reading groups and users from Okta API.
    page_size = 100

    # Optional: Set the delay for fetching batches of entities from Okta. Okta has rate limiting in place.
    delay_seconds = 0.01


@dataclass
class OktaSourceReport(SourceReport):
    filtered: List[str] = field(default_factory=list)

    def report_filtered(self, name: str) -> None:
        self.filtered.append(name)


#
# Source Connector that extracts Users, Groups, and Group Membership in batch using the Okta Python SDK.
#
# Validated against Okta API Versions:
#   - 2021.07.2
#
# Validated against load:
#   - User Count: 1000
#   - Group Count: 100
#   - Group Membership Edges: 1000 (1 per User)
#   - Run Time (Wall Clock): 2min 7sec
#
class OktaSource(Source):
    """Ingest Okta Users & Groups into Datahub"""

    @classmethod
    def create(cls, config_dict, ctx):
        config = OktaConfig.parse_obj(config_dict)
        return cls(config, ctx)

    def __init__(self, config: OktaConfig, ctx: PipelineContext):