from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import yaml

from tsfm_bench.data.interface import TimeSeriesDataSource
from tsfm_bench.data.open_meteo import OpenMeteoDataSource
from tsfm_bench.data.synthetic import DEFAULT_SYNTHETIC_SPECS, SyntheticDataSource, _SyntheticSpec
from tsfm_bench.data.ts_bench import TsBenchDataSource


def load_data_source(config_path: str | Path | None = None) -> TimeSeriesDataSource:
    """Load a data source from config. Defaults to built-in synthetic demo."""
    if config_path is None:
        return SyntheticDataSource()

    path = Path(config_path)
    if not path.exists():
        raise FileNotFoundError(f"Dataset config not found: {path}")

    if path.suffix in {".yaml", ".yml"}:
        payload = yaml.safe_load(path.read_text())
    elif path.suffix == ".json":
        payload = json.loads(path.read_text())
    else:
        raise ValueError(f"Unsupported config format: {path.suffix}")
    payload["__config_dir"] = str(path.resolve().parent)

    source_type = payload.get("source_type", "synthetic")
    if source_type == "synthetic":
        specs = tuple(_parse_spec(entry) for entry in payload.get("datasets", []))
        return SyntheticDataSource(specs=specs or DEFAULT_SYNTHETIC_SPECS)
    if source_type == "open_meteo":
        return OpenMeteoDataSource.from_config(payload)
    if source_type == "ts_bench":
        return TsBenchDataSource.from_config(payload)

    raise NotImplementedError(
        f"Source type '{source_type}' is not implemented yet. "
        "Implement TimeSeriesDataSource and register it in registry.py."
    )


def load_dataset_properties(config_path: str | Path | None = None) -> dict[str, dict[str, Any]]:
    """Build dataset_properties map keyed by dataset stem (before frequency suffix)."""
    source = load_data_source(config_path)
    properties: dict[str, dict[str, Any]] = {}

    for name in source.list_datasets():
        meta = source.get_metadata(name)
        if isinstance(source, TsBenchDataSource):
            ds_key = source.leaderboard_dataset_name(name).split("/")[0].lower()
        else:
            ds_key = name.split("/")[0].lower()
        properties[ds_key] = {
            "domain": meta.domain,
            "frequency": meta.frequency,
            "num_variates": meta.num_variates,
            "terms": list(meta.terms),
        }
    return properties


def _parse_spec(entry: dict[str, Any]) -> _SyntheticSpec:
    return _SyntheticSpec(
        name=entry["name"],
        domain=entry["domain"],
        frequency=entry["frequency"],
        num_variates=int(entry["num_variates"]),
        terms=tuple(entry.get("terms", ["short"])),
        n_series=int(entry.get("n_series", 30)),
        series_length=int(entry.get("series_length", 500)),
        seed=int(entry.get("seed", 0)),
    )