"""
Utility functions for kepler-pulse.

This module provides common utility functions for date handling and
environment configuration.
"""

import os
from typing import Union, List, Optional
import pandas as pd

from .config import config


def dt_handle(date: Union[str, int, pd.Timestamp]) -> str:
    """
    Convert various date formats to standardized YYYYMMDD string format.

    Args:
        date: Date in various formats (str, int, pandas Timestamp, datetime)

    Returns:
        Date string in YYYYMMDD format

    Examples:
        >>> dt_handle('2023-01-01')
        '20230101'
        >>> dt_handle(20230101)
        '20230101'
        >>> dt_handle(pd.Timestamp('2023-01-01'))
        '20230101'
    """
    try:
        return pd.to_datetime(date).strftime(config.default_country == 'CN' and '%Y%m%d' or '%Y-%m-%d')
    except (ValueError, TypeError):
        # Return as-is if conversion fails, assuming it's already in correct format
        return str(date)


def set(country: Optional[Union[str, List[str]]] = None) -> None:
    """
    Set the default country for calendar operations.

    This function configures the environment variable that determines
    which country's trading calendar to use.

    Args:
        country: Country code(s) to set. Can be:
                - None: Use default configuration
                - str: Single country code (e.g., 'CN', 'US')
                - list: Multiple country codes for multi-market support

    Examples:
        >>> set('CN')  # Set to China market
        >>> set(['CN', 'US'])  # Set to China and US markets
        >>> set()  # Use default configuration
    """
    if country is not None:
        if isinstance(country, str):
            country = [country]

        # Update configuration
        config.default_country = country[0]  # Use first country as primary
        os.environ['kepler.pulse'] = ','.join(country)
    else:
        # Reset to default
        config.default_country = config.country
        os.environ['kepler.pulse'] = config.country