from datetime import datetime
from typing import List, Optional, Union

from dagster import (
    EnvVar,
    TimeWindowPartitionMapping,
    IdentityPartitionMapping,
    AllPartitionMapping,
    LastPartitionMapping,
    StaticPartitionMapping,
)

DBT_MIN_PARTITION_DATE = datetime.strptime('2024-01-01', '%Y-%m-%d')

DBT_PARTITION_MAPPING_MAP = {
    'time_window': TimeWindowPartitionMapping,
    'identity': IdentityPartitionMapping,
    'all': AllPartitionMapping,
    'last': LastPartitionMapping,
    'static': StaticPartitionMapping,
}


def get_dbt_tag_selection_string(tags: Union[str, List[str]], tag_match_condition: Optional[str] = None) -> str:
    """
    Get the dbt tag selection string for the given tags and tag match condition.
    Note that tag match condition ALL is used to select models that have ALL of the tags,
    while tag match condition ANY is used to select models that have ANY of the tags.

    Args:
        tags: List[str]
        tag_match_condition: str
    Returns:
        str
    """
    if isinstance(tags, str):
        return 'tag:' + tags
    elif isinstance(tags, List):
        if tag_match_condition is None:
            raise ValueError("Tag match condition is required when tags is a list")

        if tag_match_condition == 'ALL':
            return ','.join(['tag:' + tag for tag in tags])
        elif tag_match_condition == 'ANY':
            return ' '.join(['tag:' + tag for tag in tags])
        else:
            raise ValueError(f"Invalid tag match condition {tag_match_condition}, must be ALL or ANY")


def get_dbt_schema_name(schema: str) -> str:
    """
    Get the dbt schema name for the given schema.
    """
    dbt_target = EnvVar('DBT_TARGET').get_value()
    if dbt_target == 'dev':
        default_schema = EnvVar('DBT_USER').get_value()
    else:
        default_schema = dbt_target

    return f"{default_schema}_{schema}"