from typing import (
    Any,
    Dict,
    Iterable,
    List,
    Literal,
    Optional,
    Set,
    TypedDict,
    TypeVar,
)

from pydantic import BaseModel, ConfigDict, Field, model_validator

from datahub.ingestion.source.common.subtypes import DataFlowSubTypes
from datahub.ingestion.source.sap_datasphere.config import ConnectionPlatformConfig
from datahub.ingestion.source.sap_datasphere.constants import (
    MALFORMED_COL_SENTINEL,
    REMOTE_IDENTIFIER_QUOTE,
    REMOTE_NULL_SEGMENT,
    UNNAMED_COL_SENTINEL,
)
from datahub.metadata.schema_classes import (
    FineGrainedLineageClass,
    SchemaFieldClass,
)
from datahub.utilities.str_enum import StrEnum

_T = TypeVar("_T")

# A parsed JSON object at an API boundary (response body / CSN node). Values are
# arbitrary JSON, so callers must narrow before use.
JsonDict = Dict[str, Any]


def dedup_preserving_order(items: Iterable[_T]) -> List[_T]:
    seen: Set[_T] = set()
    out: List[_T] = []
    for item in items:
        if item not in seen:
            seen.add(item)
            out.append(item)
    return out


def _require_exactly_one(
    model: str, primary_field: str, primary: object, alt_field: str, alt: object
) -> None:
    if (primary is None) == (alt is None):
        raise ValueError(
            f"{model} must carry exactly one of {primary_field}/{alt_field} "
            f"({primary_field}_is_none={primary is None}, {alt_field}={alt!r})"
        )


class TagDefinition(BaseModel):
    model_config = ConfigDict(frozen=True)

    name: str
    description: str


class UnknownColumnType(BaseModel):
    """A column whose source type literal (CDS or EDM) is not in the parser's
    type map. Emitted by both parsers so the source can report it uniformly."""

    model_config = ConfigDict(frozen=True)

    type: str
    column: str


class ConnectionRecord(TypedDict, total=False):
    # The SAP connections API returns more fields (id, description, owner, ...)
    # but the connector only consumes name + typeId; total=False stays permissive
    # at the JSON boundary.
    name: str
    typeId: str


class ResolveSkipReason(StrEnum):
    UNKNOWN_TYPEID = "unknown_typeid"
    UNKNOWN_CONNECTION = "unknown_connection"
    DISABLED = "disabled"


class ResolvedPlatform(BaseModel):
    # A disabled mapping is turned into a skip during resolution, so holding a
    # ResolvedPlatform always implies enabled.
    model_config = ConfigDict(frozen=True)

    platform: str
    platform_instance: Optional[str]
    env: str
    # Per-platform casing for external (federated / flow-target) URNs, matching how
    # the sibling native connector cases its URNs. Independent of the connector's
    # top-level flag (which governs managed assets).
    convert_urns_to_lowercase: bool = True
    # Optional leading name segment for external URNs that the Datasphere API
    # can't supply — chiefly the BigQuery GCP project, which is absent from the
    # connections/flow payloads. When set, it is prepended ahead of the
    # container-derived schema so the URN becomes `database.schema.table`.
    database: Optional[str] = None


class ResolveResult(BaseModel):
    # Exactly one side is populated: a ResolvedPlatform on success, or a
    # skip_reason when the asset's connection could not be resolved.
    model_config = ConfigDict(frozen=True)

    platform: Optional[ResolvedPlatform]
    skip_reason: Optional[ResolveSkipReason]

    @model_validator(mode="after")
    def _check_invariant(self) -> "ResolveResult":
        _require_exactly_one(
            "ResolveResult", "platform", self.platform, "skip_reason", self.skip_reason
        )
        return self


class TypeIdDefault(BaseModel):
    # Result of mapping a connection's typeId to its platform default: the raw
    # (still enabled/disabled-unchecked) config on success, else a skip reason.
    model_config = ConfigDict(frozen=True)

    config: Optional[ConnectionPlatformConfig]
    skip_reason: Optional[ResolveSkipReason]

    @model_validator(mode="after")
    def _check_invariant(self) -> "TypeIdDefault":
        _require_exactly_one(
            "TypeIdDefault", "config", self.config, "skip_reason", self.skip_reason
        )
        return self


class AssetCsn(BaseModel):
    # The design-time CSN of a View / Analytic Model asset, resolved once so the
    # emit path doesn't re-fetch or re-navigate. ``connection_name`` is the routing
    # connection (the managed HANA key unless the CSN declared a federated
    # @remote.source).
    model_config = ConfigDict(frozen=True)

    connection_name: str
    csn_obj: Optional[JsonDict] = None
    csn_def: Optional[JsonDict] = None


class EdmxFetchReason(StrEnum):
    OK = "ok"
    FORBIDDEN = "forbidden"  # 403 — principal lacks OData read (already warned)
    NOT_FOUND = "not_found"  # 404 — asset not exposed via OData (benign)
    ERROR = "error"  # network / non-2xx / other


class EdmxFetchResult(BaseModel):
    model_config = ConfigDict(frozen=True)

    xml: Optional[str]
    reason: EdmxFetchReason

    @model_validator(mode="after")
    def _check_invariant(self) -> "EdmxFetchResult":
        # Make the illegal states (OK-with-no-xml, non-OK carrying xml)
        # unrepresentable so callers can branch on reason and trust xml.
        if (self.reason is EdmxFetchReason.OK) != (self.xml is not None):
            raise ValueError(
                f"EdmxFetchResult invariant violated: xml must be set iff "
                f"reason is OK (got reason={self.reason}, xml_is_none="
                f"{self.xml is None})"
            )
        return self


class EdmxParseResult(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    fields: List[SchemaFieldClass]
    field_custom_props: Dict[str, Dict[str, str]]  # fieldPath → {key: value}
    entity_label: Optional[str]
    entity_custom_props: Dict[str, str]
    error: Optional[str] = None  # set when parse failed; None on success
    unknown_edm_types: List[UnknownColumnType] = Field(default_factory=list)

    @model_validator(mode="after")
    def _check_failure_carries_no_payload(self) -> "EdmxParseResult":
        # A failure result must not smuggle partial schema through: callers branch
        # on ``error`` and skip the payload, so any fields/props set alongside an
        # error would be silently dropped and mask a parser bug.
        if self.error is not None and (
            self.fields
            or self.field_custom_props
            or self.entity_label is not None
            or self.entity_custom_props
            or self.unknown_edm_types
        ):
            raise ValueError(
                "EdmxParseResult with an error must carry no fields/props "
                f"(error={self.error!r})"
            )
        return self


class CsnSchemaResult(BaseModel):
    # CSN-elements schema parse output: the fields plus any columns whose CDS
    # type literal was unmapped (surfaced in the report, not silently dropped).
    model_config = ConfigDict(arbitrary_types_allowed=True)

    fields: List[SchemaFieldClass]
    unknown_types: List[UnknownColumnType] = Field(default_factory=list)
    # Names of association/composition elements skipped from the scalar schema;
    # their targets are turned into lineage edges by the lineage extractor.
    navigation_elements: List[str] = Field(default_factory=list)
    # Columns whose CSN element carried no (or an empty) type key — a structural
    # concern distinct from an unknown-but-present type. They emit as NullType
    # (native "UNKNOWN"), but are surfaced so an operator can tell a genuine typed
    # column from one whose type the API failed to emit.
    columns_missing_type: List[str] = Field(default_factory=list)


class BusinessLayer(BaseModel):
    model_config = ConfigDict(frozen=True)

    fact_source_keys: List[str]
    dimension_source_keys: List[str]
    measure_names: List[str]
    attribute_names: List[str]
    variable_names: List[str]

    @property
    def upstream_keys(self) -> List[str]:
        # Derived (facts first) rather than stored so it can never drift.
        return dedup_preserving_order(
            self.fact_source_keys + self.dimension_source_keys
        )


class UpstreamRef(BaseModel):
    """A table-level upstream reference the source layer resolves to a URN.

    ``qualified`` distinguishes a bare same-space technical name (space-prefixed
    at URN build time) from an already space-qualified name — e.g. a cross-space
    or built-in association target like ``SAP.TIME.M_TIME_DIMENSION_TDAY`` — which
    is used as-is.
    """

    model_config = ConfigDict(frozen=True)

    name: str
    qualified: bool = False


class UpstreamColRef(BaseModel):
    model_config = ConfigDict(frozen=True)

    qname: str
    col: str
    # True when ``qname`` is already space-qualified (association target); a bare
    # FROM-derived ref is space-prefixed by the source layer.
    qualified: bool = False


class SourceColumnRef(BaseModel):
    # Where an analytic-model element gets its scalar type: the source object it
    # projects from and the column within it. Analytic-model elements carry no
    # inline CSN type, so the type is copied from this source column's schema.
    model_config = ConfigDict(frozen=True)

    source_object: str
    column: str


# Transform-operation labels DataHub's UI recognises for its transformation badge.
TransformOp = Literal["IDENTITY", "RENAME", "AGGREGATE", "TRANSFORMATION", "EXPRESSION"]


class ColumnLineagePair(BaseModel):
    # downstream_col may be MALFORMED_COL_SENTINEL / UNNAMED_COL_SENTINEL for
    # structurally-broken or unnamed expressions. upstream_refs is empty for
    # pure-literal columns. Resolution of CSN qualified-names to DataHub URNs is
    # the source layer's job. unresolved_refs carries diagnostics for refs the
    # walker could not attribute, surfaced via report.column_lineage_unresolved.
    model_config = ConfigDict(frozen=True)

    downstream_col: str
    upstream_refs: List[UpstreamColRef] = Field(default_factory=list)
    transform_op: Optional[TransformOp] = None
    unresolved_refs: List[str] = Field(default_factory=list)

    @model_validator(mode="after")
    def _check_sentinel(self) -> "ColumnLineagePair":
        # A sentinel-valued pair with upstream refs would make the source build a
        # schemaField URN containing the sentinel string.
        if self.downstream_col in (MALFORMED_COL_SENTINEL, UNNAMED_COL_SENTINEL):
            if self.upstream_refs:
                raise ValueError(
                    f"Sentinel ColumnLineagePair must have empty upstream_refs; "
                    f"got downstream_col={self.downstream_col!r}, "
                    f"upstream_refs={self.upstream_refs!r}"
                )
        return self


class CsnSelectEnvelope(BaseModel):
    # Tri-state result of validating a CSN query envelope: ``select`` set to
    # proceed, ``malformed`` set (with a diagnostic pair) for broken CSN, or both
    # None for a legitimate non-SELECT entity.
    model_config = ConfigDict(frozen=True)

    select: Optional[JsonDict] = None
    malformed: Optional["ColumnLineagePair"] = None

    @model_validator(mode="after")
    def _check_tristate(self) -> "CsnSelectEnvelope":
        if self.select is not None and self.malformed is not None:
            raise ValueError("CsnSelectEnvelope cannot set both select and malformed")
        return self


class ColumnLineageContext(BaseModel):
    # The pairs and the downstream URN are coupled (the URN is needed to build
    # the downstream field URN of every pair); bundling them keeps callers from
    # passing inconsistent combinations.
    model_config = ConfigDict(frozen=True)

    pairs: List[ColumnLineagePair]
    downstream_dataset_urn: str


class FlowEndpoint(BaseModel):
    # One input/output object of a flow. is_local (connection is None) means it
    # lives in the tenant's own managed HANA Cloud and is emitted space-prefixed
    # on sap_datasphere; otherwise connection/connection_type drive platform
    # mapping and container (the source/target system's schema/dataset path, e.g.
    # "/CDS_EXTRACTION") schema-qualifies the external URN.
    model_config = ConfigDict(frozen=True)

    object_name: str
    is_local: bool
    connection: Optional[str] = None
    connection_type: Optional[str] = None
    container: Optional[str] = None

    @model_validator(mode="after")
    def _check_local_connection(self) -> "FlowEndpoint":
        # Local iff no connection: the two are derived from the same source id, so
        # a mismatch means the parser mis-classified the endpoint.
        if self.is_local != (self.connection is None):
            raise ValueError(
                "FlowEndpoint.is_local must match connection being None "
                f"(is_local={self.is_local}, connection={self.connection!r})"
            )
        # A local endpoint has no external routing metadata; carrying a
        # connection_type/container while local would mean the parser leaked the
        # external system identity onto a managed object.
        if self.is_local and (
            self.connection_type is not None or self.container is not None
        ):
            raise ValueError(
                "FlowEndpoint marked local must not carry connection_type/container "
                f"(connection_type={self.connection_type!r}, "
                f"container={self.container!r})"
            )
        return self


class FlowColumnMapping(BaseModel):
    # object names must match a FlowEndpoint.object_name so the source layer can
    # resolve them to the same dataset URNs used for the table-level edges.
    model_config = ConfigDict(frozen=True)

    downstream_object: str
    downstream_col: str
    upstream_object: str
    upstream_col: str


class FlowTableMapping(BaseModel):
    # A single (upstream_object -> downstream_object) task pairing. Replication
    # flows carry independent tasks, so the source layer must attribute each
    # target's upstreams to its own paired source rather than to every flow input.
    model_config = ConfigDict(frozen=True)

    upstream_object: str
    downstream_object: str


class SystemIdentity(BaseModel):
    # The external connection id + type of a replication flow's source/target
    # system, used to route its objects to a DataHub platform. ``container`` is
    # the system-wide schema/dataset path shared by every object on that side.
    model_config = ConfigDict(frozen=True)

    connection: Optional[str] = None
    connection_type: Optional[str] = None
    container: Optional[str] = None


class AttrMapping(BaseModel):
    # A single (downstream_col, upstream_col) column mapping parsed from a flow's
    # attributeMappings before it is attributed to concrete objects.
    model_config = ConfigDict(frozen=True)

    downstream_col: str
    upstream_col: str


class ProducerColumns(BaseModel):
    # A producer node's output object plus its column mappings, buffered until
    # the flow's sole input is known (column lineage needs both endpoints).
    model_config = ConfigDict(frozen=True)

    object_name: str
    mappings: List[AttrMapping] = Field(default_factory=list)


class ParsedFlow(BaseModel):
    # A flow reduced to what DataHub needs; platform/URN resolution is the source
    # layer's job. table_mappings preserves per-task source->target pairing
    # (populated for replication flows) so an independent task's target isn't
    # given a cross-product of every flow input as its upstreams. Empty
    # table_mappings signals a connected-graph flow (data/transformation), where
    # every input legitimately feeds every output.
    model_config = ConfigDict(frozen=True)

    technical_name: str
    subtype: DataFlowSubTypes
    inputs: List[FlowEndpoint] = Field(default_factory=list)
    outputs: List[FlowEndpoint] = Field(default_factory=list)
    column_mappings: List[FlowColumnMapping] = Field(default_factory=list)
    table_mappings: List[FlowTableMapping] = Field(default_factory=list)
    # True when column mappings existed but were dropped because the flow has
    # multiple inputs (column attribution would be ambiguous); surfaced as a
    # report counter so operators see the coarsening.
    cll_suppressed_multi_input: bool = False
    # Count of malformed process nodes / replication tasks skipped during parse
    # (wrong shape, missing component/config, or unresolvable endpoint). A
    # partial parse can silently drop lineage edges, so the source surfaces this
    # as a report counter rather than leaving it debug-only.
    dropped_node_count: int = 0


class FlowTask(BaseModel):
    # One target-anchored task of a flow: the DataJob is named after the target
    # it produces (unique within the flow), its inlets are the sources feeding
    # that target, and its fine-grained edges are the column mappings landing on
    # it. Splitting a flow into per-target tasks keeps a replication flow's N
    # source->target pairs as N distinct jobs and stops a single-target flow's
    # job from duplicating the flow's own name.
    model_config = ConfigDict(arbitrary_types_allowed=True)

    target_object: str
    target_urn: str
    upstream_urns: List[str] = Field(default_factory=list)
    fine_grained: List[FineGrainedLineageClass] = Field(default_factory=list)


class RemoteTableSource(BaseModel):
    """The external origin of a federated Remote Table, parsed from the CSN
    ``@DataWarehouse.remote.*`` annotations."""

    model_config = ConfigDict(frozen=True)

    connection: str
    path_parts: List[str]

    @property
    def qualified_name(self) -> str:
        # Drop the "<NULL>" database sentinel (and empties), then strip SQL
        # identifier quotes so the external URN name is the unquoted
        # "schema.table" the sibling connector emits. The quoted single-segment
        # form ('"SCHEMA"."/BIC/TABLE"') reduces to SCHEMA./BIC/TABLE, matching
        # the delimiter form (<NULL>\x7fSCHEMA\x7f/BIC/TABLE).
        meaningful = [p for p in self.path_parts if p and p != REMOTE_NULL_SEGMENT]
        return ".".join(
            part.replace(REMOTE_IDENTIFIER_QUOTE, "") for part in meaningful
        )


class UnionMergeSlot(BaseModel):
    # Accumulator for merging one UNION output column's upstream refs across the
    # branch SELECTs (mutated during the merge, so not frozen).
    refs: List[UpstreamColRef] = Field(default_factory=list)
    unresolved: List[str] = Field(default_factory=list)
    op: Optional[TransformOp] = None
