import time
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence

from dbt.artifacts.resources import SourceFileMetadata
from dbt.artifacts.resources.base import GraphResource
from dbt.artifacts.resources.v1.components import DependsOn, RefArgs
from dbt_common.contracts.config.base import BaseConfig, CompareBehavior, MergeBehavior
from dbt_common.dataclass_schema import dbtClassMixin
from dbt_semantic_interfaces.references import (
    DimensionReference,
    EntityReference,
    LinkableElementReference,
    MeasureReference,
    SemanticModelReference,
    TimeDimensionReference,
)
from dbt_semantic_interfaces.type_enums import (
    AggregationType,
    DimensionType,
    EntityType,
    TimeGranularity,
)

"""
The classes in this file are dataclasses which are used to construct the Semantic
Model node in dbt-core. Additionally, these classes need to at a minimum support
what is specified in their protocol definitions in dbt-semantic-interfaces.
Their protocol definitions can be found here:
https://github.com/dbt-labs/dbt-semantic-interfaces/blob/main/dbt_semantic_interfaces/protocols/semantic_model.py
"""


@dataclass
class SemanticLayerElementConfig(dbtClassMixin):
    meta: Dict[str, Any] = field(
        default_factory=dict,
        metadata=MergeBehavior.Update.meta(),
    )


@dataclass
class Defaults(dbtClassMixin):
    agg_time_dimension: Optional[str] = None


@dataclass
class NodeRelation(dbtClassMixin):
    alias: str
    schema_name: str  # TODO: Could this be called simply "schema" so we could reuse StateRelation?
    database: Optional[str] = None
    relation_name: Optional[str] = ""


# ====================================
# Dimension objects
# Dimension protocols: https://github.com/dbt-labs/dbt-semantic-interfaces/blob/main/dbt_semantic_interfaces/protocols/dimension.py
# ====================================


@dataclass
class DimensionValidityParams(dbtClassMixin):
    is_start: bool = False
    is_end: bool = False


@dataclass
class DimensionTypeParams(dbtClassMixin):
    time_granularity: TimeGranularity
    validity_params: Optional[DimensionValidityParams] = None


@dataclass
class Dimension(dbtClassMixin):
    name: str
    type: DimensionType
    description: Optional[str] = None
    label: Optional[str] = None
    is_partition: bool = False
    type_params: Optional[DimensionTypeParams] = None
    expr: Optional[str] = None
    metadata: Optional[SourceFileMetadata] = None
    config: Optional[SemanticLayerElementConfig] = None

    @property
    def reference(self) -> DimensionReference:
        return DimensionReference(element_name=self.name)

    @property
    def time_dimension_reference(self) -> Optional[TimeDimensionReference]:
        if self.type == DimensionType.TIME:
            return TimeDimensionReference(element_name=self.name)
        else:
            return None

    @property
    def validity_params(self) -> Optional[DimensionValidityParams]:
        if self.type_params:
            return self.type_params.validity_params
        else:
            return None


# ====================================
# Entity objects
# Entity protocols: https://github.com/dbt-labs/dbt-semantic-interfaces/blob/main/dbt_semantic_interfaces/protocols/entity.py
# ====================================


@dataclass
class Entity(dbtClassMixin):
    name: str
    type: EntityType
    description: Optional[str] = None
    label: Optional[str] = None
    role: Optional[str] = None
    expr: Optional[str] = None
    config: Optional[SemanticLayerElementConfig] = None

    @property
    def reference(self) -> EntityReference:
        return EntityReference(element_name=self.name)

    @property
    def is_linkable_entity_type(self) -> bool:
        return self.type in (EntityType.PRIMARY, EntityType.UNIQUE, EntityType.NATURAL)


# ====================================
# Measure objects
# Measure protocols: https://github.com/dbt-labs/dbt-semantic-interfaces/blob/main/dbt_semantic_interfaces/protocols/measure.py
# ====================================


@dataclass
class MeasureAggregationParameters(dbtClassMixin):
    percentile: Optional[float] = None
    use_discrete_percentile: bool = False
    use_approximate_percentile: bool = False


@dataclass
class NonAdditiveDimension(dbtClassMixin):
    name: str
    window_choice: AggregationType
    window_groupings: List[str]


@dataclass
class Measure(dbtClassMixin):
    name: str
    agg: AggregationType
    description: Optional[str] = None
    label: Optional[str] = None
    create_metric: bool = False
    expr: Optional[str] = None
    agg_params: Optional[MeasureAggregationParameters] = None
    non_additive_dimension: Optional[NonAdditiveDimension] = None
    agg_time_dimension: Optional[str] = None
    config: Optional[SemanticLayerElementConfig] = None

    @property
    def reference(self) -> MeasureReference:
        return MeasureReference(element_name=self.name)


# ====================================
# SemanticModel final parts
# ====================================


@dataclass
class SemanticModelConfig(BaseConfig):
    enabled: bool = True
    group: Optional[str] = field(
        default=None,
        metadata=CompareBehavior.Exclude.meta(),
    )
    meta: Dict[str, Any] = field(
        default_factory=dict,
        metadata=MergeBehavior.Update.meta(),
    )


@dataclass
class SemanticModel(GraphResource):
    model: str
    node_relation: Optional[NodeRelation]
    description: Optional[str] = None
    label: Optional[str] = None
    defaults: Optional[Defaults] = None
    entities: Sequence[Entity] = field(default_factory=list)
    measures: Sequence[Measure] = field(default_factory=list)
    dimensions: Sequence[Dimension] = field(default_factory=list)
    metadata: Optional[SourceFileMetadata] = None
    depends_on: DependsOn = field(default_factory=DependsOn)
    refs: List[RefArgs] = field(default_factory=list)
    created_at: float = field(default_factory=lambda: time.time())
    config: SemanticModelConfig = field(default_factory=SemanticModelConfig)
    unrendered_config: Dict[str, Any] = field(default_factory=dict)
    primary_entity: Optional[str] = None
    group: Optional[str] = None

    @property
    def entity_references(self) -> List[LinkableElementReference]:
        return [entity.reference for entity in self.entities]

    @property
    def dimension_references(self) -> List[LinkableElementReference]:
        return [dimension.reference for dimension in self.dimensions]

    @property
    def measure_references(self) -> List[MeasureReference]:
        return [measure.reference for measure in self.measures]

    @property
    def has_validity_dimensions(self) -> bool:
        return any([dim.validity_params is not None for dim in self.dimensions])

    @property
    def validity_start_dimension(self) -> Optional[Dimension]:
        validity_start_dims = [
            dim for dim in self.dimensions if dim.validity_params and dim.validity_params.is_start
        ]
        if not validity_start_dims:
            return None
        return validity_start_dims[0]

    @property
    def validity_end_dimension(self) -> Optional[Dimension]:
        validity_end_dims = [
            dim for dim in self.dimensions if dim.validity_params and dim.validity_params.is_end
        ]
        if not validity_end_dims:
            return None
        return validity_end_dims[0]

    @property
    def partitions(self) -> List[Dimension]:  # noqa: D
        return [dim for dim in self.dimensions or [] if dim.is_partition]

    @property
    def partition(self) -> Optional[Dimension]:
        partitions = self.partitions
        if not partitions:
            return None
        return partitions[0]

    @property
    def reference(self) -> SemanticModelReference:
        return SemanticModelReference(semantic_model_name=self.name)

    def checked_agg_time_dimension_for_measure(
        self, measure_reference: MeasureReference
    ) -> TimeDimensionReference:
        measure: Optional[Measure] = None
        for measure in self.measures:
            if measure.reference == measure_reference:
                measure = measure

        assert (
            measure is not None
        ), f"No measure with name ({measure_reference.element_name}) in semantic_model with name ({self.name})"

        default_agg_time_dimension = (
            self.defaults.agg_time_dimension if self.defaults is not None else None
        )

        agg_time_dimension_name = measure.agg_time_dimension or default_agg_time_dimension
        assert agg_time_dimension_name is not None, (
            f"Aggregation time dimension for measure {measure.name} on semantic model {self.name} is not set! "
            "To fix this either specify a default `agg_time_dimension` for the semantic model or define an "
            "`agg_time_dimension` on the measure directly."
        )
        return TimeDimensionReference(element_name=agg_time_dimension_name)

    @property
    def primary_entity_reference(self) -> Optional[EntityReference]:
        return (
            EntityReference(element_name=self.primary_entity)
            if self.primary_entity is not None
            else None
        )
