import json
import logging
import time
import uuid
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple

from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from acryl_datahub_cloud.metadata.schema_classes import (
    ActionRequestInfoClass,
    ActionRequestParamsClass,
    ActionRequestStatusClass,
    AuditStampClass,
    DocumentChangeClass,
    DocumentProposalTypeClass,
    DocumentsProposalClass,
)
from datahub.configuration import ConfigModel
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.ingestion.api.common import PipelineContext
from datahub.ingestion.api.decorators import (
    SupportStatus,
    config_class,
    platform_name,
    support_status,
)
from datahub.ingestion.api.source import Source, SourceReport
from datahub.ingestion.api.workunit import MetadataWorkUnit
from datahub.ingestion.graph.client import DataHubGraph

logger = logging.getLogger(__name__)

ADMIN_ROLE_URN = "urn:li:dataHubRole:Admin"
SYSTEM_URN = "urn:li:corpuser:__datahub_system"
UNPUBLISHED_STAGE_URN = "urn:li:lifecycleStageType:UNPUBLISHED"
PUBLISHED_STAGE_URN = "urn:li:lifecycleStageType:PUBLISHED"
MAX_DOCS_PER_PROPOSAL = 5

# ES field names for the structured properties on document entities.
# Unversioned structured properties use the qualified name with dots replaced
# by underscores, under the `structuredProperties.` prefix.
EXPERTS_SP_FIELD = "structuredProperties.acryl_genai_sqlAnchor_experts"
DOMAINS_SP_FIELD = "structuredProperties.acryl_genai_sqlAnchor_domains"

LIST_ACTION_REQUESTS_QUERY = """
query listActionRequests($input: ListActionRequestsInput!) {
  listActionRequests(input: $input) {
    nextScrollId
    count
    total
    actionRequests {
      urn
      type
    }
  }
}
"""

SCROLL_UNPUBLISHED_DOCS_QUERY = """
query scrollUnpublishedDocs($scrollId: String, $count: Int!, $fetchExtraFields: [String!]) {
  scrollAcrossEntities(input: {
    query: "*"
    scrollId: $scrollId
    count: $count
    types: [DOCUMENT]
    orFilters: [{
      and: [
        {
          field: "lifecycleStage"
          values: ["urn:li:lifecycleStageType:UNPUBLISHED"]
        }
        {
          field: "creator"
          values: ["urn:li:corpuser:__datahub_system"]
        }
      ]
    }]
    searchFlags: {
      skipHighlighting: true
      skipAggregates: true
      fetchExtraFields: $fetchExtraFields
    }
  }) {
    nextScrollId
    searchResults {
      entity {
        urn
      }
      extraProperties {
        name
        value
      }
    }
  }
}
"""


@dataclass
class DocInfo:
    """Holds a document URN and its associated expert/domain metadata."""

    urn: str
    expert_urns: List[str] = field(default_factory=list)
    domain_urns: List[str] = field(default_factory=list)


class DataHubContextDocProposalsSourceConfig(ConfigModel):
    batch_size: int = 500


@dataclass
class DataHubContextDocProposalsSourceReport(SourceReport):
    unpublished_docs_found: int = 0
    docs_already_covered: int = 0
    proposals_created: int = 0


@platform_name(id="datahub", platform_name="DataHub")
@config_class(DataHubContextDocProposalsSourceConfig)
@support_status(SupportStatus.INCUBATING)
class DataHubContextDocProposalsSource(Source):
    """Creates action request proposals for AI-generated documents awaiting review."""

    def __init__(
        self,
        config: DataHubContextDocProposalsSourceConfig,
        ctx: PipelineContext,
    ):
        super().__init__(ctx)
        self.config = config
        self.report = DataHubContextDocProposalsSourceReport()
        self.graph: DataHubGraph = ctx.require_graph(
            "DataHub Context Doc Proposals source"
        )

    def get_workunits(self) -> Iterable[MetadataWorkUnit]:
        logger.info("Starting context doc proposals run")

        covered_doc_urns = self._get_covered_doc_urns()
        logger.info(
            f"Found {len(covered_doc_urns)} doc URNs already covered by existing proposals"
        )

        all_docs = self._scroll_unpublished_docs()
        logger.info(f"Found {len(all_docs)} unpublished system-generated docs")

        pending_docs = [doc for doc in all_docs if doc.urn not in covered_doc_urns]
        self.report.docs_already_covered = self.report.unpublished_docs_found - len(
            pending_docs
        )
        logger.info(
            f"{len(pending_docs)} docs need proposals (after excluding already-covered)"
        )

        grouped_batches = _group_docs_into_batches(pending_docs)
        for batch in grouped_batches:
            doc_urns = [doc.urn for doc in batch]
            expert_urns = _collect_expert_urns(batch)
            logger.info(
                f"Creating proposal for batch of {len(batch)} docs "
                f"with {len(expert_urns)} expert assignees"
            )
            yield from self._build_proposal_workunits(doc_urns, expert_urns)

    def _scroll_unpublished_docs(self) -> List[DocInfo]:
        docs: List[DocInfo] = []
        scroll_id: Optional[str] = None

        while True:
            try:
                response = self._execute_graphql_with_retry(
                    SCROLL_UNPUBLISHED_DOCS_QUERY,
                    variables={
                        "scrollId": scroll_id,
                        "count": self.config.batch_size,
                        "fetchExtraFields": [EXPERTS_SP_FIELD, DOMAINS_SP_FIELD],
                    },
                )
            except Exception as e:
                self.report.report_failure(
                    title="Failed to scroll unpublished documents",
                    message="Error while scrolling for UNPUBLISHED lifecycle stage documents",
                    context=str(e),
                    exc=e,
                )
                break

            result = response.get("scrollAcrossEntities", {})
            page_results = result.get("searchResults", [])
            for search_result in page_results:
                urn = search_result.get("entity", {}).get("urn")
                if not urn:
                    continue
                extra_props = search_result.get("extraProperties") or []
                expert_urns, domain_urns = _parse_structured_properties(extra_props)
                docs.append(
                    DocInfo(
                        urn=urn,
                        expert_urns=expert_urns,
                        domain_urns=domain_urns,
                    )
                )

            scroll_id = result.get("nextScrollId")
            if not scroll_id:
                break

        self.report.unpublished_docs_found = len(docs)
        return docs

    def _get_covered_doc_urns(self) -> Set[str]:
        """Return the set of document URNs already covered by a DOCUMENTS_PROPOSAL action request."""
        covered: Set[str] = set()
        scroll_id: Optional[str] = None

        while True:
            try:
                response = self._execute_graphql_with_retry(
                    LIST_ACTION_REQUESTS_QUERY,
                    variables={
                        "input": {
                            "type": "DOCUMENTS_PROPOSAL",
                            "orFilters": [
                                {
                                    "and": [
                                        {
                                            "field": "documentProposalType",
                                            "values": ["STATE_CHANGE"],
                                        }
                                    ]
                                }
                            ],
                            "allActionRequests": True,
                            "useScroll": True,
                            "scrollId": scroll_id,
                            "count": self.config.batch_size,
                        }
                    },
                )
            except Exception as e:
                self.report.report_failure(
                    title="Failed to list pending action requests",
                    message="Error while listing action requests to find already-covered documents",
                    context=str(e),
                    exc=e,
                )
                break

            list_result = response.get("listActionRequests", {})
            for action_request in list_result.get("actionRequests", []):
                urn = action_request.get("urn")
                if not urn:
                    continue
                covered.update(self._extract_target_doc_urns(urn))

            scroll_id = list_result.get("nextScrollId")
            if not scroll_id:
                break

        return covered

    def _extract_target_doc_urns(self, action_request_urn: str) -> List[str]:
        """Fetch the actionRequestInfo aspect and return all targetDocument URNs in the proposal."""
        try:
            info = self.graph.get_aspect_v2(
                entity_urn=action_request_urn,
                aspect="actionRequestInfo",
                aspect_type=ActionRequestInfoClass,
            )
        except Exception as e:
            self.report.report_warning(
                message="Failed to fetch actionRequestInfo for action request",
                context=f"{action_request_urn}: {e}",
            )
            return []

        if info is None or info.params is None:
            return []

        # Only suppress docs covered by proposals we created — not human-initiated ones.
        if info.createdBy != SYSTEM_URN:
            return []

        docs_proposal = info.params.documentsProposal
        if docs_proposal is None:
            return []

        return [
            change.targetDocument
            for change in docs_proposal.documents
            if change.targetDocument
        ]

    def _build_proposal_workunits(
        self, doc_urns: List[str], expert_urns: List[str]
    ) -> Iterable[MetadataWorkUnit]:
        action_request_urn = f"urn:li:actionRequest:{uuid.uuid4()}"
        now_ms = _now_ms()

        documents = [
            DocumentChangeClass(
                targetDocument=doc_urn,
                targetLifecycleStage=PUBLISHED_STAGE_URN,
            )
            for doc_urn in doc_urns
        ]

        params = ActionRequestParamsClass(
            documentsProposal=DocumentsProposalClass(
                proposalType=DocumentProposalTypeClass.STATE_CHANGE,
                documents=documents,
                rationale="AI-generated documents awaiting review and approval for publication.",
            )
        )

        assigned_users, assigned_groups = _split_user_and_group_urns(expert_urns)
        # Only fall back to the Admin role when there are no expert assignees.
        assigned_roles = [ADMIN_ROLE_URN] if not expert_urns else []

        info = ActionRequestInfoClass(
            type="DOCUMENTS_PROPOSAL",
            assignedUsers=assigned_users,
            assignedGroups=assigned_groups,
            assignedRoles=assigned_roles,
            created=now_ms,
            createdBy=SYSTEM_URN,
            params=params,
            origin="INFERRED",
        )

        status = ActionRequestStatusClass(
            status="PENDING",
            lastModified=AuditStampClass(time=now_ms, actor=SYSTEM_URN),
        )

        self.report.proposals_created += 1

        yield MetadataChangeProposalWrapper(
            entityUrn=action_request_urn,
            aspect=info,
        ).as_workunit()

        yield MetadataChangeProposalWrapper(
            entityUrn=action_request_urn,
            aspect=status,
        ).as_workunit()

    @retry(
        retry=retry_if_exception_type((Exception, ConnectionError)),
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10),
        reraise=True,
    )
    def _execute_graphql_with_retry(
        self, query: str, variables: Dict[str, Any]
    ) -> Dict[str, Any]:
        response = self.graph.execute_graphql(query, variables=variables)
        error = response.get("error")
        if error:
            raise Exception(f"GraphQL error: {error}")
        return response

    def get_report(self) -> SourceReport:
        return self.report


def _parse_structured_properties(
    extra_props: List[Dict[str, str]],
) -> Tuple[List[str], List[str]]:
    """Extract expert and domain URNs from the extra properties returned by search.

    Each structured property value is a JSON-encoded object. Expert entries
    contain a ``principal_urn`` field, domain entries contain ``domain_urn``.
    """
    expert_urns: List[str] = []
    domain_urns: List[str] = []

    for prop in extra_props:
        name = prop.get("name", "")
        raw_value = prop.get("value", "")
        if not raw_value:
            continue

        if name == EXPERTS_SP_FIELD:
            expert_urns = _parse_json_list_field(raw_value, "principal_urn")
        elif name == DOMAINS_SP_FIELD:
            domain_urns = _parse_json_list_field(raw_value, "domain_urn")

    return expert_urns, domain_urns


def _parse_json_list_field(raw_value: str, urn_key: str) -> List[str]:
    """Parse a JSON-encoded list of objects and extract the given key from each.

    The value from ES may come back as a JSON array of strings (each string
    being a JSON-encoded object), or as a single JSON-encoded string.
    """
    urns: List[str] = []
    try:
        parsed = json.loads(raw_value)
    except (json.JSONDecodeError, TypeError):
        logger.debug(f"Could not parse structured property value: {raw_value[:100]}")
        return urns

    items = parsed if isinstance(parsed, list) else [parsed]
    for item in items:
        if isinstance(item, str):
            # Each item is itself a JSON-encoded object
            try:
                item = json.loads(item)
            except (json.JSONDecodeError, TypeError):
                continue
        if isinstance(item, dict):
            urn = item.get(urn_key)
            if urn:
                urns.append(urn)
    return urns


def _grouping_key(doc: DocInfo) -> Tuple[str, ...]:
    """Compute a grouping key from a doc's experts or domains.

    Primary grouping is by experts — docs with the same expert set land in the
    same proposal so reviewers are consistent.  When a doc has no experts, we
    fall back to grouping by domain so that related docs still cluster together.
    """
    if doc.expert_urns:
        return ("experts",) + tuple(sorted(doc.expert_urns))
    if doc.domain_urns:
        return ("domains",) + tuple(sorted(doc.domain_urns))
    return ("none",)


def _group_docs_into_batches(docs: List[DocInfo]) -> List[List[DocInfo]]:
    """Group docs by experts (primary) or domains (fallback), then split into batches."""
    groups: Dict[Tuple[str, ...], List[DocInfo]] = defaultdict(list)
    for doc in docs:
        key = _grouping_key(doc)
        groups[key].append(doc)

    batches: List[List[DocInfo]] = []
    for group_docs in groups.values():
        for i in range(0, len(group_docs), MAX_DOCS_PER_PROPOSAL):
            batches.append(group_docs[i : i + MAX_DOCS_PER_PROPOSAL])
    return batches


def _collect_expert_urns(batch: List[DocInfo]) -> List[str]:
    """Return the deduplicated union of expert URNs across all docs in a batch."""
    seen: Set[str] = set()
    result: List[str] = []
    for doc in batch:
        for urn in doc.expert_urns:
            if urn not in seen:
                seen.add(urn)
                result.append(urn)
    return result


def _split_user_and_group_urns(
    urns: List[str],
) -> Tuple[List[str], List[str]]:
    """Partition expert URNs into user URNs and group URNs."""
    users: List[str] = []
    groups: List[str] = []
    for urn in urns:
        if urn.startswith("urn:li:corpGroup:"):
            groups.append(urn)
        else:
            users.append(urn)
    return users, groups


def _now_ms() -> int:
    return int(time.time() * 1000)
