# mypy: disable-error-code="union-attr"
"""Module for writing SDMX-ML 2.1 Structure Specific Time Series messages."""

from pathlib import Path
from typing import Dict, Optional, Sequence, Union

from pysdmx.errors import Invalid
from pysdmx.io.format import Format
from pysdmx.io.pd import PandasDataset
from pysdmx.io.xml.sdmx21.writer.structure_specific import (
    write as _base_write,
)
from pysdmx.model.message import Header


def write(
    datasets: Sequence[PandasDataset],
    output_path: Optional[Union[str, Path]] = None,
    prettyprint: bool = True,
    header: Optional[Header] = None,
    dimension_at_observation: Optional[Dict[str, str]] = None,
) -> Optional[str]:
    """Write data to SDMX-ML 2.1 Structure Specific Time Series format.

    This is a thin wrapper around the structure-specific writer that
    defaults ``dimension_at_observation`` to ``TIME_PERIOD`` for all
    datasets and uses the Structure Specific Time Series message type.

    Args:
        datasets: The datasets to be written.
        output_path: The path to save the file.
        prettyprint: Prettyprint or not.
        header: The header to be used (generated if None).
        dimension_at_observation:
          The mapping between the dataset and the dimension at observation.
          Defaults to TIME_PERIOD for all datasets.

    Returns:
        The XML string if path is empty, None otherwise.
    """
    if dimension_at_observation is None:
        dimension_at_observation = {
            ds.short_urn: "TIME_PERIOD" for ds in datasets
        }
    else:
        invalid_dims = {
            k: v
            for k, v in dimension_at_observation.items()
            if v != "TIME_PERIOD"
        }
        if invalid_dims:
            raise Invalid(
                "Time Series formats require "
                "dimensionAtObservation=TIME_PERIOD",
                ", ".join(f"{k}={v}" for k, v in invalid_dims.items()),
            )

    return _base_write(
        datasets=datasets,
        output_path=output_path,
        prettyprint=prettyprint,
        header=header,
        dimension_at_observation=dimension_at_observation,
        sdmx_format=Format.DATA_SDMX_ML_2_1_STRTS,
    )
