import json
import logging
from datetime import date, datetime
from enum import Enum
from typing import Any, Callable, Dict, Iterable, List, Optional

import pandas as pd
from pydantic import BaseModel

from acryl_datahub_cloud.datahub_reporting.datahub_dataset import BaseModelRow
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.ingestion.graph.client import DataHubGraph
from datahub.ingestion.graph.filters import SearchFilterRule
from datahub.metadata.schema_classes import (
    DomainPropertiesClass,
    FormInfoClass,
    FormsClass,
    FormTypeClass,
)

logger = logging.getLogger(__name__)


class FormStatus(str, Enum):
    NOT_STARTED = "not_started"
    IN_PROGRESS = "in_progress"
    COMPLETED = "complete"


class QuestionStatus(str, Enum):
    NOT_STARTED = "not_started"
    COMPLETED = "complete"


class FormType(str, Enum):
    DOCUMENTATION = "documentation"
    VERIFICATION = "verification"


class FormReportingRow(BaseModelRow):
    form_urn: str
    form_type: FormType
    form_assigned_date: date
    form_completed_date: Optional[date]
    form_status: FormStatus
    question_id: str
    question_status: QuestionStatus
    question_completed_date: Optional[date]
    assignee_urn: Optional[str]
    asset_urn: str
    platform_urn: str
    platform_instance_urn: str
    domain_urn: Optional[str]
    parent_domain_urn: Optional[str]
    asset_verified: Optional[bool]
    snapshot_date: date


class FormData:
    def get_data(
        self,
        on_asset_scanned: Optional[Callable[[str], Any]] = None,
        on_form_scanned: Optional[Callable[[str], Any]] = None,
    ) -> Iterable[FormReportingRow]:
        raise NotImplementedError

    def to_mcps(self) -> Iterable[MetadataChangeProposalWrapper]:
        raise NotImplementedError


class FormRegistry:
    def __init__(self, graph: DataHubGraph):
        self.registry: Dict[str, Optional[FormInfoClass]] = {}
        self.graph = graph

    def get_form(self, form_id: str) -> Optional[FormInfoClass]:
        if form_id not in self.registry:
            self._load_form(form_id)
        return self.registry[form_id]

    def _load_form(self, form_id: str) -> None:
        form_urn = (
            f"urn:li:form:{form_id}"
            if not form_id.startswith("urn:li:form:")
            else form_id
        )

        form_info = self.graph.get_aspect(form_urn, FormInfoClass)
        self.registry[form_id] = form_info


class DomainRegistry:
    def __init__(self, graph: DataHubGraph):
        self.registry: Dict[str, Optional[DomainPropertiesClass]] = {}
        self.graph = graph

    def get_domain(self, domain_id: str) -> Optional[DomainPropertiesClass]:
        if domain_id not in self.registry:
            self._load_domain(domain_id)
        return self.registry[domain_id]

    def _load_domain(self, domain_id: str) -> None:
        domain_urn = (
            f"urn:li:domain:{domain_id}"
            if not domain_id.startswith("urn:li:domain:")
            else domain_id
        )
        domain_props = self.graph.get_aspect(domain_urn, DomainPropertiesClass)
        self.registry[domain_id] = domain_props


class DataHubFormReportingData(FormData):
    class DataHubDatasetSearchRow(BaseModel):
        urn: str
        owners: List[str] = []
        completedForms: List[str] = []
        incompleteForms: List[str] = []
        completedFormsIncompletePromptIds: List[str] = []
        completedFormsIncompletePromptResponseTimes: List[str] = []
        completedFormsCompletedPromptIds: List[str] = []
        completedFormsCompletedPromptResponseTimes: List[str] = []
        incompleteFormsIncompletePromptIds: List[str] = []
        incompleteFormsIncompletePromptResponseTimes: List[str] = []
        incompleteFormsCompletedPromptIds: List[str] = []
        incompleteFormsCompletedPromptResponseTimes: List[str] = []
        platform: str = ""
        platformInstance: str = ""
        domains: List[str] = []

    def __init__(self, graph: DataHubGraph, allowed_forms: Optional[List[str]] = None):
        self.graph: DataHubGraph = graph
        self.form_registry = FormRegistry(graph)
        self.domain_registry = DomainRegistry(graph)
        self.snapshot_date = datetime.now().date()
        self.allowed_forms = allowed_forms

    def get_dataframe(
        self,
        on_asset_scanned: Callable[[str], Any],
        on_form_scanned: Callable[[str], Any],
    ) -> pd.DataFrame:
        return pd.DataFrame(
            x.dict()
            for x in self.get_data(
                on_asset_scanned=on_asset_scanned, on_form_scanned=on_form_scanned
            )
        )

    def get_form_existence_or_filters(self) -> List[SearchFilterRule]:
        """
        Datasets must either have completedForms or incompleteForms assigned to
        them
        """
        if self.allowed_forms:
            return [
                {
                    "field": "completedForms",
                    "condition": "EQUAL",
                    "values": self.allowed_forms,
                },
                {
                    "field": "incompleteForms",
                    "condition": "EQUAL",
                    "values": self.allowed_forms,
                },
            ]
        else:
            return [
                {
                    "field": "completedForms",
                    "condition": "EXISTS",
                },
                {
                    "field": "incompleteForms",
                    "condition": "EXISTS",
                },
            ]

    def form_assigned_date(
        self, search_row: DataHubDatasetSearchRow
    ) -> Dict[str, date]:

        form_assigned_dates: Dict[str, date] = {}
        forms = self.graph.get_aspect(search_row.urn, FormsClass)
        if not forms:
            return form_assigned_dates
        assert forms, f"Forms aspect not found for {search_row.urn}"
        for incomplete_form in forms.incompleteForms:
            form_assigned_dates[incomplete_form.urn] = datetime.fromtimestamp(
                incomplete_form.created.time / 1000 if incomplete_form.created else 0
            ).date()
        for completed_form in forms.completedForms:
            form_assigned_dates[completed_form.urn] = datetime.fromtimestamp(
                completed_form.created.time / 1000 if completed_form.created else 0
            ).date()
        return form_assigned_dates

    def form_completed_date(
        self, search_row: DataHubDatasetSearchRow
    ) -> Dict[str, date]:
        # get the maximum timestamp from the
        # incompleteFormsCompletedPromptResponseTimes and
        # completedFormsCompletedPromptResponseTimes
        form_completion_dates = {}
        for form in search_row.completedForms:
            form_info = self.form_registry.get_form(form)
            assert form_info, f"Form {form} not found"
            form_prompts = [x.id for x in form_info.prompts]
            completed_prompts_map = {
                prompt_id: response_time
                for prompt_id, response_time in zip(
                    search_row.completedFormsCompletedPromptIds,
                    search_row.completedFormsCompletedPromptResponseTimes,
                )
                if prompt_id in form_prompts
            }
            sorted_dates = sorted(completed_prompts_map.values())
            if sorted_dates:
                form_completion_dates[form] = datetime.fromtimestamp(
                    int(sorted(completed_prompts_map.values())[-1]) / 1000
                ).date()
            else:
                logger.warning(
                    f"Form {form} on asset {search_row.urn} has no completed prompts but is marked as completed"
                )
        return form_completion_dates

    def get_assignees(self, form_info: FormInfoClass, owners: List[str]) -> List[str]:
        assignees = []
        actors = form_info.actors
        if actors.owners:
            # form applies to owners
            assignees.extend(owners)
        if actors.users:
            assignees.extend(actors.users)
        if actors.groups:
            assignees.extend(actors.groups)
        return list(set(assignees))

    def get_data(  # noqa: C901
        self,
        on_asset_scanned: Optional[Callable[[str], Any]] = None,
        on_form_scanned: Optional[Callable[[str], Any]] = None,
    ) -> Iterable[FormReportingRow]:
        extra_fields = [f for f in self.DataHubDatasetSearchRow.__fields__.keys()]
        result = self.graph.get_results_by_filter(
            entity_types=["dataset"],
            extra_or_filters=self.get_form_existence_or_filters(),
            extra_source_fields=extra_fields,
            skip_cache=True,
        )
        forms_scanned = set()
        row_index = 0
        for row in result:
            row_index += 1
            if row_index % 100 == 0:
                logger.info(f"Scanned {row_index} assets")
            extra_properties = row["extraProperties"]

            extra_properties_map = {
                x["name"]: json.loads(x["value"]) for x in extra_properties
            }
            search_row = self.DataHubDatasetSearchRow(**extra_properties_map)
            if on_asset_scanned:
                on_asset_scanned(search_row.urn)
            form_assigned_dates = self.form_assigned_date(search_row)
            form_completed_dates = self.form_completed_date(search_row)
            domain = None
            parent_domain = None
            if search_row.domains:
                domain = search_row.domains[0]
                domain_props = self.domain_registry.get_domain(domain)
                if not domain_props:
                    logger.warning(
                        f"Domain {search_row.domains[0]} not found. Will record the domain urn regardless"
                    )
                else:
                    if domain_props.parentDomain:
                        parent_domain = domain_props.parentDomain
            # for owner in owners:
            incomplete_forms = (
                [x for x in search_row.incompleteForms if x in self.allowed_forms]
                if self.allowed_forms
                else search_row.incompleteForms
            )
            for form_id in incomplete_forms:
                if form_id not in forms_scanned:
                    if on_form_scanned:
                        on_form_scanned(form_id)
                    forms_scanned.add(form_id)
                form_info = self.form_registry.get_form(form_id)
                assert form_info, f"Form {form_id} not found"
                form_prompts = [x.id for x in form_info.prompts]
                form_incomplete_prompts = [
                    p
                    for p in search_row.incompleteFormsIncompletePromptIds
                    if p in form_prompts
                ]
                form_status = (
                    FormStatus.IN_PROGRESS
                    if len(form_incomplete_prompts) < len(form_prompts)
                    else FormStatus.NOT_STARTED
                )
                form_type = (
                    FormType.DOCUMENTATION
                    if form_info.type == FormTypeClass.COMPLETION
                    else FormType.VERIFICATION
                )
                assignees = self.get_assignees(form_info, search_row.owners)
                for prompt_id in [
                    p
                    for p in search_row.incompleteFormsIncompletePromptIds
                    if p in form_prompts
                ]:
                    for owner in assignees:
                        yield FormReportingRow(
                            form_urn=form_id,
                            form_assigned_date=form_assigned_dates[form_id],
                            form_completed_date=None,
                            form_status=form_status,
                            form_type=form_type,
                            assignee_urn=owner,
                            asset_urn=search_row.urn,
                            platform_urn=search_row.platform,
                            platform_instance_urn=search_row.platformInstance,
                            domain_urn=domain,
                            parent_domain_urn=parent_domain,
                            asset_verified=(
                                False if form_type == FormType.VERIFICATION else None
                            ),
                            question_id=str(prompt_id),
                            question_status=QuestionStatus.NOT_STARTED,
                            question_completed_date=None,
                            snapshot_date=self.snapshot_date,
                        )
                for prompt_id, prompt_reponse_time in [
                    (p, p_response_time)
                    for (p, p_response_time) in zip(
                        search_row.incompleteFormsCompletedPromptIds,
                        search_row.incompleteFormsCompletedPromptResponseTimes,
                    )
                    if p in form_prompts
                ]:
                    for owner in assignees:
                        yield FormReportingRow(
                            form_urn=form_id,
                            form_assigned_date=form_assigned_dates[form_id],
                            form_completed_date=None,
                            form_status=form_status,
                            form_type=form_type,
                            assignee_urn=owner,
                            asset_urn=search_row.urn,
                            platform_urn=search_row.platform,
                            platform_instance_urn=search_row.platformInstance,
                            domain_urn=domain,
                            parent_domain_urn=parent_domain,
                            asset_verified=(
                                False if form_type == FormType.VERIFICATION else None
                            ),
                            question_id=str(prompt_id),
                            question_status=QuestionStatus.COMPLETED,
                            question_completed_date=prompt_reponse_time,
                            snapshot_date=self.snapshot_date,
                        )
            complete_forms = (
                [x for x in search_row.completedForms if x in self.allowed_forms]
                if self.allowed_forms
                else search_row.completedForms
            )
            for form_id in complete_forms:
                if form_id not in forms_scanned:
                    if on_form_scanned:
                        on_form_scanned(form_id)
                    forms_scanned.add(form_id)
                form_info = self.form_registry.get_form(form_id)
                assert form_info, f"Form {form_id} not found"
                form_type = (
                    FormType.DOCUMENTATION
                    if form_info.type == FormTypeClass.COMPLETION
                    else FormType.VERIFICATION
                )
                assignees = self.get_assignees(form_info, search_row.owners)
                form_prompts = [x.id for x in form_info.prompts]
                for prompt_id in [
                    p
                    for p in search_row.completedFormsIncompletePromptIds
                    for p in form_prompts
                ]:
                    logger.warning("Unexpected incomplete prompt in completed form")
                    for owner in assignees:
                        yield FormReportingRow(
                            form_urn=form_id,
                            form_assigned_date=form_assigned_dates[form_id],
                            form_completed_date=form_completed_dates[form_id],
                            form_status=FormStatus.COMPLETED,
                            form_type=form_type,
                            assignee_urn=owner,
                            asset_urn=search_row.urn,
                            platform_urn=search_row.platform,
                            platform_instance_urn=search_row.platformInstance,
                            domain_urn=domain,
                            parent_domain_urn=parent_domain,
                            asset_verified=(
                                True if form_type == FormType.VERIFICATION else None
                            ),
                            question_id=str(prompt_id),
                            question_status=QuestionStatus.NOT_STARTED,
                            question_completed_date=None,
                            snapshot_date=self.snapshot_date,
                        )
                for prompt_id, prompt_reponse_time in [
                    (p, p_response_time)
                    for (p, p_response_time) in zip(
                        search_row.completedFormsCompletedPromptIds,
                        search_row.completedFormsCompletedPromptResponseTimes,
                    )
                    if p in form_prompts
                ]:
                    for owner in assignees:
                        yield FormReportingRow(
                            form_urn=form_id,
                            form_assigned_date=form_assigned_dates[form_id],
                            form_completed_date=form_completed_dates[form_id],
                            form_status=FormStatus.COMPLETED,
                            form_type=form_type,
                            assignee_urn=owner,
                            asset_urn=search_row.urn,
                            platform_urn=search_row.platform,
                            platform_instance_urn=search_row.platformInstance,
                            domain_urn=domain,
                            parent_domain_urn=parent_domain,
                            asset_verified=(
                                True if form_type == FormType.VERIFICATION else None
                            ),
                            question_id=str(prompt_id),
                            question_status=QuestionStatus.COMPLETED,
                            question_completed_date=prompt_reponse_time,
                            snapshot_date=self.snapshot_date,
                        )


if __name__ == "__main__":
    from datahub.ingestion.graph.client import get_default_graph

    with get_default_graph() as graph:
        form_data = DataHubFormReportingData(graph)
        for row in form_data.get_data():
            print(row)
