import logging
import uuid
from enum import Enum
from pathlib import Path
from typing import List, Optional, Union

import yaml
from pydantic import Field, validator
from ruamel.yaml import YAML
from typing_extensions import Literal

from datahub.api.entities.forms.forms_graphql_constants import (
    CREATE_DYNAMIC_FORM_ASSIGNMENT,
    FIELD_FILTER_TEMPLATE,
    UPLOAD_ENTITIES_FOR_FORMS,
)
from datahub.configuration.common import ConfigModel
from datahub.emitter.mce_builder import (
    make_container_urn,
    make_data_platform_urn,
    make_domain_urn,
    make_group_urn,
    make_tag_urn,
    make_term_urn,
    make_user_urn,
)
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
from datahub.ingestion.graph.config import ClientMode
from datahub.metadata.schema_classes import (
    FormActorAssignmentClass,
    FormInfoClass,
    FormPromptClass,
    OwnerClass,
    OwnershipClass,
    OwnershipTypeClass,
    StructuredPropertyParamsClass,
)
from datahub.utilities.urns.urn import Urn

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

FILTER_CRITERION_TYPES = "_entityType"
FILTER_CRITERION_SUB_TYPES = "typeNames.keyword"
FILTER_CRITERION_PLATFORMS = "platform.keyword"
FILTER_CRITERION_PLATFORM_INSTANCES = "dataPlatformInstance.keyword"
FILTER_CRITERION_DOMAINS = "domains.keyword"
FILTER_CRITERION_CONTAINERS = "container.keyword"
FILTER_CRITERION_OWNERS = "owners.keyword"
FILTER_CRITERION_TAGS = "tags.keyword"
FILTER_CRITERION_GLOSSARY_TERMS = "glossaryTerms.keyword"


class PromptType(Enum):
    STRUCTURED_PROPERTY = "STRUCTURED_PROPERTY"
    FIELDS_STRUCTURED_PROPERTY = "FIELDS_STRUCTURED_PROPERTY"

    @classmethod
    def has_value(cls, value):
        return value in cls._value2member_map_


class Prompt(ConfigModel):
    id: Optional[str] = None
    title: str
    description: Optional[str] = None
    type: str
    structured_property_id: Optional[str] = None
    structured_property_urn: Optional[str] = Field(default=None, validate_default=True)
    required: Optional[bool] = None

    @validator("structured_property_urn", pre=True, always=True)
    def structured_property_urn_must_be_present(cls, v, values):
        if not v and values.get("structured_property_id"):
            return Urn.make_structured_property_urn(values["structured_property_id"])
        return v


class FormType(Enum):
    VERIFICATION = "VERIFICATION"
    DOCUMENTATION = "COMPLETION"

    @classmethod
    def has_value(cls, value):
        return value in cls._value2member_map_


class Filters(ConfigModel):
    types: Optional[List[str]] = None
    sub_types: Optional[List[str]] = None
    platforms: Optional[List[str]] = None
    platform_instances: Optional[List[str]] = None
    domains: Optional[List[str]] = None
    containers: Optional[List[str]] = None
    owners: Optional[List[str]] = None
    tags: Optional[List[str]] = None
    terms: Optional[List[str]] = None


class Entities(ConfigModel):
    urns: Optional[List[str]] = None
    filters: Optional[Filters] = None


class Actors(ConfigModel):
    owners: Optional[bool] = None
    users: Optional[List[str]] = None  # can be user IDs or urns
    groups: Optional[List[str]] = None  # can be group IDs or urns


class Forms(ConfigModel):
    id: Optional[str] = None
    urn: Optional[str] = Field(default=None, validate_default=True)
    name: str
    description: Optional[str] = None
    prompts: List[Prompt] = []
    type: Optional[str] = None
    version: Optional[Literal[1]] = None
    entities: Optional[Entities] = None
    owners: Optional[List[str]] = None  # can be user IDs or urns
    group_owners: Optional[List[str]] = None  # can be group IDs or urns
    actors: Optional[Actors] = None

    @validator("urn", pre=True, always=True)
    def urn_must_be_present(cls, v, values):
        if not v:
            if values.get("id") is None:
                raise ValueError("Form id must be present if urn is not")
            return f"urn:li:form:{values['id']}"
        return v

    @staticmethod
    def create(file: str) -> None:
        emitter: DataHubGraph

        with get_default_graph(ClientMode.CLI) as emitter, open(file) as fp:
            forms: List[dict] = yaml.safe_load(fp)
            for form_raw in forms:
                form = Forms.parse_obj(form_raw)

                try:
                    if not FormType.has_value(form.type):
                        logger.error(
                            f"Form type {form.type} does not exist. Please try again with a valid type."
                        )

                    mcp = MetadataChangeProposalWrapper(
                        entityUrn=form.urn,
                        aspect=FormInfoClass(
                            name=form.name,
                            description=form.description,
                            prompts=form.validate_prompts(emitter),
                            type=form.type,
                            actors=form.create_form_actors(form.actors),
                        ),
                    )
                    emitter.emit_mcp(mcp)

                    logger.info(f"Created form {form.urn}")

                    if form.owners or form.group_owners:
                        form.add_owners(emitter)

                    if form.entities:
                        if form.entities.urns:
                            # Associate specific entities with a form
                            form.upload_entities_for_form(emitter)

                        if form.entities.filters:
                            # Associate groups of entities with a form based on filters
                            form.create_form_filters(emitter)

                except Exception as e:
                    logger.error(e)
                    return

    def validate_prompts(self, emitter: DataHubGraph) -> List[FormPromptClass]:
        prompts = []
        if self.prompts:
            for prompt in self.prompts:
                if not prompt.id:
                    prompt.id = str(uuid.uuid4())
                    logger.warning(
                        f"Prompt id not provided. Setting prompt id to {prompt.id}"
                    )
                if prompt.structured_property_urn:
                    structured_property_urn = prompt.structured_property_urn
                    if emitter.exists(structured_property_urn):
                        prompt.structured_property_urn = structured_property_urn
                    else:
                        raise Exception(
                            f"Structured property {structured_property_urn} does not exist. Unable to create form."
                        )
                elif (
                    prompt.type
                    in (
                        PromptType.STRUCTURED_PROPERTY.value,
                        PromptType.FIELDS_STRUCTURED_PROPERTY.value,
                    )
                    and not prompt.structured_property_urn
                ):
                    raise Exception(
                        f"Prompt type is {prompt.type} but no structured properties exist. Unable to create form."
                    )
                if (
                    prompt.type == PromptType.FIELDS_STRUCTURED_PROPERTY.value
                    and prompt.required
                ):
                    raise Exception(
                        "Schema field prompts cannot be marked as required. Ensure these prompts are not required."
                    )

                prompts.append(
                    FormPromptClass(
                        id=prompt.id,
                        title=prompt.title,
                        description=prompt.description,
                        type=prompt.type,
                        structuredPropertyParams=(
                            StructuredPropertyParamsClass(
                                urn=prompt.structured_property_urn
                            )
                            if prompt.structured_property_urn
                            else None
                        ),
                        required=prompt.required,
                    )
                )
        else:
            logger.warning(f"No prompts exist on form {self.urn}. Is that intended?")

        return prompts

    def create_form_actors(
        self, actors: Optional[Actors] = None
    ) -> Union[None, FormActorAssignmentClass]:
        if actors is None:
            return None

        users = None
        if actors.users is not None:
            users = Forms.format_users(actors.users)

        groups = None
        if actors.groups is not None:
            groups = Forms.format_groups(actors.groups)

        return FormActorAssignmentClass(
            owners=actors.owners, users=users, groups=groups
        )

    def upload_entities_for_form(self, emitter: DataHubGraph) -> Union[None, Exception]:
        if self.entities and self.entities.urns:
            formatted_entity_urns = ", ".join(
                [f'"{value}"' for value in self.entities.urns]
            )
            query = UPLOAD_ENTITIES_FOR_FORMS.format(
                form_urn=self.urn, entity_urns=formatted_entity_urns
            )
            result = emitter.execute_graphql(query=query)
            if not result:
                return Exception(f"Could not bulk upload entities for form {self.urn}.")

        return None

    def create_form_filters(self, emitter: DataHubGraph) -> Union[None, Exception]:
        filters_raw = []
        # Loop through each entity and assign a filter for it
        if self.entities and self.entities.filters:
            filters = self.entities.filters

            if filters.types:
                filters_raw.append(
                    Forms.format_form_filter(FILTER_CRITERION_TYPES, filters.types)
                )

            if filters.sub_types:
                filters_raw.append(
                    Forms.format_form_filter(
                        FILTER_CRITERION_SUB_TYPES, filters.sub_types
                    )
                )

            if filters.platforms:
                urns = [
                    make_data_platform_urn(platform) for platform in filters.platforms
                ]
                filters_raw.append(
                    Forms.format_form_filter(FILTER_CRITERION_PLATFORMS, urns)
                )

            if filters.platform_instances:
                urns = []
                for platform_instance in filters.platform_instances:
                    platform_instance_urn = Forms.validate_platform_instance_urn(
                        platform_instance
                    )
                    if platform_instance_urn:
                        urns.append(platform_instance_urn)
                filters_raw.append(
                    Forms.format_form_filter(FILTER_CRITERION_PLATFORM_INSTANCES, urns)
                )

            if filters.domains:
                urns = [make_domain_urn(domain) for domain in filters.domains]
                filters_raw.append(
                    Forms.format_form_filter(FILTER_CRITERION_DOMAINS, urns)
                )

            if filters.containers:
                urns = [
                    make_container_urn(container) for container in filters.containers
                ]
                filters_raw.append(
                    Forms.format_form_filter(FILTER_CRITERION_CONTAINERS, urns)
                )

            if filters.owners:
                urns = [make_user_urn(owner) for owner in filters.owners]
                filters_raw.append(
                    Forms.format_form_filter(FILTER_CRITERION_OWNERS, urns)
                )

            if filters.tags:
                urns = [make_tag_urn(tag) for tag in filters.tags]
                filters_raw.append(
                    Forms.format_form_filter(FILTER_CRITERION_TAGS, urns)
                )

            if filters.terms:
                urns = [make_term_urn(term) for term in filters.terms]
                filters_raw.append(
                    Forms.format_form_filter(FILTER_CRITERION_GLOSSARY_TERMS, urns)
                )

        filters_str = ", ".join(item for item in filters_raw)
        result = emitter.execute_graphql(
            query=CREATE_DYNAMIC_FORM_ASSIGNMENT.format(
                form_urn=self.urn, filters=filters_str
            )
        )

        if not result:
            return Exception(
                f"Could not bulk upload urns or filters for form {self.urn}."
            )

        return None

    def add_owners(self, emitter: DataHubGraph) -> Union[None, Exception]:
        owner_urns: List[str] = []
        if self.owners:
            owner_urns += Forms.format_users(self.owners)
        if self.group_owners:
            owner_urns += Forms.format_groups(self.group_owners)

        ownership = OwnershipClass(
            owners=[
                OwnerClass(owner=urn, type=OwnershipTypeClass.TECHNICAL_OWNER)
                for urn in (owner_urns or [])
            ],
        )

        try:
            mcp = MetadataChangeProposalWrapper(entityUrn=self.urn, aspect=ownership)
            emitter.emit_mcp(mcp)
        except Exception as e:
            logger.error(e)

        return None

    @staticmethod
    def format_form_filter(field: str, urns: List[str]) -> str:
        formatted_urns = ", ".join([f'"{urn}"' for urn in urns])
        return FIELD_FILTER_TEMPLATE.format(field=field, values=formatted_urns)

    @staticmethod
    def validate_platform_instance_urn(instance: str) -> Union[str, None]:
        if instance.startswith("urn:li:dataPlatformInstance:"):
            return instance

        logger.warning(
            f"{instance} is not an urn. Unable to create platform instance filter."
        )
        return None

    @staticmethod
    def from_datahub(graph: DataHubGraph, urn: str) -> "Forms":
        form: Optional[FormInfoClass] = graph.get_aspect(urn, FormInfoClass)
        if form is None:
            raise Exception("FormInfo aspect is None. Unable to create form.")
        prompts = []
        for prompt_raw in form.prompts:
            prompts.append(
                Prompt(
                    id=prompt_raw.id,
                    title=prompt_raw.title,
                    description=prompt_raw.description,
                    type=prompt_raw.type,
                    structured_property_urn=(
                        prompt_raw.structuredPropertyParams.urn
                        if prompt_raw.structuredPropertyParams
                        else None
                    ),
                )
            )
        return Forms(
            urn=urn,
            name=form.name,
            description=form.description,
            prompts=prompts,
            type=form.type,
        )

    @staticmethod
    def format_users(users: List[str]) -> List[str]:
        formatted_users: List[str] = []

        for user in users:
            if user.startswith("urn:li:"):
                formatted_users.append(user)
            else:
                formatted_users.append(make_user_urn(user))

        return formatted_users

    @staticmethod
    def format_groups(groups: List[str]) -> List[str]:
        formatted_groups: List[str] = []

        for group in groups:
            if group.startswith("urn:li:"):
                formatted_groups.append(group)
            else:
                formatted_groups.append(make_group_urn(group))

        return formatted_groups

    def to_yaml(
        self,
        file: Path,
    ) -> None:
        with open(file, "w") as fp:
            yaml = YAML(typ="rt")  # default, if not specfied, is 'rt' (round-trip)
            yaml.indent(mapping=2, sequence=4, offset=2)
            yaml.default_flow_style = False
            yaml.dump(self.dict(), fp)
